import re

import joblib
import numpy as np

from Main import Constant
from Main.AtomMass import atom_dict, atom_list, atom_mass_array, len_atom_dict, ionized_mass, isotopic_info
from Main.Parameter import Parameter


def merge_formula(all_formulas):
    all_datas = [x.get_data() for x in all_formulas if x is not None]
    if all_datas:
        new_mol = MolecularFormula(np.sum(all_datas, axis=0))
        return new_mol
    else:
        return None


def generate_formula(all_data):
    result = []
    for data in all_data:
        mol = MolecularFormula(data)
        result.append(mol)
    return result


class MolecularFormula(object):
    __slots__ = ['_data', '_hash']

    def __init__(self, data=None):
        if data is not None:
            self._data = np.array(data, Constant.numpy_formula_format, copy=True)
        else:
            self._data = np.zeros(len_atom_dict, Constant.numpy_formula_format)
        self._hash = None
        pass

    def __getstate__(self):
        return [self._data, self._hash]

    def __setstate__(self, state):
        self._data = state[0]
        self._hash = state[1]

    def __iter__(self):
        for atom in atom_dict:
            yield atom

    def __hash__(self):
        if self._hash is None:
            self._hash = int(joblib.hash(self._data.tobytes()), 16) & (2 ** 64 - 1)
        return self._hash

    def __add__(self, other):
        if other:
            result_data = self._data + other._data
            result = MolecularFormula()
            result._data = result_data
            return result
        else:
            result = MolecularFormula()
            result._data = np.copy(self._data)
            return result

    def __getitem__(self, item):
        return self._data[atom_dict[item]]

    def __setitem__(self, key, value):
        self._data[atom_dict[key]] = value

    def __eq__(self, other):
        return np.array_equal(self._data, other._data)

    def __str__(self):
        string = ''

        for atom in atom_list:
            atom_num = self[atom]
            if atom_num:
                string += atom + str(atom_num)
        return string

    def __repr__(self):
        return self.__str__()

    def refresh_hash(self):
        self._hash = int(joblib.hash(self._data.tobytes()), 16) & (2 ** 64 - 1)

    def sub_not_allow_zero(self, other):
        if other:
            result_data = self._data - other._data
            if np.sum(result_data < 0) > 0:
                return None
            if np.sum(result_data) == 0:
                return None

            result = MolecularFormula(result_data)
            return result
        else:
            result = MolecularFormula(self._data)
            return result

    def sub_allow_zero(self, other):
        if other:
            result_data = self._data - other._data
            if np.sum(result_data < 0) > 0:
                return None

            result = MolecularFormula(result_data)
            return result
        else:
            result = MolecularFormula(self._data)
            return result

    def get_mass_after_shift(self, shift_data):
        result_data = shift_data + self._data
        return np.sum(result_data * atom_mass_array, axis=1)

    def get_data(self):
        return self._data

    def from_string(self, key_string):
        all_atom_nums = re.findall('([a-zA-Z]+)([0-9]+)', key_string)
        for atom_num in all_atom_nums:
            self[atom_num[0]] = int(atom_num[1])
        pass

    def get(self, atom, default_value):
        if atom in atom_dict:
            return self[atom]
        return default_value

    def get_mass(self):
        return np.sum(atom_mass_array * self._data)

    def get_ionized_mass(self):
        return self.get_mass() + ionized_mass[Parameter.addition]

    def get_isotopic_info(self, max_relative_abundance=1.5):
        info = []
        for iso in [0, 1]:
            delta_mass = isotopic_info['mass'][iso]
            delta_abundance = self._data * isotopic_info['abundance'][iso]

            iso_abundance = np.sum(delta_abundance * delta_mass)
            sum_delta_abundance = np.sum(delta_abundance)
            if sum_delta_abundance > 0:
                iso_delta_mass = iso_abundance / np.sum(delta_abundance)
                iso_mass = self.get_ionized_mass() + iso_delta_mass

                info.append((iso_mass, iso_abundance * max_relative_abundance))
            else:
                info.append((0., 0.))

        return info

    def is_empty(self):
        return np.count_nonzero(self._data) == 0

    def contains(self, other):
        result_data = self._data - other._data
        if np.sum(result_data < 0) > 0:
            return False
        return True

    pass
