from collections import defaultdict

from rdkit import Chem

from Core.MolecularFormula import MolecularFormula
from Main import Constant
from Main.AtomMass import atom_mass, ionized_mass
from Main.Parameter import Parameter


class Molecule(object):
    def __init__(self):
        self.mol = None
        self._formula = MolecularFormula()

    def get_ionized_mass(self):
        mass = self.get_neutral_mass() + ionized_mass[Parameter.addition]
        return mass

    def get_neutral_mass(self):
        return self._formula.get_mass()

    def get_formula(self):
        return self._formula

    def get_smiles(self):
        return Chem.MolToSmiles(self.mol)

    def get_inchi(self):
        return Chem.MolToInchi(self.mol)

    def _set_mol_from_rdkit(self):
        if not self.mol:
            return False

        # Check for multiple fragments
        frags = Chem.GetMolFrags(self.mol)
        if len(frags) > 1:
            return False

        # Remove H
        self.mol = Chem.RemoveHs(self.mol)

        # if Parameter.use_mmff:
        Chem.Kekulize(self.mol)

        # Check element and generate _formula
        ele_num = defaultdict(lambda: 0)
        for atom in self.mol.GetAtoms():
            atom_symbol = atom.GetSymbol()
            if atom_symbol not in atom_mass:
                # print atom_symbol, 'is not accepted.'
                return False

            self._formula[atom_symbol] += 1
            self._formula['H'] += atom.GetTotalNumHs()
            ele_num[atom_symbol] += 1
            ele_num['H'] += atom.GetTotalNumHs()

        for value in ele_num.values():
            if value > Constant.formula_max_atom_num:
                return False

        return True

    @classmethod
    def simplify_inchi(cls, raw_inchi: str):
        raw_inchi = raw_inchi.strip()
        inchi_list = raw_inchi.split('/')
        new_inchi_list = inchi_list[:2]
        for cont in inchi_list:
            if cont[0] == 'c':
                new_inchi_list.append(cont)
            elif cont[0] == 'h':
                new_inchi_list.append(cont)
        new_inchi = '/'.join(new_inchi_list)
        # new_inchi = new_inchi.encode('ascii', 'ignore')
        return new_inchi

    def set_mol_by_identifier(self, identifier):
        if type(identifier) is bytes:
            identifier = identifier.decode()
        identifier = identifier.strip()

        if not identifier.startswith('InChI='):
            self.mol = Chem.MolFromSmiles(identifier)
            if self.mol:
                identifier = Chem.MolToInchi(self.mol)
            else:
                return False

        self.mol = Chem.MolFromInchi(self.simplify_inchi(identifier))
        result = self._set_mol_from_rdkit()
        if not result:
            self.mol = None
        return result

    pass
