from rdkit import Chem, RDLogger
import csv
import requests, time
import json
from tqdm import tqdm

RDLogger.DisableLog("rdApp.*")

def chebi_to_smiles():
    csmiles_dict = {}
    for m in tqdm(Chem.SDMolSupplier("data/chebi_lite.sdf")):
        if m is None:
            continue
        csmiles_dict[m.GetProp("ChEBI ID")] = Chem.MolToSmiles(m)

    return csmiles_dict


def pdb_to_smiles():
    all_pdb = {}
    with open("data/pdb_smiles.smi", newline="") as f:
        reader = csv.DictReader(f, delimiter="\t")
    
        for row in tqdm(reader):
            all_pdb[row['pdb'].upper()] = row['smiles']
    return all_pdb

def kegg_to_smiles(kegg_id):
    #time.sleep(0.4) # KEGG asks for modest request rates
    r = requests.get(f"https://rest.kegg.jp/get/{kegg_id}/mol", timeout=30)
    if r.status_code != 200 or not r.text.strip():
        return None
    mol = Chem.MolFromMolBlock(r.text)
    return Chem.MolToSmiles(mol) if mol else None

def read_mcsa_csv(csmiles, psmiles):
    all_compounds = {}
    success_count = 0
    total_count = 0
    unsuccessful = []
    with open("data/mcsa-curated_data.csv", newline="") as f:
        reader = csv.DictReader(f)   # each row becomes a dict keyed by column name
        
        for row in reader:
            total_count+=1
            compound_type = row["residue/reactant/product/cofactor"]
            chain_kegg = row["chain/kegg compound"]
            resid_chebi = row['resid/chebi id']
            pdbcode = row['PDB code'].upper()

            # assume chebi id
            if resid_chebi and len(resid_chebi)>3 and resid_chebi not in all_compounds:
                if f"CHEBI:{resid_chebi}" in csmiles:
                    smiles = csmiles[f"CHEBI:{resid_chebi}"]
                    all_compounds[resid_chebi] = {'chebi': resid_chebi, 'type': compound_type, 'smiles': smiles}
                    success_count+=1
                    print(f"OK:{resid_chebi}")
                    continue

            if pdbcode and pdbcode not in all_compounds:
                if pdbcode in psmiles:
                    smiles = psmiles[pdbcode]
                    all_compounds[pdbcode] = {'pdb': pdbcode, 'type': compound_type, 'smiles': smiles}
                    success_count+=1
                    print(f"OK:{pdbcode}")
                    continue

            if chain_kegg and len(chain_kegg)>5 and chain_kegg not in all_compounds:
                smiles = kegg_to_smiles(chain_kegg)
                if smiles:
                    all_compounds[chain_kegg] = {'kegg': chain_kegg, 'type': compound_type, 'smiles': smiles}
                    success_count+=1 
                    print(f"OK:{chain_kegg}")
                    continue

            if resid_chebi not in all_compounds and pdbcode not in all_compounds and chain_kegg not in all_compounds:
                print(f"NOT OK:{resid_chebi}-{pdbcode}-{chain_kegg}")
                mcsaid = row['M-CSA ID']
                uniprot = row['Uniprot IDs']
                role = row['role']
                unsuccessful.append({'resid_chebi': resid_chebi, 'pdbcode': pdbcode, \
                    'chain_kegg': chain_kegg, 'type': compound_type, \
                    'mcsa_id': mcsaid, 'uniprot id': uniprot, 'role': role})

    return (all_compounds, unsuccessful, success_count, total_count)

def main():
    # read in M-CSA csv to get all compounds used
    print("Creating CHEBI lookup table...")
    csmiles = chebi_to_smiles()
    print("Creating PDB lookup table...")
    psmiles = pdb_to_smiles()
    print("Reading M-CSA compounds...")
    (all_compounds, unsuccessful, success_count, total_count) = read_mcsa_csv(csmiles, psmiles)

    print("Done.\n")
    print("Printing to output/mcsa-smiles.json")
    with open("output/mcsa-smiles.json", "w") as f:
        json.dump(all_compounds, f) 

    print("\nPrinting unsuccessfully parsed compounds...")
    with open("output/unsuccessful.json", 'w') as f:
        json.dump(unsuccessful, f)
    
    print(f"{success_count}/{total_count} compound(s) successfully read.")


if __name__=="__main__":
    main()