import itertools

import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem

from Core.MolecularFormula import MolecularFormula
from Main import Constant
from Main.AtomMass import atom_dict, len_atom_dict


class FormulaFeatureGenerator(object):
    def __init__(self, mol):
        self._mol = Chem.RWMol(mol.mol)
        self._formula = mol.get_formula()

        self._mmff = []
        self._all_possible_combination = {}
        self._known_formula_feature_dict = {}
        self._known_formula_feature = []

        try:
            self._generate_mmff()
            bond_info_to_mmff = self._generate_bond_info_mmff()
            all_mmff_formula_shift_atom_num = self._generate_mass_shift(bond_info_to_mmff)
            self._find_all_possible_combination(all_mmff_formula_shift_atom_num)
        except AssertionError:
            pass

        pass

    @staticmethod
    def _calculate_shifted_formula(mol_peak_info, bond_info, peak, extract_name, formula):
        if extract_name not in peak:
            raise Exception()

        extract_place = Constant.extract_type_place[extract_name]
        for parent_formula in peak['anno']:
            delta_formula = parent_formula[extract_place].sub_not_allow_zero(formula)
            if delta_formula:
                mol_peak_info[extract_name].add((bond_info, delta_formula))
        return mol_peak_info

    def convert_to_formula_feature(self, spectrum_info):
        for peak in spectrum_info:
            peak['ff'] = []
            for formula_num, peak_formula_list in enumerate(peak['anno']):
                peak['ff'] += self._convert_formula_to_formula_feature(peak_formula_list)
                pass
            pass
        pass

    def _convert_formula_to_formula_feature(self, formula):
        possible_result = []

        formula_feature_remain = (0, formula[0])
        formula_feature_lost = (0, formula[1])
        self._add_formula_feature_to_list(formula_feature_remain, formula_feature_lost, possible_result)
        for mmff_formula_shift_lost in self._all_possible_combination:
            delta_formula_lost = formula[1].sub_not_allow_zero(mmff_formula_shift_lost[1])
            if delta_formula_lost is not None:
                formula_feature_lost = (mmff_formula_shift_lost[0], delta_formula_lost)

                self._add_formula_feature_to_list(formula_feature_remain, formula_feature_lost, possible_result)
            pass

        for mmff_formula_shift_remain in self._all_possible_combination:
            delta_formula_remain = formula[0].sub_not_allow_zero(mmff_formula_shift_remain[1])
            if delta_formula_remain is not None:
                formula_feature_remain = (mmff_formula_shift_remain[0], delta_formula_remain)

                formula_feature_lost = (0, formula[1])
                self._add_formula_feature_to_list(formula_feature_remain, formula_feature_lost, possible_result)
                for mmff_formula_shift_lost in self._all_possible_combination[mmff_formula_shift_remain]:
                    delta_formula_lost = formula[1].sub_not_allow_zero(mmff_formula_shift_lost[1])
                    if delta_formula_lost is not None:
                        formula_feature_lost = (mmff_formula_shift_lost[0], delta_formula_lost)

                        self._add_formula_feature_to_list(formula_feature_remain, formula_feature_lost, possible_result)
                    pass

        all_formula_feature_remain = {x[0] for x in possible_result}
        all_formula_feature_loss = {x[1] for x in possible_result}

        possible_result = [(0, x[0], x[1]) for x in all_formula_feature_remain] + \
                          [(1, x[0], x[1]) for x in all_formula_feature_loss]
        return possible_result

    def _add_formula_feature_to_list(self, formula_feature_remain, formula_feature_lost, possible_result):
        if formula_feature_remain in self._known_formula_feature_dict:
            formula_feature_remain = \
                self._known_formula_feature[self._known_formula_feature_dict[formula_feature_remain]]
        else:
            self._known_formula_feature_dict[formula_feature_remain] = len(self._known_formula_feature)
            self._known_formula_feature.append(formula_feature_remain)

        if formula_feature_lost in self._known_formula_feature_dict:
            formula_feature_lost = \
                self._known_formula_feature[self._known_formula_feature_dict[formula_feature_lost]]
        else:
            self._known_formula_feature_dict[formula_feature_lost] = len(self._known_formula_feature)
            self._known_formula_feature.append(formula_feature_lost)

        possible_result.append((formula_feature_remain, formula_feature_lost))
        pass

    def get_all_possible_combination(self):
        return self._all_possible_combination

    def find_all_possible_peaks(self, all_model_ff):
        all_possible_formula_feature = []
        for formula_feature in self._all_possible_combination:
            bond_info_formula = formula_feature[1]
            for model_ff in all_model_ff.get(formula_feature[0], []):
                if model_ff[0] == 0:
                    frag_formula = model_ff[1] + bond_info_formula

                    if self._formula.sub_not_allow_zero(frag_formula):
                        all_possible_formula_feature.append(
                            (frag_formula, model_ff[2], (model_ff[0], formula_feature[0], model_ff[1]),))
                    else:
                        pass
                else:
                    root_loss_formula = model_ff[1] + bond_info_formula
                    frag_formula = self._formula.sub_not_allow_zero(root_loss_formula)
                    if frag_formula is not None:
                        all_possible_formula_feature.append(
                            (frag_formula, model_ff[2], (model_ff[0], formula_feature[0], model_ff[1]),))
                    else:
                        pass
        return all_possible_formula_feature

    def _generate_mmff(self):
        # Generate MMFF
        mmff = AllChem.MMFFGetMoleculeProperties(self._mol)
        assert mmff, 'Error in generate MMFF atom type.'

        element_number = self._mol.GetNumAtoms()
        self._mmff = [0] * element_number
        for i in range(element_number):
            self._mmff[i] = mmff.GetMMFFAtomType(i)
        pass

    def _generate_bond_info_mmff(self):
        bond_info_to_mmff = {}
        # Generate one bond info
        for bond in self._mol.GetBonds():
            a_idx = bond.GetBeginAtomIdx()
            b_idx = bond.GetEndAtomIdx()

            mmff_a = self._mmff[a_idx]
            mmff_b = self._mmff[b_idx]
            if mmff_a > mmff_b:
                mmff_a, mmff_b, a_idx, b_idx = mmff_b, mmff_a, b_idx, a_idx

            bond_one_info = mmff_a * 256 + mmff_b

            if bond_one_info == 1 * 256 + 1:
                # Ignore C-C-C bond
                pass
            else:
                bond_info_to_mmff[(a_idx, b_idx)] = bond_one_info
                pass
            pass

        # Generate two bond info
        for atom in self._mol.GetAtoms():
            all_neighbors = list(atom.GetNeighbors())
            if len(all_neighbors) < 2:
                continue

            atom_idx = atom.GetIdx()
            mmff_atom = self._mmff[atom_idx]
            for neighbor_a, neighbor_b in itertools.combinations(all_neighbors, 2):
                a_idx = neighbor_a.GetIdx()
                b_idx = neighbor_b.GetIdx()

                # If any bond in these two bonds break can generate two fragment, ignore this two-bond.
                if (atom_idx, a_idx) not in bond_info_to_mmff \
                        and (a_idx, atom_idx) not in bond_info_to_mmff \
                        and (atom_idx, b_idx) not in bond_info_to_mmff \
                        and (b_idx, atom_idx) not in bond_info_to_mmff:
                    mmff_a = self._mmff[a_idx]
                    mmff_b = self._mmff[b_idx]

                    bond_a = mmff_atom * 256 + mmff_a if mmff_atom < mmff_a else mmff_a * 256 + mmff_atom
                    bond_b = mmff_atom * 256 + mmff_b if mmff_atom < mmff_b else mmff_b * 256 + mmff_atom

                    if bond_a > bond_b:
                        bond_a, bond_b, a_idx, b_idx = bond_b, bond_a, b_idx, a_idx

                    bond_two_info = bond_a * 256 * 256 + bond_b
                    bond_info_to_mmff[(atom_idx, a_idx, b_idx)] = bond_two_info
                else:
                    pass

        return bond_info_to_mmff

    def _generate_mass_shift(self, bond_info_to_mmff):
        # Generate mass for each heavy atom.
        atom_num = self._mol.GetNumAtoms()
        atom_list = np.zeros((atom_num, len_atom_dict), Constant.numpy_formula_format)
        for i in range(atom_num):
            atom = self._mol.GetAtomWithIdx(i)
            atom_list[i, atom_dict[atom.GetSymbol()]] = 1
            atom_list[i, atom_dict['H']] += atom.GetTotalNumHs()

        all_mmff_formula_shift_atom_num = []
        # Generate mass shift for one bond break.
        for bond in bond_info_to_mmff:
            if len(bond) == 2:
                bond_type = self._mol.GetBondBetweenAtoms(bond[0], bond[1]).GetBondType()
                self._mol.RemoveBond(bond[0], bond[1])
                all_frags = Chem.GetMolFrags(self._mol)

                self._mol.AddBond(bond[0], bond[1], bond_type)

            elif len(bond) == 3:
                bond_type_1 = self._mol.GetBondBetweenAtoms(bond[0], bond[1]).GetBondType()
                bond_type_2 = self._mol.GetBondBetweenAtoms(bond[0], bond[2]).GetBondType()
                self._mol.RemoveBond(bond[0], bond[1])
                self._mol.RemoveBond(bond[0], bond[2])
                all_frags = Chem.GetMolFrags(self._mol)

                self._mol.AddBond(bond[0], bond[1], bond_type_1)
                self._mol.AddBond(bond[0], bond[2], bond_type_2)

            else:
                pass

            if len(all_frags) == 2:
                mass_shift_1 = MolecularFormula(np.sum(atom_list[list(all_frags[0])], axis=0))
                mass_shift_2 = MolecularFormula(np.sum(atom_list[list(all_frags[1])], axis=0))

                all_mmff_formula_shift_atom_num.append(((bond_info_to_mmff[bond], mass_shift_1), set(all_frags[0])))
                all_mmff_formula_shift_atom_num.append(((bond_info_to_mmff[bond], mass_shift_2), set(all_frags[1])))
            pass
        return all_mmff_formula_shift_atom_num

    def _find_all_possible_combination(self, all_mmff_formula_shift_atom_num):
        for i, item_i in enumerate(all_mmff_formula_shift_atom_num):
            for j, item_j in enumerate(all_mmff_formula_shift_atom_num[(i + 1):]):
                if item_i[1] & item_j[1]:
                    pass
                else:
                    new_i = item_i[0]
                    new_j = item_j[0]
                    if new_i in self._all_possible_combination:
                        self._all_possible_combination[new_i].add(new_j)
                    else:
                        self._all_possible_combination[new_i] = {new_j}

                    if new_j in self._all_possible_combination:
                        self._all_possible_combination[new_j].add(new_i)
                    else:
                        self._all_possible_combination[new_j] = {new_i}
                pass
        pass


class SimpleFormulaFeatureGenerator(object):
    def __init__(self, mol):
        pass

    def convert_to_formula_feature(self, spectrum_info):
        for peak in spectrum_info:
            peak['ff'] = []
            for formula_num, peak_formula_list in enumerate(peak['anno']):
                peak['ff'] += self._convert_formula_to_formula_feature(peak_formula_list)
                pass
            pass
        pass

    def _convert_formula_to_formula_feature(self, formula):
        possible_result = [(0, 0, formula[0]),
                           (1, 0, formula[1]), ]
        return possible_result
