# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/af3/00_json.ipynb.

# %% auto 0
__all__ = ['dump_json', 'get_protein_json', 'read_json', 'get_protein_smiles_json', 'get_protein_ccdcode_json',
           'assign_atom_names_from_graph', 'mol_to_ccd_text', 'sdf2ccd', 'get_protein_ccd_json', 'split_nfolder']

# %% ../../nbs/af3/00_json.ipynb 2
import re, shutil, json, pandas as pd, numpy as np
from pathlib import Path

from rdkit import Chem as rd_chem
from rdkit.Chem import AllChem,rdmolfiles
from rdkit import Chem

from Bio.PDB import PDBParser

# %% ../../nbs/af3/00_json.ipynb 4
def dump_json(data, save_path):
    "Save json data into a file"
    with open(save_path,'w') as f: 
        json.dump(data,f,indent=4)

# %% ../../nbs/af3/00_json.ipynb 5
def get_protein_json(name, # job name
                     seq, # aa sequence
                     save_path=None, # .json
                     seeds=[1]
                     ):
    "Generate json of single protein sequence for input of docker command"
    
    json_data = {
        "name": name,
        "modelSeeds": seeds,
        "sequences": [
            {
                "protein": {
                    "id": "A",
                    "sequence": seq,
                }
            },
        ],
        "bondedAtomPairs": [],
        "dialect": "alphafold3",
        "version": 3
    }
    if save_path:
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        dump_json(json_data,save_path)
    return json_data

# %% ../../nbs/af3/00_json.ipynb 9
def read_json(file_path):
    with open(file_path,'r') as f: 
        data = json.load(f)
    return data

# %% ../../nbs/af3/00_json.ipynb 11
def get_protein_smiles_json(smi_id:str, 
                            SMILES:str, 
                            protein_json, # json type
                            save_path=None, # .json
                            seeds=[1]
                            ):
    
    "Get json for protein-ligand docking task"
    raw_smiles = r"{}".format(SMILES) # JSON escaping, \ to \\
    protein_index = next(i for i, item in enumerate(protein_json["sequences"]) if "protein" in item)
    json_data = {
        "name": smi_id,
        "modelSeeds": seeds,
        "sequences": [
            {
                "ligand": {
                    "id": "L",
                    "smiles": raw_smiles,
                }
            }, 
            {
                "protein": protein_json["sequences"][protein_index]["protein"]
            },
        ],
        "bondedAtomPairs": [],
        "dialect": "alphafold3",
        "version": 2
    }
    if save_path:
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        dump_json(json_data,save_path)
    return json_data

# %% ../../nbs/af3/00_json.ipynb 18
def get_protein_ccdcode_json(protein_json,  # dict with protein sequence
                              ccd_code,      # str or list of str
                              job_id: str,   # job/task ID
                              save_path=None,  # optional output path
                              seeds=[1]):      # optional random seeds
    "Create AlphaFold3 docking JSON with CCD code(s)."
    
    # Normalize ccd_code to a list
    if isinstance(ccd_code, str):
        ccd_code = [ccd_code]
    elif not isinstance(ccd_code, list):
        raise TypeError("ccd_code must be a string or a list of strings.")

    protein_index = next(i for i, item in enumerate(protein_json["sequences"]) if "protein" in item)

    json_data = {
        "name": job_id,
        "modelSeeds": seeds,
        "sequences": [
            {
                "ligand": {
                    "id": "L",
                    "ccdCodes": ccd_code
                }
            },
            {
                "protein": protein_json["sequences"][protein_index]["protein"]
            },
        ],
        "dialect": "alphafold3",
        "version": 3
    }

    if save_path:
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        dump_json(json_data, save_path)

    return json_data

# %% ../../nbs/af3/00_json.ipynb 22
# Mapping bond types to mmCIF-compatible values
_RDKIT_BOND_TYPE_TO_MMCIF = {
    rd_chem.BondType.SINGLE: 'SING',
    rd_chem.BondType.DOUBLE: 'DOUB',
    rd_chem.BondType.TRIPLE: 'TRIP',
    rd_chem.BondType.AROMATIC: 'AROM'
}

def assign_atom_names_from_graph(mol):
    for i, atom in enumerate(mol.GetAtoms()):
        atom.SetProp('atom_name', f"{atom.GetSymbol()}{i+1}")
    return mol

def mol_to_ccd_text(mol, component_id, pdbx_smiles=None, include_hydrogens=False):
    mol = rd_chem.Mol(mol)
    if include_hydrogens:
        mol = rd_chem.AddHs(mol)
    rd_chem.Kekulize(mol, clearAromaticFlags=True)

    if mol.GetNumConformers() == 0:
        raise ValueError('The molecule has no conformers')
    conf = mol.GetConformer()
    coords = conf.GetPositions()

    mol = assign_atom_names_from_graph(mol)
    atom_map = {atom.GetIdx(): atom.GetProp('atom_name') for atom in mol.GetAtoms()}

    lines = [
        f"data_{component_id}",
        "#",
        f"_chem_comp.id {component_id}",
        f"_chem_comp.name '{component_id}'",
        "_chem_comp.type non-polymer",
        "_chem_comp.formula '?'",
        "_chem_comp.mon_nstd_parent_comp_id ?",
        "_chem_comp.pdbx_synonyms ?",
        "_chem_comp.formula_weight '?'",
    ]
    if pdbx_smiles:
        lines.append(f"_chem_comp.pdbx_smiles {pdbx_smiles}")
    lines += [
        "#",
        "loop_",
        "_chem_comp_atom.comp_id",
        "_chem_comp_atom.atom_id",
        "_chem_comp_atom.type_symbol",
        "_chem_comp_atom.charge",
        "_chem_comp_atom.pdbx_leaving_atom_flag",
        "_chem_comp_atom.pdbx_model_Cartn_x_ideal",
        "_chem_comp_atom.pdbx_model_Cartn_y_ideal",
        "_chem_comp_atom.pdbx_model_Cartn_z_ideal"
    ]

    for i, atom in enumerate(mol.GetAtoms()):
        if not include_hydrogens and atom.GetSymbol() == 'H':
            continue
        x, y, z = coords[i]
        lines.append(f"{component_id} {atom_map[atom.GetIdx()]} {atom.GetSymbol()} {atom.GetFormalCharge()} N {x:.3f} {y:.3f} {z:.3f}")

    lines += [
        "#",
        "loop_",
        "_chem_comp_bond.atom_id_1",
        "_chem_comp_bond.atom_id_2",
        "_chem_comp_bond.value_order",
        "_chem_comp_bond.pdbx_aromatic_flag"
    ]

    for bond in mol.GetBonds():
        a1, a2 = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
        if not include_hydrogens and (mol.GetAtomWithIdx(a1).GetSymbol() == 'H' or mol.GetAtomWithIdx(a2).GetSymbol() == 'H'):
            continue
        bond_type = _RDKIT_BOND_TYPE_TO_MMCIF[bond.GetBondType()]
        aromatic_flag = 'Y' if bond.GetIsAromatic() else 'N'
        lines.append(f"{atom_map[a1]} {atom_map[a2]} {bond_type} {aromatic_flag}")
    lines.append("#")

    return "\n".join(lines)

# %% ../../nbs/af3/00_json.ipynb 23
def sdf2ccd(sdf_path,
            CCD_name='lig-1', # do not use '_'; use as less letter as possible, 'lig-any' leads to extra ligands
            ):

    "Convert the compound to the AF3 required CCD format"
    supplier = Chem.SDMolSupplier(sdf_path)
    mol = supplier[0]  # Get the first molecule
    return mol_to_ccd_text(mol,CCD_name)

# %% ../../nbs/af3/00_json.ipynb 26
def get_protein_ccd_json(protein_json, # dict with protein sequence
                         rec_residue_num:int, # 1-indexed, for bondedAtomPairs, e.g., ["A", 145, "SG"] 
                         rec_atom_id:str, # for bondedAtomPairs, e.g., ["A", 145, "SG"] 
                         lig_sdf_path, # ccd text
                         lig_atom_id:str, # 0-indexed, for bondedAtomPairs, ["L", 1, "C04"]
                         job_id:str, # str, job/task ID
                         save_path=None,# optional output path
                         seeds=[1], # optional random seeds
                         ):               
    "Create AlphaFold3 docking JSON with customized CCD ligand and bondedAtomPairs."

    # get userCCD
    userCCD=sdf2ccd(lig_sdf_path)
    ccd_id = re.search(r"_chem_comp.id\s+([^\s#]+)", userCCD).group(1)
    
    protein_index = next(i for i, item in enumerate(protein_json["sequences"]) if "protein" in item)

    json_data = {
        "name": job_id,
        "modelSeeds": seeds,
        "sequences": [
            {
                "ligand": {
                    "id": "L",
                    "ccdCodes": [ccd_id]
                }
            },
            {
                "protein": protein_json["sequences"][protein_index]["protein"]
            },
        ],
        "bondedAtomPairs": [[["A", rec_residue_num, rec_atom_id],["L", 1, lig_atom_id]]],
        "userCCD": userCCD,
        "dialect": "alphafold3",
        "version": 3
    }

    if save_path:
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        dump_json(json_data, save_path)

    return json_data

# %% ../../nbs/af3/00_json.ipynb 30
def split_nfolder(folder_dir, 
                  n=4):
    "Move json files from a folder into subfolders (folder_0, folder_1, ..., folder_N)."
    
    folder_dir = Path(folder_dir)

    files = sorted(folder_dir.glob("*.json"))
    # print(len(files))
    subfolders = [folder_dir / f"folder_{i}" for i in range(n)]
    for folder in subfolders:
        folder.mkdir(exist_ok=True)

    for idx, file in enumerate(files):
        target_folder = subfolders[idx % n]
        shutil.move(str(file), target_folder / file.name)

    print(f"Distributed {len(files)} files into {n} folders.")
