import argparse
import pandas as pd
import re
import os
from collections import defaultdict
import numpy as np
import sys
sys.path.append(os.getcwd())
from utils.lca_determination.determine_summary_lca import create_lca_from_summary


class ReadSummary:
    def __init__(self, fname, sep):
        """
        ## large RAM requirement, complete summary file there
            :param fname: path to summary file
            :param sep: seperator data frame
            :param tasks: list of lists all tasks with same id in one sublist
                    [['task_0::tax_from_nr_20180808_qcover90_diamond::superkingdom',... ],[],...]
        """
        self.fname = fname
        self.df = pd.read_csv(self.fname, sep, dtype=str, header=0).replace(np.nan, '-', regex=True)
        self.regex = re.compile("task_[0-9]+::")
        self.task_indices = self._find_column_start_indices_per_task()
        self.tasks = self._get_all_task_ids()
        self.members_identifier_pos = list(self.df).index('members_identifier')
        self.pgid_pos = list(self.df).index('#pg')
        self.member_pos = list(self.df).index('level')
        try:
            spectra_columns = [i for i, e in enumerate(list(self.df)) if re.match(".*:: spectra_IDs", e)]
            self.spectralIDs_start_pos = spectra_columns[0]
            self.spectralIDs_end_pos = spectra_columns[-1]
        except IndexError:
            self.spectralIDs_start_pos = None
        self.is_with_spectra = True if self.spectralIDs_start_pos is not None else False

    @staticmethod
    def flatten_list(l):
        """
            :return one flat list from list of lists
        """
        return [item for sublist in l for item in sublist]

    def get_member(self, pgid):
        """
            :param pgid: protein group id
            :return filtered dataframe with col #pg == pgid and col level == member
        """
        return self.df[(self.df['#pg'] == pgid) & (self.df['level'] == "member")]

    def _find_column_start_indices_per_task(self):
        j = None
        start_indices = []
        task_indices = []
        k = 0
        for i, item in enumerate(self.df.columns):
            if self.regex.match(item) and item.split("::")[0] != j:
                k = 0
                j = item.split("::")[0]
                start_indices.append(i)
            if self.regex.match(item) and item.split("::")[0] == j:
                k += 1
        start_indices.append(start_indices[-1] + k)
        for i in range(len(start_indices) - 1):
            task_indices.append((start_indices[i], start_indices[i + 1]))
        return task_indices

    def _get_all_task_ids(self):
        """
            :return tasks: list of tax tasks []
        """
        tasks = []
        i = None
        for task in [x for x in self.df.columns if self.regex.match(x)]:
            task_id = task.split("::")[0]
            if i != task_id:
                tasks.append([])
                i = task_id
            tasks[-1].append(task)
        return tasks

    def get_spectra_dict(self, id_sep):
        """
        :param id_sep: string separator between protein/spectra accessions
        :return: acc_to_spectra_dict
        """
        acc_to_spectra_dict = defaultdict(list)
        for row in self.df.values.tolist():
            if row[self.member_pos] == 'member':
                if row[self.members_identifier_pos] not in acc_to_spectra_dict.keys():
                    if self.spectralIDs_start_pos is not None:
                        spectra_ID_lists = [spectra_per_sample.split(id_sep) for spectra_per_sample in
                                            row[self.spectralIDs_start_pos:self.spectralIDs_end_pos + 1]]
                        spectra_IDs = set(
                            [spectraID for sublist in spectra_ID_lists for spectraID in sublist if spectraID != '-'])
                        acc_to_spectra_dict[row[self.members_identifier_pos]] = spectra_IDs
                    # if no spectra IDs available use protein accession instead. Equals protein-support value of 1
                    else:
                        spectra_IDs = {row[self.members_identifier_pos]}
                        acc_to_spectra_dict[row[self.members_identifier_pos]] = spectra_IDs
        return acc_to_spectra_dict

    def add_lineages_per_task_to_dict(self, row, task_to_acc_to_lineages_dict):
        for taskid in range(len(self.tasks)):
            lineage = row[self.task_indices[taskid][0]:self.task_indices[taskid][1]]
            if row[self.members_identifier_pos] in task_to_acc_to_lineages_dict[taskid].keys():
                task_to_acc_to_lineages_dict[taskid][row[self.members_identifier_pos]].append(lineage)
            else:
                task_to_acc_to_lineages_dict[taskid][row[self.members_identifier_pos]] = [lineage]
        return task_to_acc_to_lineages_dict

    def add_accs_to_dict(self, row, id_sep, pgid_to_accs_dict):
        pgid_to_accs_dict[int(row[self.pgid_pos])] = row[self.members_identifier_pos].split(id_sep)
        return pgid_to_accs_dict

    def read_lineages_from_summary_df_in_dict(self, id_sep):
        """
        :param id_sep: string separator between protein/spectra accessions
        :return: task_to_acc_to_lineages_dict, pgid_to_accs_dict
        """
        task_to_acc_to_lineages_dict = defaultdict(dict)
        pgid_to_accs_dict = {}
        for row in self.df.values.tolist():
            if row[self.member_pos] == 'member':
                task_to_acc_to_lineages_dict = self.add_lineages_per_task_to_dict(row, task_to_acc_to_lineages_dict)
            if row[self.member_pos] == 'group':
                pgid_to_accs_dict = self.add_accs_to_dict(row, id_sep, pgid_to_accs_dict)
        return task_to_acc_to_lineages_dict, pgid_to_accs_dict


def main():
    parser = argparse.ArgumentParser(prog='wheight_lcas.py',
                                     description='introduce weighted LCAs in Prophane\'s summary reports')
    parser.add_argument('-r', '--report', metavar="SUMMARY_FILE", help="summary report", type=str)
    parser.add_argument('-o', '--outfile', metavar="OUT_FILE", help="output file (will be overwritten!)", type=str)
    parser.add_argument('-m', '--method', metavar="STRING", help="method used for lcas [lca = weighted per group, "
                                                                 "democratic_lca = democratic (default: lca)", type=str,
                        default='lca')
    parser.add_argument('-t', '--threshold', metavar="FLOAT", help="minimal relative support [0, 1] (default: 0.5001),"
                                                                   " threshold for lca", type=float, default=0.5001)
    parser.add_argument('-x', '--taxmap', metavar="TAX_MAP", help="path to tax_map", type=str)
    parser.add_argument('-i', '--id_sep', metavar="STRING", help="separator of protein/spectra accessions", type=str,
                        default=';')
    parser.add_argument('-s', '--sep', metavar="STRING", help="separator of columns", type=str, default='\t')
    parser.add_argument('--ignore_unclassified', action='store_true', help="remove all complete unclassified lineages from analysis.",
                        default=False)
    parser.add_argument('--min_nb_annot', metavar="INTEGER", help="minimal number of classified entries per lineage to be "
                                                                  "considered in analysis (default: 1).", type=int, default=1)
    args = parser.parse_args()

    if (args.threshold > 1) | (args.threshold < 0):
        raise ValueError("LCA support must be between 0 (0%) and 1 (100%).")
    if args.method not in ['lca', 'democratic_lca']:
        raise ValueError("Unknown lca strategy. lca= group based LCA with threshold, 1 = democratic LCA.")
    if not os.path.isfile(args.report):
        raise IOError(f"summary.txt path ({args.report}) do not exist")

    reader = ReadSummary(args.report, args.sep)
    accs_to_spectra_dict = reader.get_spectra_dict(args.id_sep)
    task_to_acc_to_lineages_dict, pgid_to_accs_dict = reader.read_lineages_from_summary_df_in_dict(args.id_sep)
    create_lca_from_summary(args.outfile, reader, args.threshold, args.method, pgid_to_accs_dict, accs_to_spectra_dict,
                            task_to_acc_to_lineages_dict, args.taxmap, args.sep, args.ignore_unclassified, args.min_nb_annot)


if __name__ == "__main__":
    main()
