import os
import pandas as pd
import re
import sys
import yaml

from utils.exceptions import SearchResultParsingError, StyleError
from utils.search_result_parsing.style_handling import Style
from utils.search_result_parsing.parse_generic_table import GenericTableReader
from utils.search_result_parsing.parse_mpa_metaprotein_report import MpaPortableTableReader
from utils.search_result_parsing.parse_mpa_multisample_metaprotein_report import MpaPortableMultisampleTableReader
from utils.search_result_parsing.parse_scaffold_protein_report import ScaffoldTableReader
from utils.search_result_parsing.parse_mzident_report import MzidentTableReader
from utils.search_result_parsing.parse_mztab import MzTabReader
from utils.search_result_parsing.parse_proteome_discoverer import PDParser


class ProteomicSearchResultParser(object):
    """
    Class for parsing the result files of proteomic search engines.
    After initialising the class, a protein group yaml can be written using 'write_protein_groups_yaml'
    params result_table: results search engines, eg. mzid
    params f_style: path to yaml file eg. path/to/mpa_1_8_x.yaml
    param out_file: write yaml file to out_file eg. tests/resources/test_full_analysis_and_file_presence/test_mzIdent_parser/expected-results/protein_groups.yaml
    param file_reader_class: return reader class, eg. MpaPortableTableReader
    param bool_read_sample_descriptors_from_table:
    param: required_sample_descriptors
    param: file_reader: right file_reader_class execute parsing with params result_table, f_style and required_sample_descriptors
    param: style_dict: dict for elements contained in result yaml
    parm: df_standard: resulting dataframe from parsing search results
    """
    mandatory_style_fields = ['proteins_sep', 'field_sep', 'sample_cols', 'proteins_col', 'quant_col']
    supported_styles_dict = {
        "generic.yaml": GenericTableReader,
        "mpa_1_8_x.yaml": MpaPortableTableReader,
        "mpa_server_multisample.yaml": MpaPortableMultisampleTableReader,
        "scaffold_4_8_x.yaml": ScaffoldTableReader,
        "mzident_1_2_0.yaml": MzidentTableReader,
        "mztab_1_0_0.yaml": MzTabReader,
        "proteome_discoverer_2_5.yaml": PDParser
    }

    def __init__(self, result_table, style, out_file, decoy_regex, required_sample_descriptors):
        self.result_table = result_table
        self.f_style = style
        self.out_file = out_file
        self.decoy_regex = decoy_regex
        self.file_reader_class = self.get_file_parser_for_style()
        self.bool_read_sample_descriptors_from_table = True if not required_sample_descriptors else False
        self.required_sample_descriptors = set(required_sample_descriptors)
        self.file_reader = self.file_reader_class(self.result_table, self.f_style, self.required_sample_descriptors)
        self.style_dict = Style(self.f_style)
        self.df_standard = self.parse_search_result()

    def get_file_parser_for_style(self):
        style_basename = os.path.basename(self.f_style)
        if style_basename in self.supported_styles_dict:
            parser_class = self.supported_styles_dict[style_basename]
        else:
            raise StyleError("Provided style for search result parsing is not supported. Provided style:\n"
                             f"{style_basename}\n"
                             "Supported styles are:\n"
                             f"\t{self.supported_styles_dict.keys()}")
        return parser_class

    def write_protein_groups_yaml(self):
        print("writing protein report ...", file=sys.stderr)
        self.write_to_yaml()
        self.print_summary()

    def parse_search_result(self):
        print("reading protein report ...", file=sys.stderr)
        df = self.file_reader.read_file_into_dataframe()
        self.sanity_check_standard_df(df)
        if self.decoy_regex:
            df = self.filter_out_decoys(df)
        self.sanity_check_protein_acc_in_df(df)
        if self.bool_read_sample_descriptors_from_table:
            self.required_sample_descriptors = self.read_sample_descriptors_from_table(df)
        self.check_consensus_of_cfg_and_table_derived_sample_descriptors(df)
        df = self.sort_protein_cell_content(df)
        df = self.add_pg_string(df)
        self.check_value_uniqueness(df)
        return df

    def check_value_uniqueness(self, df=None):
        if df is None:
            df = self.df_standard
        groups = df.groupby(['sample', 'pg_string'])
        samples_with_multiple_quant = groups.size()[~(groups.size() == 1)].to_frame('count')
        if not (groups.size() == 1).all():
            raise SearchResultParsingError(
                'Expected a single quantification value for each protein-group/sample combination.' +
                'The following combinations have more than one:\n' +
                str(samples_with_multiple_quant)
            )

    def add_pg_string(self, df):
        df['pg_string'] = df['proteins'].apply(lambda x: self.style_dict['proteins_sep'].join(x))
        return df

    def sanity_check_standard_df(self, df=None):
        if df is None:
            df = self.df_standard
        required_cols = ['proteins', 'sample', 'quant']
        try:
            df[required_cols]
        except KeyError:
            raise AttributeError('Required fields are not present in self.standard_dataframe\n' +
                                 'Required: {}\n'.format(required_cols) +
                                 'Present: {}'.format(df.columns.values))

        # at least one row
        if df.shape[0] == 0:
            sys.exit("input format error in " + self.result_table + ":\nno data rows")

    def sanity_check_protein_acc_in_df(self, df=None):
        # not allowed to occur at all
        start_of_heading = chr(1)
        illegal_chars = {start_of_heading, " "}
        # not allowed to occur more than once
        illegal_more_than_once = {'|'}

        def does_acc_contain_illegal_chars(acc):
            for c in illegal_chars:
                if c in acc:
                    return True
            for c in illegal_more_than_once:
                if acc.count(c) > 1:
                    return True
            return False

        def get_set_of_illegal_accessions(series_of_accession_lists):
            set_accs = {acc for accs in series_of_accession_lists for acc in accs}
            illegal_accs = {acc for acc in set_accs if does_acc_contain_illegal_chars(acc)}
            return illegal_accs

        if df is None:
            df = self.df_standard
        illegal_accs = get_set_of_illegal_accessions(df["proteins"])
        if illegal_accs:
            raise AttributeError(
"""Input error: at least one protein accessions contains illegal characters (separator of header, space or multiple 
pipes)
illegal accessions:
{}""".format('\n'.join(illegal_accs))
            )

    def check_consensus_of_cfg_and_table_derived_sample_descriptors(self, df=None):
        if df is None:
            df = self.df_standard
        # contains all required sample descriptors
        present_sample_descriptors = set(df['sample'].unique())
        if not self.required_sample_descriptors.issubset(present_sample_descriptors):
            raise AttributeError(
                "Missing required sample descriptors:\n"
                + "\t" + "\n\t".join(self.required_sample_descriptors - present_sample_descriptors) + "\n"
                + "Present sample descriptors are:\n"
                + "\t" + "\n\t".join(present_sample_descriptors)
            )
        # if table contains sample descriptors not provided by user
        elif present_sample_descriptors - self.required_sample_descriptors:
            raise AttributeError(
                "Some input table derived sample descriptors are missing in config section sample_groups.\n"
                f"\tInput table: {self.result_table}\n"
                "Missing sample descriptors are:\n"
                + "\t" + "\n\t".join(present_sample_descriptors - self.required_sample_descriptors) + "\n"
            )

    def read_sample_descriptors_from_table(self, df):
        set_sample_descriptors = set(df["sample"].unique())
        return set_sample_descriptors

    def get_sample_groups_set(self):
        return self.required_sample_descriptors

    # template for parsing. Result: dataframe with columns -proteins (=list of accessions)
    # -quant(=number identified spectra for all proteins in the group) and -sample (=sample name, eg. name spectra_file)
    def read_file_into_dataframe(self):
        raise NotImplemented('This is a template class. '
                             'Please define your own child class with method "read_file_into_dataframe".')
        # some Code for reading a table, please adapt to your needs
        required_col_names = self.style_dict['sample_cols'] \
                             + [self.style_dict['proteins_col'], self.style_dict['quant_col']]
        raw_df = pd.read_csv(self.result_table, sep=self.style_dict['field_sep'])
        df = raw_df[required_col_names].copy()
        df['sample'] = df[self.style_dict['sample_cols']] \
            .apply(lambda x: '::'.join(x), axis=1)

        df['proteins'] = df[self.style_dict['proteins_col']].apply(lambda x: x.split(self.style_dict['proteins_sep']))

        df_standard = pd.DataFrame()
        df_standard['proteins'] = df['proteins']
        df_standard['sample'] = df['sample']
        df_standard['quant'] = df[self.style_dict['quant_col']]

        return df_standard

    def write_to_yaml(self):
        samples = sorted(self.df_standard['sample'].unique())
        include_single_acc_quant = ('spectra_IDs' in self.df_standard)
        yaml_data = {}
        # samples
        yaml_data['samples'] = samples
        # style definitions
        yaml_data['style'] = {}
        yaml_data['style']['protein_sep'] = self.style_dict['proteins_sep']
        yaml_data['style']['quant_type'] = self.style_dict['quant_type']
        yaml_data['protein_groups'] = []
        groups = self.df_standard.groupby('pg_string')
        lst_pg_strings = self.df_standard['pg_string'].unique()
        for proteins in sorted(lst_pg_strings):
            df_sample = groups.get_group(proteins)
            protein_group_dict = {'proteins': proteins, 'quant': []}
            if include_single_acc_quant:
                protein_group_dict['spectra'] = []
            for sample in samples:
                # if sample in protein-group specific dataframe
                if df_sample['sample'].str.match(re.escape(sample)+'$').any():
                    quant = df_sample[df_sample['sample'] == sample]['quant'].values
                    protein_group_dict['quant'].append(float(quant[0]))
                    if include_single_acc_quant:
                        row_nb = int(str(df_sample[df_sample['sample'] == sample]['spectra_IDs']).split()[0])
                        ID_str = str(self.df_standard['spectra_IDs'][row_nb])
                        protein_group_dict['spectra'].append(ID_str.replace("'", ""))
                else:
                    protein_group_dict['quant'].append(0.)
                    if include_single_acc_quant:
                        protein_group_dict['spectra'].append('-')
            yaml_data['protein_groups'].append(protein_group_dict)
        # writing
        with open(self.out_file, "w") as handle:
            handle.write(yaml.dump(yaml_data, default_flow_style=False))
        return

    def print_summary(self):
        n_groups = len(self.df_standard['proteins'].apply(lambda x: ";".join(x)).unique())
        n_proteins = len({p for lst in self.df_standard['proteins'] for p in lst})
        print("protein groups:", n_groups, file=sys.stderr)
        print("unique proteins:", n_proteins, file=sys.stderr)

    def filter_out_decoys(self, df):
        """
        remove accessions from list of protein accessions that match to self.decoy_regex.
        If resulting list is empty, remove entire row.

        :param df: dataframe that will be filtered
        :return: dataframe with decoy accessions removed
        """

        def filter_decoys_from_protein_list(protein_list: list, regex_str):
            re_decoy = re.compile(regex_str)
            filtered_lst = [acc for acc in protein_list if not re_decoy.search(acc)]
            return filtered_lst

        df['proteins'] = df['proteins'].apply(lambda x: filter_decoys_from_protein_list(x, self.decoy_regex))
        # drop rows containing only decoys
        df = df[df['proteins'].apply(len) != 0]
        return df

    def sort_protein_cell_content(self, df=None):
        if df is None:
            df = self.df_standard
        df['proteins'] = df['proteins'].apply(sorted)
        return df
