import numpy as np

from utils.protein_group_sql_db import ProteinGroupSqlDbInteractor


rule quant_task:
  input:
    pg_db='pgs/protein_groups_db.sql',
    quant='tasks/quant.tsv',
    lca="tasks/{annot_type}_annot_by_{tool}_on_{db_type}.v{db_version}.task{taskid}.lca"
  output:
    'tasks/{annot_type}_annot_by_{tool}_on_{db_type}.v{db_version}.task{taskid}.quant'
  message:
    "add quantification to task ..."
  version:
    "0.1"
  params:
    quant_methods = {"min_nsaf", "max_nsaf", "mean_nsaf", "raw"}
  run:

    samples = ProteinGroupSqlDbInteractor(input.pg_db).get_samples()

    #read quant
    quants = {}
    with open(input[1], "r") as handle:
        for line in handle:
            line = line.strip()
            if len(line) == 0 or line[0] == "#":
                continue
            fields = line.split("\t")
            pg = fields[0]
            if pg not in quants:
                quants[pg] = {}
            quants[pg][fields[1]] = fields[-2]

    #read lca
    lcas = {}
    with open(input[2], "r") as handle:
        headline = handle.readline().strip().split("\t")
        headline = "\t".join(headline[1:])
        for line in handle:
            line = line.strip()
            if len(line) == 0:
                continue
            fields = line.split("\t")
            pg = fields[0]
            lcas[fields[0]] = fields[1:]

    #calc
    lca_quant = {}

    for pg, lca in lcas.items():
        for sample in samples:
            lca_str = "\t".join(lca)
            if lca_str not in lca_quant:
                lca_quant[lca_str] = {}
            if sample not in lca_quant[lca_str]:
                lca_quant[lca_str][sample] = 0
            lca_quant[lca_str][sample] += float(quants[pg][sample])
    lca_order = sorted(lca_quant.keys())

    if 'sample_groups' in config:
        sg_indices = {}
        groups = sorted(config['sample_groups'].keys())
        for sg, samplelist in config['sample_groups'].items():
            sg_indices[sg] = [samples.index(x) for x in set(samplelist)]

    with open(output[0], "w") as handle:
        handle.write("sample\t" + headline + "\tquant\tsd\n")
        if 'sample_groups' in config:
            for g in groups:
                for lca in lca_order:
                    g_quant = [lca_quant[lca][x] for x in config['sample_groups'][g]]
                    quant_mean = np.mean(g_quant)
                    quant_std=np.std(g_quant)
                    handle.write(f"{g} (mean)\t{lca}\t{quant_mean}\t{quant_std}\n")
        for sample in samples:
            for lca in lca_order:
                quant = lca_quant[lca][sample]
                handle.write(f"{sample}\t{lca}\t{quant}\t-\n")
