import os
import pandas as pd
import re
import sys
import math

from utils.exceptions import SearchResultParsingError, StyleError
from utils.protein_group_sql_db import create_protein_group_database
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 create_protein_groups_db(self):
        print("writing protein report ...", file=sys.stderr)
        self.create_pg_sql_db()
        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)
        self.check_for_quant_values(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 check_for_quant_values(self, df=None):
        if df is None:
            df = self.df_standard
        quant_df = df[['pg_string', 'quant']].groupby(['pg_string'], as_index=False).agg({'quant': lambda x: list(x.unique())})
        for pg_quant_list, pg_string in zip(quant_df['quant'].tolist(), quant_df['pg_string'].tolist()):
            if len(pg_quant_list) == 1 and math.isclose(pg_quant_list[0], 0, rel_tol=1e-5):
                raise SearchResultParsingError(
                    "Expected at least one quantification value of one sample for every protein group"+
                    " The following protein group is without quantification value:\n" + pg_string
                )

    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 get_protein_group_quant_data_dict(self, groups, proteins, with_spectra, samples):
        """
        convert information of pg_dataframe (protein, quant, spectra) into dict and fill missing quant info with 0 for
        'quant' or '-' for spectra
        :param groups: pandas group by object, pg-dataframe grouped by 'pg_string'
        :param proteins: protein strings concat with separator (one string)
        :param with_spectra: True if spectra information contained in df, else False
        :param samples: list of samples
        :return protein_group_dict{'proteins':string of concat protein IDs, 'quant': list of floats,
        'spectra': string of concat spectra IDs}
        """
        df_sample = groups.get_group(proteins)
        protein_group_dict = {'proteins': proteins, 'quant': []}
        if with_spectra:
            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 with_spectra:
                    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)
            else:
                protein_group_dict['quant'].append(0.)
                if with_spectra:
                    protein_group_dict['spectra'].append('-')
        return protein_group_dict

    def create_pg_sql_db(self):
        samples = sorted(self.df_standard['sample'].unique())
        with_spectra = ('spectra_IDs' in self.df_standard)
        include_single_acc_quant = with_spectra
        # style definitions
        protein_info_data = {'samples': samples, 'protein_sep': self.style_dict['proteins_sep'], 'with_spectra': with_spectra,
                             'quant_type': self.style_dict['quant_type']}
        protein_group_data = []
        groups = self.df_standard.groupby('pg_string')
        lst_pg_strings = self.df_standard['pg_string'].unique()
        for proteins in sorted(lst_pg_strings):
            protein_group_dict = self.get_protein_group_quant_data_dict(groups, proteins, include_single_acc_quant, samples)
            protein_group_data.append(protein_group_dict)
        protein_group_db_obj = create_protein_group_database(self.out_file, protein_group_data, protein_info_data)

    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
