wrapper.py

"""
Luke Cirne
ColabFold Wrapper
Template Cycling with Experimental Distance Restraint Data
Ma Lab
"""
import os
import sys
import subprocess
import shutil
import json
import argparse
import getpass
from pathlib import Path
import numpy as np
import data_engine as engine

# ---- Global Variables ----
MOD_COUNTS_FILEPATH = "/gpfs1/home/l/c/lcirne/ColabFoldWrapper/mod_counts.json"

num_iterations = 0

# Setters for global vars
def set_num_iterations(iters):
    global num_iterations
    num_iterations = iters
# --------------------------

def initialize_project(jobs):
    """
    Initializes a ColabFold project by gathering user input,
    creating necessary variables, generating a shell script,
    and appending metadata to a JSON log.

    Args:
        jobs (str): Path to the JSON file tracking job metadata.

    Returns:
        str: Path to the generated shell script for running ColabFold.
        str: Content of shell script generated by this function for running ColabFold.
    """
    key_values = []
    # Obtain username
    try:
        username = getpass.getuser()
        print(f">>> LOGGING AS: {username}")
    except OSError:
        print(">>> COULD NOT OBTAIN USERNAME")
        print(">>> LOGGING AS: DEFAULT")
        username = "DEFAULT"
    key_values.append(("user", username))

    # Obtain job ID
    try:
        with open(jobs, "r") as file:
            print(">>> READING JSON")
            job_dict = json.load(file)
            user_JIDs = [int(job["JID"]) for job in job_dict["jobs"] if job["user"] == username]
            last_user_JID = max(user_JIDs) if user_JIDs else 0
            current_JID = last_user_JID + 1
            key_values.append(("JID", current_JID))
    except FileNotFoundError:
        current_JID = 0
        key_values.append(("JID", current_JID))

    # Obtain input file
    while True:
        input_file = input("Input file name in the format 'name.fasta': ")
        if os.path.isfile(f"./{input_file}"):
            break
        else:
            print("###### INVALID FILEPATH ######")
    key_values.append(("input_file", input_file))

    # Obtain template directory
    while True:
        temp_dir = input("Input template directory: ")
        if os.path.isdir(temp_dir):
            break
        else:
            print("###### INVALID DIRECTORY ######")
    key_values.append(("temp_dir", temp_dir))

    # Obtain num recycles
    while True:
        num_c = input("Desired number of recycles (integer) (max 5, min 0): ")
        try:
            if 0 <= int(num_c) < 6:
                break
            else:
                print("###### Invalid input ######")
        except ValueError:
                print("###### Invalid input ######")
    key_values.append(("num_recycles", num_c))    

    # Obtain num seeds
    while True:
        num_s = input("Desired number of seeds (integer) (min 1): ")
        try:
            if 0 < int(num_s):
                num_s = int(num_s)
                break
            else:
                print("###### Invalid input #######")
        except ValueError:
                print("###### Invalid input #######")
    key_values.append(("num_s", num_s))

    # Obtain num iterations
    while True:
        iters = input("Desired number of iterations for the wrapper (integer) (min 1): ")
        try:
            if int(iters) >= 1:
                set_num_iterations(int(iters))
                break
            else:
                print("###### Invalid Input ######")
        except ValueError:
                print("###### Invalid Input ######")

    # Obtain value for n
    while True:
        n = input("Desired value for n (integer) (min 10): ")
        try:
            if 9 < int(n):
                n = int(n)
                break
            else:
                print("###### Invalid input #######")
        except ValueError:
                print("###### Invalid input #######")
    key_values.append(("n", n))

    # Obtain value for m_e_msa
    while True:
        m_e_msa = input("Desired value for max number of msa's (integer) (min 1): ")
        try:
            if 1 < int(m_e_msa):
                m_e_msa = int(m_e_msa)
                break
            else:
                print("###### Invalid input #######")
        except ValueError:
                print("###### Invalid input #######")
    key_values.append(("m_e_msa", m_e_msa))
    m_msa = m_e_msa // 2
    key_values.append(("m_msa", m_msa))


    while True:
        num_models = input("Desired number of AlphaFold models (integer) (min 1) (max 5): ")
        try:
            if 1 <= int(num_models) and int(num_models) <= 5:
                num_models = int(num_models)
                break
            else:
                print("###### Invalid input #######")
        except ValueError:
                print("###### Invalid input #######")
    key_values.append(("num_models", num_models))

    # Variable for full output directory by user, job id, and max msa's
    dir_name = f"{username}{current_JID}mm{m_msa}"
    container_name = f"{username}{current_JID}mm{m_msa}-container"
    outputdir = os.path.abspath(f"{container_name}/{dir_name}")
    key_values.append(("outputdir", outputdir))

    # Create shell script to run colabfold_batch
    current_dir = os.path.dirname(os.path.abspath(__file__))
    script_name = "wrapper.sh"
    script_path = os.path.join(current_dir, script_name)

    # -------------------- Writing script --------------------
    script_content = f"""#!/bin/bash
JID={current_JID}
num_c={num_c}
seed=1
num_s={num_s}
m_e_msa={m_e_msa}
m_msa={m_msa}
inputfile=./{input_file}
outputdir={outputdir}
temp_dir={temp_dir}
num_models={num_models}

export PATH="/gpfs1/home/l/c/lcirne/localcolabfold/.pixi/envs/default/bin:${{PATH}}"

colabfold_batch --pair-mode unpaired_paired --templates \\
--msa-mode mmseqs2_uniref_env \\
--custom-template-path $temp_dir \\
--max-msa $m_msa:$m_e_msa \\
--use-dropout \\
--num-seeds $num_s \\
--num-recycle $num_c \\
--num-models $num_models \\
$inputfile $outputdir
    """
    # -------------------------------------------------------

    # Create shell script to execute ColabFold
    with open(script_path, 'w') as file:
        print(">>> WRITING SHELL SCRIPT")
        file.write(script_content)
    append_jobs_json("jobs.json", key_values)
    return script_path, script_content


def delete_directory(dir_path):
    """
    Deletes a directory and all its contents.

    Args:
        dir_path (str): The path to the directory to delete.

    Returns:
        bool: True if the directory was successfully deleted, False otherwise.
    """
    try:
        shutil.rmtree(dir_path)
        return True
    except OSError:
        return False


def clear_directory(dir_path):
    """
    Clears the contents of a directory without deleting the directory itself.

    Args:
        dir_path (str): Path to the directory to be cleared.
    """
    for item in os.listdir(dir_path):
        item_path = os.path.join(dir_path, item)
        if os.path.isfile(item_path):
            os.remove(item_path)  # Remove files
        elif os.path.isdir(item_path):
            shutil.rmtree(item_path)


def append_jobs_json(jobs, key_values):
    """
    Appends job metadata to a JSON file. If the file does not exist,
    it creates a new one.

    Args:
        jobs (str): Path to the JSON file.
        key_values (list): List of (key, value) tuples to append.
    """
    # Convert list of tuples into a proper dictionary
    new_entry = {key: value for key, value in key_values}
    try:
        with open(jobs, "r") as f:
            job_dict = json.load(f)
    except FileNotFoundError:
        # Create new JSON structure if file doesn't exist
        job_dict = {"jobs": []}

    # Append the new dictionary entry
    job_dict["jobs"].append(new_entry)

    with open(jobs, "w") as f:
        print(f">>> APPENDING JSON TO {jobs}")
        json.dump(job_dict, f, indent=4)
        print("###### COMPLETE ######")


def append_mods_json(mods_file, mods_dict):
    """
    Append distribution modifications from the distribution building
    process to the json logs.

    Args:
        mods_file (str): Path to JSON file.
        mods_json (dict): Dictionary containing json to append.
    """
    mods_json = None
    while not mods_json:
        try:
            with open(mods_file, "r") as f:
                mods_json = json.load(f)
        except FileNotFoundError as e:
            print(f">>> EXCEPTION WHEN APPENDING TO {mods_file}: {e}")
            subprocess.run(["touch", mods_file])
        except json.JSONDecodeError:
            mods_json = {"mods": []}

    mods_json["mods"].append(mods_dict)
    with open(mods_file, "w") as f:
        print(f">>> APPENDING JSON TO {mods_file}")
        print(mods_dict)
        json.dump(mods_json, f, indent=4)
        print("###### COMPLETE ######")


def get_from_current_job(jobs_file, items) -> list:
    try:
        with open(jobs_file, "r") as file:
            jobs_dict = json.load(file)
            current_job_info = list(jobs_dict["jobs"][-1].values())
            print(current_job_info)
            requested_items = []
            for item in items:
                match item:
                    case "user":
                        requested_items.append(current_job_info[0])
                    case "JID":
                        requested_items.append(current_job_info[1])
                    case "input_file":
                        requested_items.append(current_job_info[2])
                    case "temp_dir":
                        requested_items.append(current_job_info[3])
                    case "num_recycles":
                        requested_items.append(current_job_info[4])
                    case "num_s":
                        requested_items.append(current_job_info[5])
                    case "n":
                        requested_items.append(current_job_info[6])
                    case "m_e_msa":
                        requested_items.append(current_job_info[7])
                    case "m_msa":
                        requested_items.append(current_job_info[8])
                    case "run_number":
                        requested_items.append(current_job_info[9])
                    case "outputdir":
                        requested_items.append(current_job_info[10])
        return requested_items
    except FileNotFoundError:
        print("###### JSON FILE NOT FOUND ######")
    return 0


def filter_output(run_number, jobs, script_path, n):
    """
    Filters PDB output files based on proximity to target distances.
    Updates template directory for next ColabFold iteration accordingly.

    Args:
        run_number (int): The current iteration number.
        jobs (str): Path to the job metadata JSON file.
        script_path (str): Path to the ColabFold execution script.
    """
    # Load json and obtain outputdir and temp_dir
    outputdir, temp_dir = get_from_current_job(jobs, ["outputdir", "temp_dir"])
    print(outputdir)

    # 1. Create the pool dir
    output_path = Path(outputdir)
    parent_dir = output_path.parent
    output_pool = parent_dir / "output_pool"
    output_pool.mkdir(exist_ok=True)

    # 2. Copy outputdir contents to the pool dir (append)
    subprocess.run(["cp", "-r", f"{outputdir}/", f"{output_pool}/"])
    output_name = os.path.basename(outputdir)
    current_iteration = f"iteration{run_number+1}"
    subprocess.run([
        "mv",
        f"{output_pool}/{output_name}",
        f"{output_pool}/{current_iteration}"
    ])

    # 3. Pass the pool dir to build_distribution
    colabfold_output = []
    # Traverse output_pool with os.walk, checking each iteration subdirectory
    for root, dirs, files in os.walk(output_pool):
        for file in files:
            if file.endswith(".pdb"): # Filter for pdbs only
                abs_path = os.path.abspath(os.path.join(root, file))
                colabfold_output.append(abs_path)

    old_iteration_distances = {}
    current_iteration_distances = {}
    # Find distances between "probes"
    for file_abs_path in colabfold_output:
        distance = float(run_distance_finder(f"{file_abs_path}", "100", "473"))
        # Separate into new and old distances to graph current distribution in
        # in isolation before sampling from all
        if current_iteration in file_abs_path:
            current_iteration_distances[file_abs_path] = distance
        else:
            old_iteration_distances[file_abs_path] = distance

    # Convert distances to efficiencies
    old_distances_to_convert = np.array(list(old_iteration_distances.values()))
    old_e_conversions = engine.compute_E(old_distances_to_convert)

    for filename, eff in zip(
        old_iteration_distances.keys(),
        old_e_conversions
    ):
        old_iteration_distances[filename] = eff

    current_distances_to_convert = np.array(list(current_iteration_distances.values()))
    current_e_conversions = engine.compute_E(current_distances_to_convert)

    for filename, eff in zip(
        current_iteration_distances.keys(),
        current_e_conversions
    ):
        current_iteration_distances[filename] = eff

    distances = old_iteration_distances | current_iteration_distances

    # Execute algorithm to determine which files to extract from distances
    # and add to included_distances to fit a normal distribution
    # Remember to normalize data points
    # Duplicate / discard templates if necessary to fit proper distribution.
    y_exp = 0.291
    sigma = 0.083
    included_distances, bins, bin_centers, mod_count = engine.build_distribution(file_eff_dict=distances, mean=y_exp, std=sigma, n=n)

    # Plot and save original distances using bins from build_distribution
    plot_and_save_distances(current_iteration_distances, run_number, bins, n)

    # Plot and save filtered distances
    plot_and_save_distances(included_distances, run_number, bins, n, dirname="filtered_distributions")

    # If included_distances dictionary is still empty after checks,
    # proceed to next iteration with user provided templates 
    if not included_distances:
        print("###### NO VALID TEMPLATES PRODUCED ######")
        update_temp_dir(script_path, temp_dir)
    else:
        # Creating dir for selected templates
        temp_dir = f"iteration{run_number + 1}"
        try:
            os.mkdir("iterations")
        except FileExistsError:
            print(">>> APPENDING TO ITERATIONS DIRECTORY")
        try:
            os.mkdir(f"iterations/{temp_dir}")
        except FileExistsError:
            shutil.rmtree(f"iterations/{temp_dir}")
            os.mkdir(f"iterations/{temp_dir}")

        template_number = 0
        for filepath, distance in included_distances.items():
            # Logic to check if a filename is duplicated or not
            # if so, cp the original file with new name and add to temp_dir
            # if not, just cp original file to temp_dir
            filepath = Path(filepath)
            filename = filepath.name

            if "_dupe" in filename:
                seperator = "_dupe"
                filename_parts = filename.split(seperator, 1)

                original_filename = f"{filename_parts[0]}.pdb"
                original_path = filepath.parent / original_filename

                dupe_path = filepath.parent / filename
                subprocess.run(["cp", str(original_path), str(dupe_path)])

            subprocess.run([
                "cp",
                str(filepath),
                f"iterations/{temp_dir}"
            ])

            template_number_str = f"{template_number:04d}"

            subprocess.run([
                "mv",
                f"iterations/{temp_dir}/{filename}",
                f"iterations/{temp_dir}/{template_number_str}.pdb"
            ])

            print(f"###### {filename} ADDED TO {temp_dir} ({distance} A) ######")
            template_number += 1

        update_temp_dir(script_path, f"iterations/{temp_dir}")

        # Clear ouput directory
        if run_number < num_iterations - 1:
            clear_directory(outputdir)
    return mod_count


def run_distance_finder(structure_file, p1, p2):
    """
    Runs an external script to calculate the distance between two residues
    in a protein structure.

    Args:
        structure_file (str): Path to the .pdb file.
        p1 (str): Residue index 1.
        p2 (str): Residue index 2.

    Returns:
        str or None: Distance in angstroms as a string, or None if failed.
    """
    distance = subprocess.run(
        [sys.executable, "distance_finder.py", structure_file, p1, p2],
        capture_output=True,
        text=True
    )
    if distance.returncode != 0:
        print("Error:", distance.stderr)
        return None
    return distance.stdout.strip()


def update_temp_dir(script_path, dir_name):
    """
    Updates the 'temp_dir' line in the shell script to point to a new directory.

    Args:
        script_path (str): Path to the shell script.
        dir_name (str): Name of the new template directory.
    """
    with open(script_path, 'r') as file:
        lines = file.readlines()

    with open(script_path, 'w') as file:
        for line in lines:
            if line.startswith("temp_dir="):
                file.write(f"temp_dir={dir_name}\n")
            else:
                file.write(line)


def plot_and_save_distances(distances, run_number, bin_edges, n, dirname="distance_distributions"):
    os.makedirs(dirname, exist_ok=True)
    plot_name = f"{engine.graph_output_accuracy_bar(distances, bins=bin_edges, N=n)}"
    subprocess.run(["mv", f"{plot_name}.png", f"./{dirname}/{plot_name}{run_number+1}.png"])
    return 0


def plot_fret_efficiencies(distances: dict, run_number: int, bin_centers: list[float], n: int):
    os.makedirs("distribution_graphs", exist_ok=True)
    plot_name = f"{engine.graph_output_accuracy(distances, bins=bin_centers, N=n)}"
    subprocess.run(["mv", f"{plot_name}.png", "graphing-utils/distribution_graphs/"])
    return 0


def main():
    """
    Entry point for the ColabFold Wrapper.
    Initializes project setup and executes the template filtering loop.
    """

    # Welcome message
    print("*" * 31)
    print()
    print((" " * 7) + "ColabFold Wrapper")
    print()
    print("*" * 31)

    jobs = os.path.abspath("jobs.json")
    iterations_script_path, script_content = initialize_project(jobs)
    i0_script_path = iterations_script_path

    # ------------------------ CLI flags -------------------------
    parser = argparse.ArgumentParser()

    # ColabFold_batch flags
    parser.add_argument("--pair-mode",
                        required=True,
                        help="ColabFold_batch required flag.",
                        default="unpaired_paired")
    parser.add_argument("--templates",
                        action="store_true",
                        required=True,
                        help="ColabFold_batch required flag.")
    parser.add_argument("--msa-mode",
                        default="mmseqs2_uniref_env")
    parser.add_argument("--custom-template-path")
    parser.add_argument("--max-msa",
                        required=True)
    parser.add_argument("--use-dropout",
                        action="store_true")
    parser.add_argument("--num-seeds",
                        required=True)
    parser.add_argument("--num-recycle",
                        required=True)
    parser.add_argument("--num-models",
                        required=True)

    # Wrapper flags
    parser.add_argument("--no-templates",
                        action="store_true",
                        help="Turn off custom templates for the inital iteration of ColabFold Wrapper")
    parser.add_argument("--num-iterations",
                        default=10)
    parser.add_argument("--num-n", "-n",
                        required=True)
    # ------------------------------------------------------------
    args = parser.parse_args()
    no_templates = getattr(args, "no_templates", False)
    if no_templates:
        # Generate a new script without --custom-template-path flag
        # and change i0_script_path
        i0_script_path = f"{iterations_script_path[:-3]}_i0.sh"
        with open(f"{i0_script_path}", "w") as file:
            for line in script_content.splitlines():
                if line.startswith("num_models="):
                    # When using no custom templates, default to 5 models
                    line = "num_models=5"
                if line.startswith("--custom-template-path"):
                    continue
                file.write(line + "\n")

    # Remove output pool from any previous wrapper run
    n, outputdir = get_from_current_job(jobs, ["n", "outputdir"])
    output_path = Path(outputdir)
    parent_dir = output_path.parent
    print(f"RUNNING FROM PARENT DIR: {parent_dir}")
    output_pool = parent_dir / "output_pool"
    if output_pool.exists():
        shutil.rmtree(output_pool)

    print(">>> ATTEMPTING TO RUN COLABFOLD\n")
    outputdir_container = f"{outputdir}-container"
    n = int(n) # n = number of templates passed as input through iterations
    mod_counts = {outputdir: {}}

    # Create output directory container and cd into it
    subprocess.run(["mkdir", "-p", outputdir_container])
    #subprocess.run(["cd", outputdir_container])
    #mods = os.path.abspath("mod_counts.json")
    mods = MOD_COUNTS_FILEPATH

    for run_number in range(num_iterations):
        """
        Start with three iterations for testing
        Once running, continue iterating until an ideal structure is output
        """
        if run_number == 0:
            script_path = i0_script_path
        else:
            script_path = iterations_script_path
        print("-"*30)
        print("Running script from: ", script_path)
        print("-"*30)
        os.chmod(script_path, 0o755)
        subprocess.run([script_path], check=True)
        iteration_mod_count = filter_output(run_number, jobs, script_path, n)
        mod_counts[outputdir][run_number] = int(iteration_mod_count) # Ensure count is NOT np.int64

    subprocess.run(["mv", "./iterations/", outputdir])
    subprocess.run(["mv", './*distributions/', outputdir])
    subprocess.run(["rm", "-rf", outputdir_container])
    subprocess.run(["mv", output_pool, parent_dir.parent])
    subprocess.run(["rm", "-rf", parent_dir])
    append_mods_json(mods, mod_counts)

if __name__ == '__main__':
    main()