# Add a drug ligands to an AlphaFold 3 input json file making a separate input for each drug.
# The drugs are specified by a smiles string.  Name the json input file to match the drug.

def make_drug_json_input_files(receptor_json_path, drugs_path, drug_chain_id = 'D'):
    import json
    with open(receptor_json_path, 'r') as f:
        rj = json.load(f)

    with open(drugs_path, 'r') as f:
        drugs = f.readlines()

    ligand_spec = {"ligand": {"id": [drug_chain_id]}}
    rj['sequences'].append(ligand_spec)
    for line in drugs:
        drug_name, smiles = line.split(',')
        name = drug_name.replace(' ', '_')
        rj['name'] = name
        ligand_spec['ligand']['smiles'] = smiles.strip()
        with open(name + '.json', 'w') as f:
            json.dump(rj, f)

import sys
if len(sys.argv) == 3:
    make_drug_json_input_files(sys.argv[1], sys.argv[2])
else:
    print('Syntax: python3 af3_drug_json_input.py <receptor_json_path> <drug_smiles_path>')