import sys

import numpy as np
from rdkit.Chem.Fingerprints import FingerprintMols
from rdkit.Chem.rdMolDescriptors import *

from Main.Parameter import Parameter


def generate_fingerprint(mol):
    fingerprint = []
    for cur_fp_class_str, cur_fp_bits in Parameter.fp_type_bits:
        wanted_class = getattr(sys.modules[__name__], cur_fp_class_str)
        mol_to_fp = wanted_class(fp_bits=cur_fp_bits)
        cur_fingerprint = mol_to_fp.get_fp_for_mol(mol.mol)
        fingerprint.append(cur_fingerprint)
    fingerprint = np.concatenate(fingerprint)
    return fingerprint


class FingerPrint(object):
    def __init__(self, fp_bits):
        self._fp_bits = fp_bits


class MolToRDKitFP(FingerPrint):
    def __init__(self, *args, **kwargs):
        super(MolToRDKitFP, self).__init__(*args, **kwargs)
        pass

    def get_fp_for_mol(self, mol):
        fp = list(FingerprintMols.FingerprintMol(mol, minPath=1, maxPath=8, fpSize=self._fp_bits, bitsPerHash=2,
                                                 useHs=False, tgtDensity=0.3, minSize=self._fp_bits))
        fp_result = np.array(fp, dtype=np.int8)
        return fp_result

    pass


class MolToMorganFP(FingerPrint):
    def __init__(self, *args, **kwargs):
        super(MolToMorganFP, self).__init__(*args, **kwargs)

    def get_fp_for_mol(self, mol):
        fp = list(GetMorganFingerprintAsBitVect(mol, radius=2, nBits=self._fp_bits))
        fp_result = np.array(fp, dtype=np.int8)
        return fp_result

    pass


class MolToHashedTopologicalTorsionFP(FingerPrint):
    def __init__(self, *args, **kwargs):
        super(MolToHashedTopologicalTorsionFP, self).__init__(*args, **kwargs)

    def get_fp_for_mol(self, mol):
        fp = list(GetHashedTopologicalTorsionFingerprintAsBitVect(mol, nBits=self._fp_bits))
        fp_result = np.array(fp, dtype=np.int8)
        return fp_result

    pass


class MolToHashedAtomPairFP(FingerPrint):
    def __init__(self, *args, **kwargs):
        super(MolToHashedAtomPairFP, self).__init__(*args, **kwargs)

    def get_fp_for_mol(self, mol):
        fp = list(GetHashedAtomPairFingerprintAsBitVect(mol, nBits=self._fp_bits))
        fp_result = np.array(fp, dtype=np.int8)
        return fp_result

    pass


class MolToAvalonFP(FingerPrint):
    def __init__(self, *args, **kwargs):
        super(MolToAvalonFP, self).__init__(*args, **kwargs)

    def get_fp_for_mol(self, mol):
        from rdkit.Avalon import pyAvalonTools
        fp = list(pyAvalonTools.GetAvalonFP(mol, nBits=self._fp_bits))
        fp_result = np.array(fp, dtype=np.int8)
        return fp_result

    pass
