"""Ident base parser class."""
import csv
import multiprocessing as mp

import ahocorasick
import numpy as np
import pandas as pd
from chemical_composition.chemical_composition_kb import PROTON
from loguru import logger
from peptide_mapper.mapper import UPeptideMapper


from pyiohat.parsers.base_parser import BaseParser
from pyiohat.parsers.misc import (
    get_compositions_and_monoisotopic_masses,
    get_isotopologue_accuracy,
    init_custom_cc,
    sort_mods,
)
from pyiohat.utils import merge_and_join_dicts

mp.set_start_method(method="fork")


class IdentBaseParser(BaseParser):
    """Base class of all ident parsers."""

    def __init__(self, *args, **kwargs):
        """Initialize parser.

        Reads in data file and provides mappings.
        """
        super().__init__(*args, **kwargs)
        self.DELIMITER = self.params.get("delimiter", "<|>")
        self.PROTON = PROTON
        self.IUPAC_AAS = tuple("ACDEFGHIKLMNPQRSTUVWY")
        self.df = None

        successfully_mapped_mods = set(
            mod["name"] for mod in self.params["mapped_mods"]["fix"]
        ) | set(mod["name"] for mod in self.params["mapped_mods"]["opt"])
        self.non_mappable_mods = set(
            mod["name"] for mod in self.params.get("modifications", [])
        ).difference(successfully_mapped_mods)
        self.mod_dict = self._create_mod_dicts()
        self.rt_truncate_precision = 2
        self.reference_dict = {
            "search_engine": None,
            "spectrum_id": None,
            "modifications": None,
            "retention_time_seconds": None,
        }
        self.required_headers = self._load_model("ident_parser_model.json")
        self.col_order = pd.Series(self.required_headers.keys())

    def _calc_mz(self, mass, charge):
        """Calculate mass-to-charge ratio.

        Args:
            mass (pd.Series): masses
            charge (pd.Series): charges
        Returns:
            (pd.Series): m/z
        """
        return (
            mass.astype(float) + (charge.astype(int) * self.PROTON)
        ) / charge.astype(int)

    def _calc_mass(self, mz, charge):
        """Calculate mass from mass-to-charge ratio.

        Args:
            mz (pd.Series): mass-to-charge
            charge (pd.Series): charges
        Returns:
            (pd.Series): mass
        """
        return mz.astype(float) * charge.astype(int) - (
            charge.astype(int) * self.PROTON
        )

    def clean_up_modifications(self):
        """Sanitizes modstrings generated by engine parsers.

        Modifications are sorted by position and leading, repeated or trailing delimiters are removed
        Operations are performed inplace on self.df
        """
        # Remove any trailing, leading, repeated, or only delimiters modstrings
        self.df.loc[:, "modifications"] = self.df.loc[:, "modifications"].str.replace(
            r"^;+(?=\w)|(?<=\w);+$|^;+$|;+(?=;)", "", regex=True
        )

        # Ensure same order of modifications
        self.df.loc[:, "modifications"] = (
            self.df["modifications"].fillna("").str.split(";").apply(sort_mods)
        )

    def assert_only_iupac_and_missing_aas(self):
        """Assert that only IUPAC nomenclature one letter amino acids are used in sequence.

        Non-IUPAC designations are dropped (except for selenocysteine).
        Operations are performed inplace.
        """
        self.df["sequence"] = self.df["sequence"].str.upper()
        # Added X for missing AAs
        iupac_conform_seqs = self.df["sequence"].apply(
            lambda seq: set(seq).issubset(self.IUPAC_AAS)
        )
        if any(~iupac_conform_seqs):
            self.df = self.df.loc[iupac_conform_seqs, :]
            logger.warning(
                f"sequences are not IUPAC conform. {(~iupac_conform_seqs).sum()} PSMs were dropped."
            )

    def add_protein_ids(self):
        """Add all Protein IDs that matching the sequence.

        Operations are performed inplace on self.df
        """
        peptide_mapper = UPeptideMapper(self.params["database"])
        mapped_peptides = peptide_mapper.map_peptides(self.df["sequence"].tolist())

        peptide_mappings = [
            merge_and_join_dicts(mapped_peptides[seq], self.DELIMITER)
            for seq in self.df["sequence"]
        ]

        columns_translations = {
            "start": "sequence_start",
            "end": "sequence_stop",
            "post": "sequence_post_aa",
            "id": "protein_id",
            "pre": "sequence_pre_aa",
        }
        new_columns = pd.DataFrame(peptide_mappings)
        new_columns.rename(columns=columns_translations, inplace=True)

        self.df.loc[:, new_columns.columns] = new_columns.values
        new_columns = new_columns.dropna(axis=0, how="all")
        if len(new_columns) != len(self.df):
            logger.warning(
                f"{len(self.df)-len(new_columns)} PSMs were dropped because their respective sequences could not be mapped."
            )
        self.df = self.df.iloc[new_columns.index, :].reset_index(drop=True)

    def check_enzyme_specificity(self):
        """Check consistency of N/C-terminal cleavage sites.

        Calculates number of missed cleavage sites.
        Operations are performed inplace.
        """
        if self.params["enzyme"] == ".^":
            self.df.loc[:, ["enzn", "enzc"]] = True
            self.df.loc[:, "missed_cleavages"] = 0
            return None

        enzyme_pattern = self.params["enzyme"]
        integrity_strictness = self.params["terminal_cleavage_site_integrity"]

        pren_seq = (
            pd.concat(
                [
                    self.df["sequence_pre_aa"].str.split(rf"{self.DELIMITER}"),
                    self.df["sequence"].str[:1],
                ],
                axis=1,
            )
            .explode("sequence_pre_aa")
            .sum(axis=1)
        )
        self.df.loc[:, "enzn"] = (
            pren_seq.str.split(rf"{enzyme_pattern}").str[0].str.len() == 1
        ).groupby(pren_seq.index).agg(integrity_strictness) | (
            pren_seq.str[0] == "-"
        ).groupby(
            pren_seq.index
        ).agg(
            integrity_strictness
        )
        postc_seq = (
            pd.concat(
                [
                    self.df["sequence"].str[-1:],
                    self.df["sequence_post_aa"].str.split("<\\|>"),
                ],
                axis=1,
            )
            .explode("sequence_post_aa")
            .sum(axis=1)
        )
        self.df.loc[:, "enzc"] = (
            postc_seq.str.split(rf"{enzyme_pattern}").str[0].str.len() == 1
        ).groupby(postc_seq.index).agg(integrity_strictness) | (
            postc_seq.str[-1] == "-"
        ).groupby(
            postc_seq.index
        ).agg(
            integrity_strictness
        )

        internal_cuts = self.df["sequence"].str.split(rf"{enzyme_pattern}")
        self.df.loc[:, "missed_cleavages"] = (
            internal_cuts.apply(len)
            - internal_cuts.apply(lambda row: "" in row).astype(int)
            - 1
        )

    def calc_masses_offsets_and_composition(self):
        """Calculate chemical composition theoretical masses, and mass-to-charge ratios.

        Offsets are calculated between theoretical and experimental mass-to-charge ratio.
        Operations are performed inplace on self.df
        """
        all_compositions = {}
        for aa in self.IUPAC_AAS:
            self.cc.use(sequence=aa)
            all_compositions[aa] = self.cc.copy()
        for mod in self.mod_dict.keys():
            all_compositions[mod] = self.mod_mapper.name_to_composition(mod)[0]

        compositions, mono_masses = get_compositions_and_monoisotopic_masses(
            sequences=self.df["sequence"].to_numpy(dtype=str),
            modifications=self.df["modifications"].to_numpy(dtype=str),
            compositions=all_compositions,
            isotopic_distributions=self.cc.isotopic_distributions,
        )
        self.df.loc[:, "chemical_composition"] = compositions
        self.df.loc[:, "ucalc_mass"] = mono_masses

        with mp.Pool(
            self.params.get("cpus", mp.cpu_count() - 1),
            initializer=init_custom_cc,
            initargs=(
                get_isotopologue_accuracy,
                self.PROTON,
            ),
        ) as pool:
            acc = pool.starmap(
                get_isotopologue_accuracy,
                zip(
                    self.df["chemical_composition"].values,
                    self.df["charge"].astype(int).values,
                    self.df["exp_mz"].values,
                ),
            )
        self.df.loc[:, "accuracy_ppm"] = acc
        self.df.loc[:, "ucalc_mz"] = self._calc_mz(
            mass=self.df["ucalc_mass"], charge=self.df["charge"]
        )
        self.df.loc[:, "accuracy_ppm_C12"] = (
            (self.df["exp_mz"].astype(float) - self.df["ucalc_mz"])
            / self.df["ucalc_mz"]
            * 1e6
        )
        # Clear rows with non-mappable mods
        for mod in self.non_mappable_mods:
            self.df.loc[
                self.df["modifications"].str.contains(mod),
                (
                    "chemical_composition",
                    "ucalc_mass",
                    "ucalc_mz",
                    "accuracy_ppm",
                    "accuracy_ppm_C12",
                ),
            ] = np.nan
        self.df.loc[:, "exp_mass"] = self._calc_mass(
            mz=self.df["exp_mz"], charge=self.df["charge"]
        )
        self.df.loc[:, "mass_delta"] = self.df["exp_mass"] - self.df["ucalc_mass"]

    def get_meta_info(self):
        """Extract meta information.

        Experimental mass-to-charge ratios, retention times, file names,
        and spectrum titles are added.
        Operations are performed inplace on self.df
        """
        rt_lookup = self._read_meta_info_lookup_file()
        self.df["spectrum_id"] = self.df["spectrum_id"].astype(int)
        if self.style in ("comet_style_1", "omssa_style_1"):
            logger.warning(
                "This engine does not provide retention time information. Grouping only by Spectrum ID. This may cause problems when working with multi-file inputs."
            )
            for name, grp in self.df.groupby("spectrum_id"):
                mappable_within_precision = list(rt_lookup[name].keys())
                if len(mappable_within_precision) == 1:
                    self.df.loc[
                        grp.index,
                        ("raw_data_location", "exp_mz", "retention_time_seconds"),
                    ] = rt_lookup[name][mappable_within_precision[0]] + [
                        mappable_within_precision[0]
                    ]
                else:
                    logger.error(
                        f"Could not uniquely assign meta data to spectrum id {name}."
                    )
        else:
            self.df["retention_time_seconds"] = self.df[
                "retention_time_seconds"
            ].astype(float)
            for name, grp in self.df.groupby(["spectrum_id", "retention_time_seconds"]):
                meta_rts = rt_lookup[name[0]].keys()
                mappable_within_precision = [
                    rt
                    for rt in meta_rts
                    if abs(name[1] - rt) <= 10**-self.rt_truncate_precision
                ]
                if len(mappable_within_precision) == 1:
                    self.df.loc[
                        grp.index,
                        ("raw_data_location", "exp_mz", "retention_time_seconds"),
                    ] = rt_lookup[name[0]][mappable_within_precision[0]] + [
                        mappable_within_precision[0]
                    ]
                else:
                    logger.error(
                        f"Could not uniquely assign meta data to spectrum id, retention time {name}."
                    )

        self.df.loc[:, "spectrum_title"] = (
            self.df["raw_data_location"]
            + "."
            + self.df["spectrum_id"].astype(str)
            + "."
            + self.df["spectrum_id"].astype(str)
            + "."
            + self.df["charge"].astype(str)
        )

    def add_ranks(self):
        """Ranks are calculated based on the engine scoring column at Spectrum ID level.

        Operations are performed inplace on self.df
        """
        eng_name = self.df["search_engine"].unique()[0]
        score_col = self.params["validation_score_field"][eng_name]
        top_is_highest = self.params["bigger_scores_better"][eng_name]
        ranking_needs_to_be_ascending = False if top_is_highest is True else True
        self.df.loc[:, score_col] = self.df[score_col].astype(float)
        self.df.loc[:, "rank"] = (
            self.df.groupby("spectrum_id")[score_col]
            .rank(ascending=ranking_needs_to_be_ascending, method="min")
            .astype(int)
        )

    def add_decoy_identity(self):
        """Add boolean decoy state if designated decoy prefix is in Protein IDs.

        Also marks peptides which were assigned decoy but were immutable during
        target decoy generation.
        Operations are performed inplace on self.df
        """
        decoy_tag = self.params.get("decoy_tag", "decoy_")
        self.df.loc[:, "is_decoy"] = self.df["protein_id"].str.contains(decoy_tag)
        if self.immutable_peptides is not None:
            auto = ahocorasick.Automaton()
            for seq in self.immutable_peptides:
                auto.add_word(seq, seq)
            auto.make_automaton()
            self.df.loc[:, "is_immutable"] = [
                sum([len(match) for _, match in auto.iter_long(seq)]) == len(seq)
                for seq in self.df["sequence"]
            ]
        else:
            self.df.loc[:, "is_immutable"] = False

    def process_unify_style(self):
        """Combine all additional operations that are needed to calculate new columns and sanitize the dataframe.

        Operations are performed inplace on self.df
        """
        self.df.drop_duplicates(inplace=True, ignore_index=True)
        self.clean_up_modifications()
        self.assert_only_iupac_and_missing_aas()
        self.add_protein_ids()
        self.get_meta_info()
        self.calc_masses_offsets_and_composition()
        self.check_enzyme_specificity()
        self.add_ranks()
        self.add_decoy_identity()
        self.sanitize()
