import csv
from collections import defaultdict
from pathlib import Path
import pandas as pd

from utils.config import INCONCLUSIVE_ANNOTATIONS
from utils.input_output import iter_file, get_headline
from utils.protein_group_sql_db import ProteinGroupSqlDbInteractor


class ProteinGroupSummary:
    def __init__(self, pgid, pg_db, accs_per_group, taskids, taskid_to_lca_dict, lca_support_results, mafft_results,
                 quant_data, samples, task_dfs_filtered_for_accs, sample_group_to_sample_list_dict,
                 sample_id_to_sorted_spectra_list=None):
        """
        :param samples: sorted list of samples from database
        :param accs_per_group: list of accs of group pgid
        """
        self.pgid = pgid
        self.pg_db = pg_db
        self.accs_per_group = sorted(accs_per_group)
        self.taskids = taskids
        self.taskid_to_lca_dict = taskid_to_lca_dict
        self.lca_support_results = lca_support_results
        self.mafft_results = mafft_results
        self.quant_data = quant_data
        self.sample_id_to_sorted_spectra_list = sample_id_to_sorted_spectra_list
        self.samples = samples

        self.task_dfs_of_group_accs = task_dfs_filtered_for_accs
        self.with_spectra = sample_id_to_sorted_spectra_list is not None
        self.sample_group_to_sample_list_dict = sample_group_to_sample_list_dict
        self.len_mafft = 5
        self.len_quant = 4 * len(self.sample_group_to_sample_list_dict) \
                         + 4 * sum([len(sample_list) for sample_list in self.sample_group_to_sample_list_dict.values()])

    def get_group_row_elements(self):
        group_row = ['group', str(self.pgid), str(len(self.accs_per_group))]
        group_row.extend(self.mafft_results)
        group_row.extend([';'.join(self.accs_per_group)])
        for taskid in range(self.taskids):
            group_row.extend(self.taskid_to_lca_dict[taskid])
        group_row.extend(self.quant_data)
        if self.sample_id_to_sorted_spectra_list:
            for sample in self.samples:
                group_row.extend([self.sample_id_to_sorted_spectra_list[sample]])
        for taskid in range(self.taskids):
            group_row.extend(self.lca_support_results[taskid])
        return group_row

    def _combine_all_member_entries_into_one_list(self, task_row, spectra=None):
        member_row = ['member', str(self.pgid)] + ['-'] * self.len_mafft + task_row + ['-'] * self.len_quant
        if spectra:
            member_row.extend(spectra)
        return member_row

    def _get_df_per_accs_per_task(self, df):
        acc_to_df_dict = {}
        for acc, acc_df in df.groupby("#acc"):
            acc_df = acc_df.sort_values(by=list(df.columns)).reset_index(drop=True)
            acc_df = acc_df.drop('#acc', axis=1)
            acc_to_df_dict[acc] = acc_df
        return acc_to_df_dict

    def _create_task_df(self, list_of_dict_acc_to_df_per_task):
        """
        # merge by accs not possible (all rows were filled)
        # methode here: for every sorted acc of the group accs all task results are concatenated side by side (=df_acc)
        # all df_acc are concatenated one below the other
        :return final_task_df, df with all task results of one protein group
        """
        final_task_df=pd.DataFrame()
        for acc in self.accs_per_group:
            df_acc = pd.DataFrame()
            # concat side by side all task results for one accs
            for task_dict in list_of_dict_acc_to_df_per_task:
                df_acc = pd.concat([df_acc, task_dict[acc]], axis=1)
            df_acc.insert(loc=0, column='#acc', value='')
            df_acc['#acc'] = acc
            # concat below
            final_task_df = pd.concat([final_task_df, df_acc], ignore_index=True)
        final_task_df = final_task_df.fillna('-')
        return final_task_df

    def member_data_generator(self):
        """
        :return member_rows: list of lists, sublist=member_row, one element=one cell in summary
        """
        list_of_dict_acc_to_df_per_task= [self._get_df_per_accs_per_task(df) for df in self.task_dfs_of_group_accs]
        final_task_df = self._create_task_df(list_of_dict_acc_to_df_per_task)
        task_rows = list(final_task_df.values.tolist())
        if self.with_spectra:
            acc_to_spectra_dict = ProteinGroupSqlDbInteractor(self.pg_db).get_spectra_per_acc(self.accs_per_group, self.pgid)
        for task_row in task_rows:
            if self.with_spectra:
                acc = task_row[0]
                member_row = self._combine_all_member_entries_into_one_list(task_row, acc_to_spectra_dict[acc])
            else:
                member_row = self._combine_all_member_entries_into_one_list(task_row)
            yield member_row


class Summary:
    def __init__(self, pg_db, task_lca_files, task_lca_support_files, task_map_files, quant_file, mafft_files,
                 tasks, sample_group_to_sample_list_dict):
        """
        :param pg_db: path to protein group db
        :param task_lca_files: list of lca result files for every task, ordered by taskid  'tasks/*.task_nb.lca'
        :param task_lca_support_files: list of lca support result files for every task,
                                ordered by taskid 'tasks/*.task_nb.lca_support'
        :param task_map_files: list of task map result files for every task, ordered by taskid 'tasks/*.task_nb.map'
        :param quant_file: path to quant tsv '/tasks/quant.tsv'
        :param mafft_files: path to mafft report files '/algn/report/mafft.pgid.txt
        :param tasks: list of task dicts with keys: 'prog', 'type', 'shortname', 'db_type', 'params'
        """
        self.pg_db = pg_db
        self.lca_task_files = task_lca_files
        self.lca_support_files = task_lca_support_files
        self.task_map_files = task_map_files
        self.quant_file = quant_file
        self.mafft_report_files = mafft_files
        self.task_dicts = tasks
        self.sample_group_to_sample_list_dict = sample_group_to_sample_list_dict

        self.taskids = len(tasks)
        self.sql_db_interactor = ProteinGroupSqlDbInteractor(pg_db)
        self.samples = self.sql_db_interactor.get_samples()
        self.with_spectra = self.sql_db_interactor.is_with_spectra()
        self.taskid_to_levelnb_dict = self._get_length_of_annotations_per_task()
        self.accs = [acc for accs in self.sql_db_interactor.get_acc_list_per_pg().values() for acc in accs]

        self.pgid_to_mafft_results_dict = self._read_mafft_report_files()
        self.pgid_to_taskid_to_lca_lineage_dict = self._read_task_lca_files_into_dict()
        self.pgid_to_taskid_to_lca_support_dict = self._read_task_lca_support_files_into_dict()
        self.quant_data_dict = self._read_quant_data_file_into_dict()
        self.task_df_list = self._read_task_map_files()

    def _clean_map_df(self, df):
        """
         remove nan and -, replace by 'unclassified'
        """
        df = df.replace('-', INCONCLUSIVE_ANNOTATIONS['unclassified'])
        df = df.replace('nan', INCONCLUSIVE_ANNOTATIONS['unclassified'])
        df = df.fillna(INCONCLUSIVE_ANNOTATIONS['unclassified'])
        return df

    def _get_length_of_annotations_per_task(self):
        """
        length of task level = length of lca header - 1 (#pg)
        :return task_to_levelnb_dict: {taskid: level length}
        """
        task_to_levelnb_dict = {}
        for taskid in range(self.taskids):
            task_to_levelnb_dict[taskid] = len(get_headline(self.lca_task_files[taskid]).split('\t')) - 1
        return task_to_levelnb_dict

    def _read_quant_data_from_quant_tsv(self):
        pgid_to_sample_to_quant_values_dict = defaultdict(dict)
        for line in iter_file(self.quant_file):
            fields = line.split("\t")
            pgid_to_sample_to_quant_values_dict[int(fields[0])][fields[1]] = fields[-4:]
        return pgid_to_sample_to_quant_values_dict

    def _sort_quant_data(self, pgid_to_sample_to_quant_values_dict):
        """
        sort quant data into correct order for printing in group row
        : return quant_data_dict: dict{pgid:[quant values]}
        """
        quant_data_dict_sorted = defaultdict(list)
        for pgid, sample_to_quant_values_dict in pgid_to_sample_to_quant_values_dict.items():
            if self.sample_group_to_sample_list_dict:
                for group in list(self.sample_group_to_sample_list_dict.keys()):
                    quant_data_dict_sorted[pgid].extend(sample_to_quant_values_dict[group + " (mean)"])
                    for sample in self.sample_group_to_sample_list_dict[group]:
                        quant_data_dict_sorted[pgid].extend(sample_to_quant_values_dict[sample])
            else:
                for sample in self.samples:
                    quant_data_dict_sorted[pgid].extend(sample_to_quant_values_dict[sample])
        return quant_data_dict_sorted

    def _read_quant_data_file_into_dict(self):
        pgid_to_sample_to_quant_values_dict = self._read_quant_data_from_quant_tsv()
        quant_data_dict = self._sort_quant_data(pgid_to_sample_to_quant_values_dict)
        return quant_data_dict

    def _test_for_unknown_accs(self, task_df, taskid):
        acc_set_of_df = set(task_df['#acc'])
        if len(acc_set_of_df - set(self.accs)) > 0:
            raise KeyError(f"Task results of task {taskid} contain accessions {', '.join(acc_set_of_df - set(self.accs))}"
                           f" not contained in the protein group yaml.")

    def _add_unclassified_for_missing_accs(self, task_df):
        """
        :param task_df: columns: #acc | level1 | level2 | ...
        """
        missing_accs = set(self.accs) - set(task_df['#acc'])
        task_len = len(task_df.columns)-1
        missing_lineages = [[missing_acc] + [INCONCLUSIVE_ANNOTATIONS['unclassified']]*task_len for missing_acc in missing_accs]
        new_part_df = pd.DataFrame(missing_lineages, columns=task_df.columns)
        task_df = task_df.append(new_part_df)
        return task_df

    def _read_task_map_files(self):
        df_list = []
        for taskid in range(self.taskids):
            task_df = pd.read_csv(self.task_map_files[taskid], sep='\t')
            # remove not needed columns, keep first (acc) column and last x columns
            task_df.drop(task_df.iloc[:, 1:len(task_df.columns)-self.taskid_to_levelnb_dict[taskid]], axis=1, inplace=True)
            # column names
            task_df.columns = [list(task_df.columns)[0]] + [f"task_{taskid}_{column_name}" for column_name in list(task_df.columns)[1:]]
            self._test_for_unknown_accs(task_df, taskid)
            task_df = self._clean_map_df(task_df)
            task_df = self._add_unclassified_for_missing_accs(task_df)
            task_df = task_df.drop_duplicates()
            df_list.append(task_df)
        return df_list

    def _read_files(self, file_list):
        pgid_to_taskid_to_value_list_dict = defaultdict(dict)
        for taskid in range(self.taskids):
            file = file_list[taskid]
            for line in iter_file(file):
                pg_id, *lca_results = line.split("\t")
                pgid_to_taskid_to_value_list_dict[int(pg_id)][taskid] = lca_results
        return pgid_to_taskid_to_value_list_dict

    def _read_task_lca_files_into_dict(self):
        taskid_to_pgid_to_lca_lineage_dict = self._read_files(self.lca_task_files)
        return taskid_to_pgid_to_lca_lineage_dict

    def _read_task_lca_support_files_into_dict(self):
        taskid_to_pgid_to_lca_support_dict = self._read_files(self.lca_support_files)
        return taskid_to_pgid_to_lca_support_dict

    def _read_mafft_report_files(self):
        pgid_to_mafft_results_dict = {}
        for pgid in range(self.sql_db_interactor.get_number_of_pgs()):
            mafft_file_of_pg = [file for file in self.mafft_report_files if pgid == int(Path(file).stem.split('.')[1])]
            if len(mafft_file_of_pg) == 1:
                with open(mafft_file_of_pg[0], "r") as handle:
                    min_seqlen, max_seqlen, min_ident, max_ident = handle.readlines()[1].strip().split("\t")
            elif len(mafft_file_of_pg) == 0:
                min_seqlen = max_seqlen = min_ident = max_ident = "-"
            pgid_to_mafft_results_dict[pgid] = [min_seqlen, max_seqlen, min_ident, max_ident]
        return pgid_to_mafft_results_dict

    def _get_headline_task_data(self):
        task_header = []
        for taskid in range(self.taskids):
            headline_task = [f"task_{taskid}::{self.task_dicts[taskid]['shortname'].replace(' ', '_')}::" +
                             x.replace('_', ' ') for x in get_headline(self.lca_task_files[taskid]).split('\t')]
            task_header.extend(headline_task[1:])
        return task_header

    def _get_headline_lca_support(self, tasks_header):
        headline_lca_support =  [f"{entry}::lca_support" for entry in tasks_header]
        return headline_lca_support

    def _get_headline_quant_data(self):
        quant_headline = []
        headline_quant = get_headline(self.quant_file).split("\t")[-4:]
        quant_cols = []
        if self.sample_group_to_sample_list_dict:
            for group in list(self.sample_group_to_sample_list_dict.keys()):
                quant_cols.append(group + " (mean)")
                for sample in self.sample_group_to_sample_list_dict[group]:
                    quant_cols.append(sample)
        else:
            for sample in self.samples:
                quant_cols.append(sample)
        for sample in quant_cols:
            quant_headline.extend([x + "::" + sample for x in headline_quant])
        return quant_headline

    def _get_headline_spectra(self):
        if self.with_spectra:
            return self.sql_db_interactor.get_spectra_header()
        else:
            return []

    def _get_spectra_per_pg_per_sample_group(self, pgid):
        sample_to_spectra_set = self.sql_db_interactor.get_set_all_spectra_per_sample_for_pg(pgid)
        return {sample: ';'.join(sorted(list(spectra_set))) if len(spectra_set) > 0 else '-' for sample, spectra_set
                in sample_to_spectra_set.items()}

    def _get_summary_header(self):
        headline_list = ["level", "#pg", "members_count", "min_seqlen", "max_seqlen", "min_rel_pairw_ident",
                         "max_rel_pairw_ident", "members_identifier"]
        tasks_header = self._get_headline_task_data()
        headline_list.extend(tasks_header)
        headline_list.extend(self._get_headline_quant_data())
        headline_list.extend(self._get_headline_spectra())
        headline_list.extend(self._get_headline_lca_support(tasks_header))
        return headline_list

    def _protein_group_rows_generator(self, pgid):
        accs_per_group = self.sql_db_interactor.get_acc_list_per_pg([pgid])[pgid]
        sample_id_to_sorted_spectra_list = self._get_spectra_per_pg_per_sample_group(pgid) if self.with_spectra else None
        task_dfs_filtered_for_acc = [task_df.loc[task_df['#acc'].isin(accs_per_group)] for task_df in self.task_df_list]
        pg_row_obj = ProteinGroupSummary(pgid, self.pg_db, accs_per_group, self.taskids,
                                         self.pgid_to_taskid_to_lca_lineage_dict[pgid],
                                         self.pgid_to_taskid_to_lca_support_dict[pgid],
                                         self.pgid_to_mafft_results_dict[pgid],
                                         self.quant_data_dict[pgid], self.samples, task_dfs_filtered_for_acc,
                                         self.sample_group_to_sample_list_dict, sample_id_to_sorted_spectra_list)
        yield pg_row_obj.get_group_row_elements()
        for row in pg_row_obj.member_data_generator():
            yield row

    def _write_summary_header(self, output, headline_list):
        """
        :param headline_list: list of header elements
        """
        with open(output, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile, delimiter='\t')
            writer.writerow(headline_list)

    def _write_summary_protein_group_row(self, output, row_elements_list):
        """
        :param row_elements_list: list containing all elements of one row
        the elements are concatenated with tab-characters and written as additional line to output
        """
        with open(output, 'a', newline='') as csvfile:
            writer = csv.writer(csvfile, delimiter='\t')
            writer.writerow(row_elements_list)

    def write_summary_to_output(self, output):
        headline_list = self._get_summary_header()
        self._write_summary_header(output, headline_list)
        for pgid in range(self.sql_db_interactor.get_number_of_pgs()):
            for row in self._protein_group_rows_generator(pgid):
                self._write_summary_protein_group_row(output, row)
