# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/protenix/07_protenix.ipynb.

# %% auto 0
__all__ = ['get_single_job', 'get_single_protein_ligand_json', 'get_protein_ligand_df_json', 'get_virtual_screening_json']

# %% ../../nbs/protenix/07_protenix.ipynb 6
import json
from pathlib import Path

# %% ../../nbs/protenix/07_protenix.ipynb 9
def get_single_job(job_name, protein_seq, msa_dir, SMILES=None,CCD=None):
    "Get protenix json format of protein and ligand."
    
    if SMILES and CCD:
        raise ValueError("Please provide only one of SMILES or CCD, not both.")
    if not SMILES and not CCD:
        raise ValueError("You must provide either SMILES or CCD.")

    ligand_value = SMILES if SMILES else f"CCD_{CCD}"
    
    return {
        "name": job_name,
        "sequences": [
            {
                "proteinChain": {
                    "count": 1,
                    "sequence": protein_seq,
                    "msa": {
                        "precomputed_msa_dir": msa_dir,
                        "pairing_db": "uniref100"
                    }
                }
            },
            {
                "ligand": {
                    "count": 1,
                    "ligand": ligand_value
                }
            }
        ]
    }

# %% ../../nbs/protenix/07_protenix.ipynb 11
def get_single_protein_ligand_json(job_name, 
                                   protein_seq, 
                                   msa_dir, 
                                   SMILES=None, 
                                   CCD=None, 
                                   json_path=None):
    "Generate json input for one protein-ligand job."
    data = [get_single_job(job_name, protein_seq, msa_dir, SMILES=SMILES, CCD=CCD)]

    if json_path:
        save_path = Path(json_path)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        with save_path.open("w") as f:
            json.dump(data, f, indent=4)
        print(f"JSON saved to {save_path}")

    return data

# %% ../../nbs/protenix/07_protenix.ipynb 16
def get_protein_ligand_df_json(df,
                               id_col,
                               seq_col, 
                               msa_col, 
                               smi_col=None, 
                               ccd_col=None, 
                               save_json=None):
    "Get json file of protein and ligand in a dataframe."
    
    if smi_col and ccd_col:
        raise ValueError("Provide only one of smi_col or ccd_col, not both.")
    if not smi_col and not ccd_col:
        raise ValueError("You must provide either smi_col or ccd_col.")

    use_smiles = smi_col is not None

    def build_job(row):
        job_name = row[id_col]
        protein_seq = row[seq_col]
        msa_dir = row[msa_col]
        SMILES = row[smi_col] if use_smiles else None
        CCD = None if use_smiles else row[ccd_col]
        return get_single_job(job_name, protein_seq, msa_dir, SMILES=SMILES, CCD=CCD)

    all_jobs = df.apply(build_job, axis=1).tolist()

    if save_json:
        save_path = Path(save_json)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        with save_path.open("w") as f:
            json.dump(all_jobs, f, indent=4)
        print(f"JSON saved to {save_path}")

    return all_jobs


# %% ../../nbs/protenix/07_protenix.ipynb 19
def get_virtual_screening_json(df, 
                               protein_seq, 
                               msa_dir, 
                               id_col,
                               smi_col=None, 
                               ccd_col=None, 
                               save_json=None):
    "Get json file of single protein against multiple SMILES in a dataframe."
    if smi_col and ccd_col:
        raise ValueError("Provide only one of smi_col or ccd_col, not both.")
    if not smi_col and not ccd_col:
        raise ValueError("You must provide either smi_col or ccd_col.")

    use_smiles = smi_col is not None

    def build_job(row):
        job_name = row[id_col]
        SMILES = row[smi_col] if use_smiles else None
        CCD = None if use_smiles else row[ccd_col]
        return get_single_job(job_name, protein_seq, msa_dir, SMILES=SMILES, CCD=CCD)

    all_jobs = df.apply(build_job, axis=1).tolist()

    if save_json:
        save_path = Path(save_json)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        with save_path.open("w") as f:
            json.dump(all_jobs, f, indent=4)
        print(f"JSON saved to {save_path}")

    return all_jobs
