# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/af3/02_protein_pairs.ipynb.

# %% auto 0
__all__ = ['get_colabfold_cmd', 'copy_a3m', 'a3m_to_seq', 'get_protein_subjson', 'dump_json_folder', 'get_multi_protein_json',
           'generate_pair_df']

# %% ../../nbs/af3/02_protein_pairs.ipynb 4
import os, json, shutil, pandas as pd
from tqdm import tqdm
from itertools import combinations
from pathlib import Path
from .json import *
from .docker import *

# %% ../../nbs/af3/02_protein_pairs.ipynb 9
def get_colabfold_cmd(csv_path,project_name):
    print('Run below in terminal:')
    print(f'\n colabfold_batch {csv_path} msa_{project_name} --msa-only')

# %% ../../nbs/af3/02_protein_pairs.ipynb 13
def copy_a3m(a3m_dir: str, # Path to the source directory containing .a3m files.
             dest_dir: str, # Path to the destination directory where files will be copied
             ):
    "Copies all .a3m files from the source directory to the destination directory."
    
    a3m_dir,dest_dir = Path(a3m_dir),Path(dest_dir)
    dest_dir.mkdir(parents=True, exist_ok=True)

    files = list(a3m_dir.glob('*.a3m'))

    for file in tqdm(files, desc="Copying files", unit="file"):
        shutil.copy(file, dest_dir / file.name)

    print(f"Copied {len(files)} a3m files from {a3m_dir} to {dest_dir}")

# %% ../../nbs/af3/02_protein_pairs.ipynb 17
def a3m_to_seq(file_path:Path):
    "Get protein sequence from a3m file"
    return file_path.read_text().splitlines()[2] # protein sequence is located on line 2

# %% ../../nbs/af3/02_protein_pairs.ipynb 19
def get_protein_subjson(gene_name, a3m_dir=".",idx = 'A',run_template=True):
    "Get subjson (protein part) with colabfold unpairedMSA .a3m path"
    file_path = Path(a3m_dir)/f"{gene_name}.a3m"
    protein_sequence = a3m_to_seq(file_path)
    
    json_data = {
        'id': idx,
        'sequence': protein_sequence, 
        'modifications': [],
        'unpairedMsaPath': str("/root"/file_path), # for docker path, ECD under af_input
        'pairedMsa': '',
        'templates': None if run_template else []
    }

    return json_data

# %% ../../nbs/af3/02_protein_pairs.ipynb 22
def dump_json_folder(json_data, folder):
    "Save json under a folder"
    file_path = Path(folder)/f"{json_data['name']}.json"
    with open(file_path,'w') as f: json.dump(json_data,f,indent=4)

# %% ../../nbs/af3/02_protein_pairs.ipynb 23
def get_multi_protein_json(gene_list,a3m_dir,run_template=True,save_folder=None):
    'Get json of multiple proteins, with unpaired MSA path indicated (from colabfold MSA)'
    sequences = []
    alphabets = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
    for index, gene in enumerate(gene_list):
        sub_json=get_protein_subjson(gene,a3m_dir,idx=alphabets[index],run_template=run_template)
        sequences.append({'protein':sub_json})
    name = '_'.join(gene_list)
    json_data = {
            "name": name,
            "modelSeeds": [1],
            "sequences": sequences,
            "bondedAtomPairs": [],
            "dialect": "alphafold3",
            "version": 2
        }
    if save_folder:
        dump_json_folder(json_data,save_folder)
    return json_data

# %% ../../nbs/af3/02_protein_pairs.ipynb 27
def generate_pair_df(gene_list,self_pair=True):
    "Unique pair genes in a gene list"
    pairs = list(combinations(gene_list, 2))
    pair_df = pd.DataFrame(pairs,columns=["Gene1", "Gene2"])
    
    if self_pair:
        self_pair_df = pd.DataFrame({'Gene1':gene_list, 'Gene2':gene_list})
        pair_df = pd.concat([pair_df,self_pair_df])

    return pair_df.reset_index(drop=True)
