from rdkit import Chem
from rdkit.Chem import AllChem
import argparse, os, gzip
import glob
import numpy as np
import rdkit
from rdkit import Chem
from rdkit.Chem import AllChem, rdShapeHelpers
from rdkit.Chem.FeatMaps import FeatMaps
from rdkit import RDConfig
from rdkit.Chem import Draw
from joblib import Parallel, delayed
import multiprocessing
import operator

fdef = AllChem.BuildFeatureFactory(os.path.join(RDConfig.RDDataDir, 'BaseFeatures.fdef'))

fmParams = {}
for k in fdef.GetFeatureFamilies():
    fparams = FeatMaps.FeatMapParams()
    fmParams[k] = fparams

keep = ('Donor', 'Acceptor', 'NegIonizable', 'PosIonizable', 'ZnBinder',
        'Aromatic', 'Hydrophobe', 'LumpedHydrophobe')


def get_FeatureMapScore(small_m, large_m, score_mode=FeatMaps.FeatMapScoreMode.All):
    featLists = []
    for m in [small_m, large_m]:
        rawFeats = fdef.GetFeaturesForMol(m)
        # filter that list down to only include the ones we're intereted in
        featLists.append([f for f in rawFeats if f.GetFamily() in keep])
    fms = [FeatMaps.FeatMap(feats=x, weights=[1] * len(x), params=fmParams) for x in featLists]
    fms[0].scoreMode = score_mode
    fm_score = fms[0].ScoreFeats(featLists[1]) / min(fms[0].GetNumFeatures(), len(featLists[1]))
    return fm_score


def score(reflig, prb_mols, ids, score_mode=FeatMaps.FeatMapScoreMode.All, p=False):
    ref = Chem.AddHs(reflig)
    idx = 0

    results_sucos = {}
    results_tani = {}

    smi_mol = Chem.MolToSmiles(prb_mols)

    for i in ids:

        prb = Chem.AddHs(Chem.MolFromMolBlock(Chem.MolToMolBlock(elab_mol, confId=i)))

        fm_score = get_FeatureMapScore(ref, prb, score_mode)
        fm_score = np.clip(fm_score, 0, 1)

        protrude_dist = rdShapeHelpers.ShapeProtrudeDist(ref, prb,
                                                         allowReordering=False)
        protrude_dist = np.clip(protrude_dist, 0, 1)

        SuCOS_score = 0.5 * fm_score + 0.5 * (1 - protrude_dist)
        tanimoto_score = Chem.rdShapeHelpers.ShapeTanimotoDist(ref, prb)

        results_sucos[str(idx)] = SuCOS_score
        results_tani[str(idx)] = tanimoto_score

        if p:
            print("********************************")
            print("index: " + str(idx))
            print("SuCOS score:\t%f" % SuCOS_score)
            print("Tani score:\t%f" % tanimoto_score)
            print("********************************")

        idx += 1

    return results_sucos


def get_best_align(hit_mblock, elab_smiles):
    hit_mol = Chem.MolFromMolBlock(hit_mblock)
    elab_mol = Chem.MolFromSmiles(elab_smiles)
    ids = AllChem.EmbedMultipleConfs(elab_mol, numConfs=100, params=AllChem.ETKDG())
    for cid in ids: AllChem.MMFFOptimizeMolecule(elab_mol, confId=cid)
    results_sucos = score(hit_mol, elab_mol, ids)
    best_i = list(results_sucos.values()).index(max(results_sucos.values()))
    elab_molblock = Chem.MolToMolBlock(elab_mol, confId=best_i)

    return elab_molblock


def get_core_mol(three_d_mol, core_mol):
    repl_sidechains = AllChem.ReplaceSidechains(three_d_mol, core_mol)
    if repl_sidechains:
        return AllChem.DeleteSubstructs(repl_sidechains, Chem.MolFromSmiles("*"))


def gen_conf_from_vector(input_mol_block, elaborated_smiles):
    # Get the mol
    get_best_align(input_mol_block, elaborated_smiles)


def generate_confs_for_vector(input_smiles, input_mol_block):
    confs = []
    for elaborated_smiles in input_smiles:
        confs.append(
            gen_conf_from_vector(input_mol_block, elaborated_smiles)
        )
    return confs


def generate_confs_for_vectors(input_vectors, input_mol_block):
    confs = {}
    for vector in input_vectors:
        confs[vector] = generate_confs_for_vector(
            vector, input_vectors[vector], input_mol_block
        )
    return confs


if __name__ == "__main__":
    generate_confs_for_vectors(
        {"O=C(Nc1ccc([Xe])nc1)N1CCNCC1_1_ADDITION": ["COc1ccc(NC(=O)N2CCNCC2)cn1"]},
        "\n     RDKit          3D\n\n 15 16  0  0  0  0  0  0  0  0999 V2000\n   -8.0100   -6.1770   18.7960 O   0  0  0  0  0  0  0  0  0  0  0  0\n   -7.8350   -7.3610   19.1060 C   0  0  0  0  0  0  0  0  0  0  0  0\n   -6.7970   -8.0980   18.5710 N   0  0  0  0  0  0  0  0  0  0  0  0\n   -6.0500   -7.8150   17.3890 C   0  0  0  0  0  0  0  0  0  0  0  0\n   -6.2220   -6.6910   16.5800 C   0  0  0  0  0  0  0  0  0  0  0  0\n   -5.5070   -6.4830   15.4790 N   0  0  0  0  0  0  0  0  0  0  0  0\n   -4.5740   -7.3730   15.1470 C   0  0  0  0  0  0  0  0  0  0  0  0\n   -4.3180   -8.4870   15.8900 C   0  0  0  0  0  0  0  0  0  0  0  0\n   -5.0700   -8.7240   17.0200 C   0  0  0  0  0  0  0  0  0  0  0  0\n   -8.6530   -7.9840   20.0040 N   0  0  0  0  0  0  0  0  0  0  0  0\n   -8.8660   -9.4290   20.1410 C   0  0  0  0  0  0  0  0  0  0  0  0\n  -10.3240   -9.8000   19.9120 C   0  0  0  0  0  0  0  0  0  0  0  0\n  -11.2290   -8.6780   20.2340 N   0  0  0  0  0  0  0  0  0  0  0  0\n  -10.7100   -7.8960   21.3670 C   0  0  0  0  0  0  0  0  0  0  0  0\n   -9.4210   -7.2020   20.9820 C   0  0  0  0  0  0  0  0  0  0  0  0\n  1  2  2  0\n  2  3  1  0\n  2 10  1  0\n  3  4  1  0\n  4  5  2  0\n  4  9  1  0\n  5  6  1  0\n  6  7  2  0\n  7  8  1  0\n  8  9  2  0\n 10 11  1  0\n 10 15  1  0\n 11 12  1  0\n 12 13  1  0\n 13 14  1  0\n 14 15  1  0\nM  END\n",
    )
