import subprocess
from contextlib import ExitStack
from pathlib import Path

from utils.lca_determination.weighted_lca_class import WeightedLCA


def write_lca_results_in_new_summary(lca_files_opened, lca_support_files_opened, path_output, reader, sep):
    for row in reader.df.values.tolist():
        row = [str(r) for r in row]
        if row[reader.member_pos] == 'group':
            lca_row = []
            support_row = []
            for taskid in range(len(reader.tasks)):
                lca_row.extend(lca_files_opened[taskid].readline().rstrip().split('\t')[1:])
                support_row.extend(lca_support_files_opened[taskid].readline().rstrip().split('\t')[1:])
            row[reader.task_indices[0][0]:reader.task_indices[-1][1]] = lca_row
            row.extend(support_row)
        with open(path_output, 'a') as output:
            output.write(sep.join(row) + '\n')


def write_summary(path_output, reader,  sep):
    """
    ## write summary with new lca starting from old summary
    :param path_output: path for summary_output
    :param reader: object containing read summary
    :param sep: separator between columns (e.g.'\t')
    """
    path_output = Path(path_output)
    summary_header = list(reader.df)
    lca_support_header = [header_item + '::lca-support' for header_item in summary_header if header_item.startswith('task')]

    lca_files = [f"{path_output.parents[0]}/{path_output.stem}_task_{taskid}.lca.tmp" for taskid in range(len(reader.tasks))]
    lca_support_files = [f"{path_output.parents[0]}/{path_output.stem}_task_{taskid}.lca_support.tmp" for taskid in range(len(reader.tasks))]

    with open(path_output, 'w') as output:
        output.write(sep.join(summary_header) + '\t' + sep.join(lca_support_header) + '\n')
    # open tmp lca files written by lca.weight_and_write_lca(lca_output, lca_support_output)
    with ExitStack() as stack_lca, ExitStack() as stack_lcasupport:
        lca_files_opened = [stack_lca.enter_context(open(f)) for f in lca_files]
        lca_support_files_opened = [stack_lcasupport.enter_context(open(f)) for f in lca_support_files]
        # skip header line
        for lca_file, lca_support_file in zip(lca_files_opened, lca_support_files_opened):
            lca_file.readline()
            lca_support_file.readline()
        # write lca results (all tasks) in new summary file
        write_lca_results_in_new_summary(lca_files_opened, lca_support_files_opened, path_output, reader, sep)
    # remove tmp lca files
    for tmp_file in lca_files:
        subprocess.call('rm %s' % (str(tmp_file)), shell=True)
    for tmp_file in lca_support_files:
        subprocess.call('rm %s' % (str(tmp_file)), shell=True)


def clean_empty_lineages(acc_to_lineages_dict):
    # set empty lineages to []
    for acc, lineages in acc_to_lineages_dict.items():
        lineages_checked = []
        for lineage in lineages:
            if not set(lineage) in [{'-'}, {' '}, {'nan'}, {'-', ' '}, {' ', 'nan'}, {' ', 'nan', '-'}]:
                lineages_checked.append(lineage)
        acc_to_lineages_dict[acc] = lineages_checked
    return acc_to_lineages_dict


def get_chunks(lst_of_pgs, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst_of_pgs), n):
        yield lst_of_pgs[i:i + n]


def create_lca_from_summary(output, reader, threshold, method, pgid_to_accs_dict, accs_to_spectra_dict,
                            task_to_acc_to_lineages_dict, path_to_taxmap, sep='\t', ignore_unclassified=False,
                            minimum_number_of_annotations=0):
    """
    :param output: write summary with new lca to output
    :param reader: Object reader containing all information from summary
    :param threshold: threshold for group_lca, float value given in config
    :param method: lca method ('lca' or 'democratic lca' given in config)
    :param pgid_to_accs_dict:
    :param accs_to_spectra_dict:
    :param task_to_acc_to_lineages_dict
    :param path_to_taxmap
    :param ignore_unclassified: True -> remove all complete unclassified lineages before analysis
    :param minimum_number_of_annotations: int minimal number of classified entries contained in lineage, remove lineages with less
    :param sep: separator (e.g.'\t') between columns
    # in summary replace group rows in task columns and add lca_support
            ...
    """
    output = Path(output)
    for taskid in range(len(reader.tasks)):
        lca_output = output.parent/f"{output.stem}_task_{taskid}.lca.tmp"
        lca_support_output = output.parent/f"{output.stem}_task_{taskid}.lca_support.tmp"
        level_names = [task_name.split('::')[-1] for task_name in list(reader.df) if task_name.startswith(f"task_{taskid}")]
        lca_headline = '\t'.join(level_names)
        WeightedLCA.write_headers(lca_output, lca_support_output, lca_headline)

        acc_to_lineages_dict = task_to_acc_to_lineages_dict[taskid]
        acc_to_lineages_dict = clean_empty_lineages(acc_to_lineages_dict)

        is_taxonomic_annotation = True if level_names[0] in \
                              ['superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'] else False
        for pgid_range in list(get_chunks(range(0, len(pgid_to_accs_dict)), 1000)):
            lca = WeightedLCA(level_names, acc_to_lineages_dict, pgid_to_accs_dict, accs_to_spectra_dict,
                              method, path_to_taxmap, is_taxonomic_annotation, threshold, ignore_unclassified,
                              minimum_number_of_annotations)
            pgid_to_lcas_dict, pgid_to_supports_dict = lca.weight_lca(list(pgid_range))
            WeightedLCA.write_part_lca_result(pgid_to_lcas_dict, lca_output)
            WeightedLCA.write_part_lca_result(pgid_to_supports_dict, lca_support_output)

    write_summary(output, reader, sep)
