import pickle
import sqlite3

import numpy as np

from Core.Molecule import Molecule
from Core.Spectrum import Spectrum
from Main import AtomMass
from Main.Parameter import Parameter
from Model.AnnotateMolecule import annotate_spectrum
from Utils.GenerateMolecule import get_mol_from_mol_str
from Utils.Multiplecore import run_function


def _generate_spec_annotation_info(mol_info):
    mol_identifier = mol_info['identifier'].strip()
    mol = get_mol_from_mol_str(mol_identifier)
    if not mol:
        return None

    mol_mass = mol.get_ionized_mass()
    if mol_info['spec'].precursor_mass \
            and (abs(mol_info['spec'].precursor_mass - mol_mass) > 0.5):
        if not Parameter.not_guess_adduct:
            # The precursor mass and molecule mass not match.
            # Try other adduct mass
            adduct_mass = AtomMass.adduct_mass[Parameter.ion_charge]
            for adduct_type in adduct_mass:
                cur_add_mass = adduct_mass[adduct_type]
                cur_corrected_mass = AtomMass.ionized_mass[Parameter.addition] - cur_add_mass

                cur_corrected_precursor_mass = mol_info['spec'].precursor_mass + cur_corrected_mass
                if abs(cur_corrected_precursor_mass - mol_mass) < 0.5:
                    # Correct mass
                    for peak in mol_info['spec']:
                        peak['mass'] += cur_corrected_mass
                    mol_info['spec'].precursor_mass += cur_corrected_mass
                    break
                pass
            else:
                print('Mass not match:', mol_info['spec'].precursor_mass, mol.get_neutral_mass(),
                      mol_info['identifier'])
                return None
        else:
            print('Mass not match:', mol_info['spec'].precursor_mass, mol.get_neutral_mass(), mol_info['identifier'])
            return None

    mol_spec_info = annotate_spectrum(mol, mol_info['spec'], allow_null_spec=True)
    if not mol_spec_info:
        return None

    mol_spec_info['name'] = mol_info['name']
    mol_spec_info['smiles'] = mol.get_smiles()
    mol_spec_info['inchi'] = mol.get_inchi()
    return mol_spec_info


def _score_mol_info(model, all_mol_fingerprints, all_mol_features, all_mol_intensity, feature_intensity_matrix):
    print('Start fitting fingerprint model')
    all_fp_score = model.score_mol_info_with_fp_model(
        feature_intensity_matrix,
        all_mol_fingerprints)

    print('Start fitting peak model')
    all_peak_score = model.score_mol_info_with_peak_model(
        all_mol_intensity,
        all_mol_features,
        all_mol_fingerprints)

    return all_peak_score, all_fp_score


def get_spectrum_info(filename_input, precursor_mz=None):
    all_spec = []
    if filename_input[-4:] == '.mgf':
        mgf_str = open(filename_input).readlines()
        peak_start = False
        for line in mgf_str:
            line = line.strip()
            if line.startswith('BEGIN IONS'):
                peak_start = True
            if line.startswith('PEPMASS='):
                precursor_mz = float(line.split('=')[1].strip())

            if peak_start and line.find('=') < 0 and line.find('ION') < 0 and len(line.split()) > 1:
                items = line.split()
                all_spec.append((float(items[0]), float(items[1])))
    else:
        ori_str = open(filename_input).readlines()
        for line in ori_str:
            items = line.split()
            all_spec.append((float(items[0]), float(items[1])))

    spectrum = Spectrum(all_spec)
    spectrum.precursor_mass = precursor_mz
    return spectrum


def get_mol_info_from_db(filename_db, precursor_mz, ppm):
    precursor_mass = precursor_mz - AtomMass.ionized_mass[Parameter.addition]
    mass_min = precursor_mass - precursor_mass / 1e6 * ppm
    mass_max = precursor_mass + precursor_mass / 1e6 * ppm

    db = sqlite3.connect(filename_db)
    cur = db.cursor()
    sql = "SELECT Molecules.DBID,Molecules.InChI FROM Molecules " \
          "WHERE ({}<=Molecules.Mass) AND (Molecules.Mass<={})".format(mass_min, mass_max)
    cur.execute(sql)
    result = []
    for cur_info in cur.fetchall():
        mol_info = {
            'name': cur_info[0],
            'identifier': cur_info[1],
        }
        result.append(mol_info)
    return result


def get_mol_info(filename_candidates):
    cur_spec_all_mol = []
    # Get smile for each mol.
    fi = open(filename_candidates)
    for line in fi.readlines():
        cont = line.split()
        if len(cont) == 1:
            mol_name = cont[0].strip()
            mol_identifier = cont[0].strip()
        elif len(cont) == 2:
            mol_name = cont[0].strip()
            mol_identifier = cont[1].strip()
        else:
            continue
        mol_info = {
            'name': mol_name,
            'identifier': mol_identifier,
        }
        cur_spec_all_mol.append(mol_info)
    return cur_spec_all_mol


def annotate_spectra(all_raw_spec_info):
    all_info = run_function(_generate_spec_annotation_info, [(x,) for x in all_raw_spec_info],
                            threads_type=3, chunksize=1)

    error_info = [y for x, y in zip(all_info, all_raw_spec_info) if x is None]
    right_info = [x for x in all_info if x is not None]
    return right_info, error_info


def extract_info(annotated_spec_info, model):
    result = {'name': [x['name'] for x in annotated_spec_info],
              'smiles': [x['smiles'] for x in annotated_spec_info],
              'inchi': [x['inchi'] for x in annotated_spec_info], }

    feature_intensity_matrix = model.convert_mol_info_to_feature_intensity_matrix(annotated_spec_info)
    all_mol_intensity, all_mol_features = model.convert_mol_info_to_peak_intensity_features(annotated_spec_info)
    all_mol_fingerprints = np.array([x['fingerprint'] for x in annotated_spec_info], np.int8)

    result.update({
        'fingerprint': all_mol_fingerprints,
        'feature': all_mol_features,
        'spec_intensity': all_mol_intensity,
        'feature_intensity': feature_intensity_matrix,
    })
    return result


def score_with_peak_model(all_mol_info, model, filename_saved_result=None):
    all_mol_intensity = all_mol_info['spec_intensity']
    all_mol_features = all_mol_info['feature']
    all_mol_fingerprints = all_mol_info['fingerprint']

    print('Start scoring with peak model')
    all_peak_score = model.score_mol_info_with_peak_model(
        all_mol_intensity,
        all_mol_features,
        all_mol_fingerprints)

    all_mol_info['peak_score'] = all_peak_score
    if filename_saved_result is not None:
        all_peak_score = np.array(all_peak_score, np.float64)
        pickle.dump(all_peak_score, open(filename_saved_result, 'w'))
    return all_peak_score


def _correct_with_adduct(mol_mass, spec_info):
    if abs(spec_info.precursor_mass - mol_mass) < 0.5:
        return None

    if Parameter.not_guess_adduct:
        return False

    adduct_mass = AtomMass.adduct_mass[Parameter.ion_charge]
    for adduct_type in adduct_mass:
        cur_add_mass = adduct_mass[adduct_type]
        cur_corrected_mass = AtomMass.ionized_mass[Parameter.addition] - cur_add_mass

        cur_corrected_precursor_mass = spec_info.precursor_mass + cur_corrected_mass
        if abs(cur_corrected_precursor_mass - mol_mass) < 0.5:
            # Correct mass
            for peak in spec_info:
                peak['mass'] += cur_corrected_mass
                spec_info.precursor_mass += cur_corrected_mass
            return True
    return False


def score_with_spectra_p(all_mol_info):
    if len(all_mol_info) == 0:
        return []
    score = []
    for mol_info in all_mol_info:
        if mol_info['spec'].precursor_mass:
            mol = Molecule()
            mol_result = mol.set_mol_by_identifier(mol_info['identifier'])
            if not mol_result:
                score.append(0.)
                continue

            mol_mass = mol.get_ionized_mass()
            guess_adduct_result = _correct_with_adduct(mol_mass, mol_info['spec'])
            if guess_adduct_result is None:
                pass
            elif guess_adduct_result:
                pass
            else:
                score.append(0.)
                continue
            pass

        spec_real, spec_predict = mol_info['spec'], mol_info['spec_p']

        precursor_intensity = [x['intensity'] for x in spec_real if x['mass'] <= spec_real.precursor_mass - 0.5]
        precursor_intensity_sum = sum(precursor_intensity)

        if precursor_intensity_sum == 0.:
            score.append(0.)

        cur_score = 0.
        if spec_predict is not None:
            for peak in spec_real:
                if peak['mass'] > spec_real.precursor_mass - 0.5:
                    continue
                mass_delta = peak['mass'] * Parameter.ms2_ppm * 1e-6
                mass_min = peak['mass'] - mass_delta
                mass_max = peak['mass'] + mass_delta
                select_info = np.bitwise_and(spec_predict[:, 0] >= mass_min, spec_predict[:, 0] <= mass_max)
                select_peak = spec_predict[select_info, 1]
                if select_peak.any():
                    p = np.max(select_peak)
                    cur_score += peak['intensity'] * p
                pass
        score.append(cur_score / precursor_intensity_sum)

    return score


def output_result(df, filename_output):
    df = df.sort_values(by='score', ascending=False)
    if df['score'].max() == 0.:
        df['score'] += 1.

    if Parameter.casmi2017:
        df = df[['name', 'score']]
        df.to_csv(filename_output, sep="\t", index=False, header=False)
    else:
        df = df[['name', 'score', 'identifier']]
        df.to_csv(filename_output, sep="\t", index=False, header=True)
    pass
