import csv
from collections import defaultdict
from pathlib import Path

from utils.config import INCONCLUSIVE_ANNOTATIONS
from utils.helper_funcs_for_snakemake_workflow import get_mafft_report
from utils.input_output import (
    iter_file,
    get_headline,
    get_spectra_header_and_accs_to_spectraID_dict,
    get_accs_per_pg_from_yaml, get_samples, get_pg_dict, is_with_spectra)


def clean_member_data(lineage_tuple):
    """
    # remove nan and -, replace by 'unclassified'
    """
    return tuple([INCONCLUSIVE_ANNOTATIONS['unclassified'] if elem in ['-', 'nan'] else elem for elem in lineage_tuple])


def get_quantdata(quant_file, samples, sample_groups):
    """
    :param quant_file: quant results
    :param samples:
    :param sample_groups:
    :return quant_headline, quant_data
    """
    quant_headline = []
    headline_quant = get_headline(quant_file).split("\t")[-4:]
    quant_cols = []
    if sample_groups:
        for group in sorted(sample_groups.keys()):
            quant_cols.append(group + " (mean)")
            for sample in sample_groups[group]:
                quant_cols.append(sample)
    else:
        for sample in samples:
            quant_cols.append(sample)
    for sample in quant_cols:
        quant_headline.extend([x + "::" + sample for x in headline_quant])
    quant_data = defaultdict(list)
    for line in iter_file(quant_file):
        fields = line.split("\t")
        quant_data[int(fields[0])].extend(fields[-4:])
    return quant_headline, quant_data


def get_members_count_and_mafft_results_and_members_identifier(mafft_files, pgid_to_accs_dict, pgid):
    #mr = get_mafft_report(pgid)
    try:
        mr = [file for file in mafft_files if pgid == int(Path(file).stem.split('.')[1])][0]
        with open(mr, "r") as handle:
            min_seqlen, max_seqlen, min_ident, max_ident = handle.readlines()[-1].strip().split("\t")
    except IndexError:
        min_seqlen = max_seqlen = min_ident = max_ident = "-"
    accs_per_group = pgid_to_accs_dict[pgid]
    return [str(len(accs_per_group)), min_seqlen, max_seqlen, min_ident, max_ident, ";".join(accs_per_group)]


def get_task_header(taskids, tasks, task_lca_files):
    headline_data = []
    headline_lca_support = []
    tasks_len = []
    for taskid in range(taskids):
        headline_task = [f"task_{taskid}::{tasks[taskid]['shortname'].replace(' ', '_')}::" +
                         x.replace('_', ' ') for x in get_headline(task_lca_files[taskid]).split('\t')]
        headline_data.extend(headline_task[1:])
        headline_lca_support.extend([f"{entry}::lca_support" for entry in headline_task[1:]])
        tasks_len.append(len(headline_task) - 1)
    return tasks_len, headline_data, headline_lca_support


def get_lca_results_per_task(task_lca_file, taskid, data_dict):
    """
    :return: group_task_data[gpid][taskid]: lca
    """
    for line in iter_file(task_lca_file):
        pg_id, *lca_results = line.split("\t")
        data_dict[int(pg_id)][taskid] = lca_results
    return data_dict


def get_member_data_per_task(task_map_file, taskid, task, task_length, accs, member_data):
    """
    # if no member data present, add 'unclassified'
    :return member_data[acc][taskid]=member data (sorted list tuple_lineages)
    """
    for acc in accs:
        member_data[acc][taskid] = set()

    for line in iter_file(task_map_file):
        fields = line.split("\t")
        try:
            member_data[fields[0]][taskid].add(clean_member_data(tuple(fields[-task_length:])))
        except KeyError:
            raise KeyError(f"Task results of task {task} contain accessions not contained in the protein group yaml.")

    for acc in accs:
        # no task results for acc: add unclassified as result
        if member_data[acc][taskid] == set():
            member_data[acc][taskid] = {tuple([INCONCLUSIVE_ANNOTATIONS['unclassified']] * (task_length))}
        # sort results
        member_data[acc][taskid] = list(member_data[acc][taskid])
        member_data[acc][taskid].sort()
    return member_data


def get_support(taskids, task_lca_support_files):
    """
   :return: support[pgid][taskid]: lca_support
   """
    support = defaultdict(dict)
    for taskid in range(taskids):
        support = get_lca_results_per_task(task_lca_support_files[taskid], taskid, support)
    return support


def get_member_data(taskids, task_map_files, task_length_list, tasks, accs):
    member_data = defaultdict(dict)
    for taskid in range(taskids):
        member_data = get_member_data_per_task(task_map_files[taskid], taskid, tasks[taskid], task_length_list[taskid], accs, member_data)
    return member_data


def get_group_data(taskids, task_lca_files, mafft_files, pgid_to_accs_dict):
    group_data = defaultdict(dict)
    for taskid in range(taskids):
        group_data = get_lca_results_per_task(task_lca_files[taskid], taskid, group_data)
        for pgid in pgid_to_accs_dict.keys():
            group_data[pgid]['mafft'] = get_members_count_and_mafft_results_and_members_identifier(mafft_files, pgid_to_accs_dict, pgid)
    return group_data


def write_summary(output, headline_list, pgid_to_group_row, pgid_to_member_rows):
    """
    :param headline_list: list of header elements
    :param: pgid_to_group_row[pgid] = list of group elements
    :param: pgid_to_member_rows[pgid] = list of member_rows (= fully assembled rows = list of list of strings
    """
    with open(output, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile, delimiter='\t')
        writer.writerow(headline_list)
        for pgid in range(len(pgid_to_group_row.keys())):
            writer.writerow(pgid_to_group_row[pgid])
            for row in pgid_to_member_rows[pgid]:
                writer.writerow(row)


def get_headline_list(headline_task_data, headline_quant_data, headline_lca_support, header_spectra=[]):
    headline_list = ["level", "#pg", "members_count", "min_seqlen", "max_seqlen", "min_rel_pairw_ident", "max_rel_pairw_ident", "members_identifier"]
    headline_list.extend(headline_task_data)
    headline_list.extend(headline_quant_data)
    headline_list.extend(header_spectra)
    headline_list.extend(headline_lca_support)
    return headline_list


def get_pg_member_dict(member_data, spectra_IDs_per_accession_dict, pgid_to_accs, taskids, task_length,
                       len_quant, with_spectra):
    """
    :param member_data: member_data[acc][taskid]=member data (sorted list tuple_lineages)
    :param spectra_IDs_per_accession_dict
    :param accs: [string] protein accessions list per protein group
    :param task_range:
    :param task_length:
    :param with_spectra:
    :param len_quant:
    :return pgid_to_member_rows[pgid] = list of member_rows (= fully assembled rows = list of list of strings
"""
    len_mafft = 5
    pgid_to_member_rows = defaultdict(list)
    for pgid in pgid_to_accs.keys():
        member_rows_per_group = []
        for acc in pgid_to_accs[pgid]:
            max_nb_results_per_acc = max([len(member_data[acc][taskid]) for taskid in range(taskids)])
            # tasks with different number of identifications, to fill all member rows, add '-' task results (except error)
            for j in range(max_nb_results_per_acc):
                member_row = ['member', str(pgid)]
                member_row.extend(['-'] * len_mafft)
                member_row.append(acc)
                for taskid in range(taskids):
                    try:
                        member_row.extend(list(member_data[acc][taskid][j]))
                    except IndexError:
                        member_row.extend(['-'] * task_length[taskid])
                member_row.extend(['-'] * len_quant)
                if with_spectra:
                    member_row.extend(spectra_IDs_per_accession_dict[acc])
                member_rows_per_group.append(member_row)
        pgid_to_member_rows[pgid] = member_rows_per_group
    return pgid_to_member_rows


def get_pg_group_dict(group_data, quant_dict, lca_support_data, pgids, taskids, with_spectra, pg_dict):
    """
    :param group_data:
    """
    pgid_to_group_row = {}
    for pgid in range(pgids):
        group_row = ['group', str(pgid)]
        group_row.extend(group_data[pgid]['mafft'])
        for taskid in range(taskids):
            group_row.extend(group_data[pgid][taskid])
        group_row.extend(quant_dict[pgid])
        if with_spectra:
            spectra_list = [spectra.replace('][', ', ').replace('[', '').replace(']', '').split(', ') for spectra in pg_dict[pgid]['spectra']]
            all_spectra_per_group = list(map(sorted, [list(set(spectra_per_group)) for spectra_per_group in spectra_list]))
            group_row.extend([';'.join(l) for l in all_spectra_per_group])
        for taskid in range(taskids):
            group_row.extend(lca_support_data[pgid][taskid])
        pgid_to_group_row[pgid] = group_row
    return pgid_to_group_row


def create_summary(pg_yaml, task_lca_files, task_lca_support_files, task_map_files, quant_file, mafft_files,
                   output, tasks, sample_groups):
    pg_dict = get_pg_dict(pg_yaml)
    samples = get_samples(pg_yaml)
    taskids = len(tasks)
    with_spectra = is_with_spectra(pg_yaml)
    pgid_to_accs_dict = get_accs_per_pg_from_yaml(pg_yaml)
    accs = [acc for accs in pgid_to_accs_dict.values() for acc in accs]
    total_number_of_pgs = len(pgid_to_accs_dict.keys())
    task_length_list, headline_task_data, headline_lca_support = get_task_header(taskids, tasks, task_lca_files)
    lca_support_data = get_support(taskids, task_lca_support_files)
    member_data = get_member_data(taskids, task_map_files, task_length_list, tasks, accs)
    group_data = get_group_data(taskids, task_lca_files, mafft_files, pgid_to_accs_dict)
    headline_quant_data, quant_dict = get_quantdata(quant_file, samples, sample_groups)

    if with_spectra:
        header_spectra, spectra_IDs_per_accession_dict = get_spectra_header_and_accs_to_spectraID_dict(pg_yaml)
    else:
        spectra_IDs_per_accession_dict = {}
        header_spectra = []

    headline_list = get_headline_list(headline_task_data, headline_quant_data, headline_lca_support, header_spectra)
    pgid_to_group_row = get_pg_group_dict(group_data, quant_dict, lca_support_data, total_number_of_pgs, taskids, with_spectra, pg_dict)
    pgid_to_member_rows = get_pg_member_dict(member_data, spectra_IDs_per_accession_dict, pgid_to_accs_dict, taskids,
                                             task_length_list, len(quant_dict[0]), with_spectra)

    write_summary(output, headline_list, pgid_to_group_row, pgid_to_member_rows)

