import itertools

import numpy as np

from Core.MolecularFormula import MolecularFormula
from Main.AtomMass import atom_mass_array, ionized_mass
from Main.Parameter import Parameter


def mass_to_formula(mass, mass_error, precursor_formula):
    mass -= ionized_mass[Parameter.addition]

    lo_mass = mass - mass_error
    hi_mass = mass + mass_error

    # Generate candidate range
    precursor_data = precursor_formula.get_data()
    formula_range = [range(x + 1) for x in precursor_data]
    all_possible_candidate_formula = np.array(list(itertools.product(*formula_range)), np.int8)
    all_possible_mass = np.sum(atom_mass_array * all_possible_candidate_formula, axis=1)

    candidate_data = all_possible_candidate_formula[(lo_mass <= all_possible_mass) & (all_possible_mass <= hi_mass)]

    result = []
    for data in candidate_data:
        formula = MolecularFormula(data)
        result.append(formula)
    return result


def test_unsaturation(formula):
    unsaturation = _get_unsaturation(formula)
    return unsaturation > -1e-4


def _get_unsaturation(formula):
    n_c = formula['C']
    n_h = formula['H']
    n_n = formula['N']
    n_p = formula['P']

    n_x = formula['F']
    n_x += formula['Br']
    n_x += formula['Cl']
    n_x += formula['I']
    unsaturation = (1.0 * (2.0 * n_c + n_n + n_p - n_h - n_x) / 2.0) + 1.0
    return unsaturation
