import numpy as np

from utils.protein_group_sql_db import ProteinGroupSqlDbInteractor


def calc_nsafs(scounts, protlens):
    safs = [x[0] / x[1] for x in zip(scounts, protlens)]
    total = sum(safs)
    return [x/total for x in safs]

rule quant_pg:
  input:
    fasta = 'seqs/all.faa',
    pg_db = 'pgs/protein_groups_db.sql'
  output:
    'tasks/quant.tsv'
  message:
    "creating FASTA files for protein groups ..."
  version:
    "0.1"
  params:
    quant_methods = {"min_nsaf", "max_nsaf", "mean_nsaf", "raw"}
  run:
    pg_db_obj = ProteinGroupSqlDbInteractor(input.pg_db)
    sep = pg_db_obj.get_protein_sep()
    quant_method = config['quant_method']
    if quant_method not in params.quant_methods:
        trace("config error:\nunkown quantification method (quant_method=" + quant_method + ")", True)

    #store raw quant data
    samples = pg_db_obj.get_samples()
    pg_members = []
    pg_raw_quant = []
    all_members = set()
    pg_id_to_accs_dict= pg_db_obj.get_pg_id_to_accs_dict()
    pg_id_to_quant_dict= pg_db_obj.get_quant_values_of_pgs()
    for pgid in pg_id_to_accs_dict.keys():
        pg_members.append(pg_id_to_accs_dict[pgid].split(sep))
        all_members.update(pg_members[-1])
        pg_raw_quant.append(pg_id_to_quant_dict[pgid])

    #read seqs
    seqlens = {}
    for line in iter_file(input.fasta):
        if line[0] == ">":
            acc = line[1:]
        elif acc in all_members:
            l = len(line)
            if acc in seqlens and l != seqlens[acc]:
                trace("quantification error:\n"+ 'ambigious sequences for ' + acc + " in " + input.fasta, True)
            seqlens[acc] = l

    #get member seqlens
    pg_lens = []
    final_quant = []
    for i in range(len(pg_members)):
        pg_lens.append([seqlens[x] for x in pg_members[i]])
        final_quant.append([])

    if quant_method == "min_nsaf":
        processed_pg_lens = [min(x) for x in pg_lens]
    elif quant_method == "max_nsaf":
        processed_pg_lens = [max(x) for x in pg_lens]
    elif quant_method == "mean_nsaf":
        processed_pg_lens = [sum(x)/len(x) for x in pg_lens]

    #calc quant
    if quant_method == "raw":
        final_quant = pg_raw_quant
    else:
        for i in range(len(samples)):
            raw_quants = [x[i] for x in pg_raw_quant]
            i = 0
            for nsaf in calc_nsafs(raw_quants, processed_pg_lens):
                final_quant[i].append(nsaf)
                i += 1

    #get sample indices for mean quant
    if 'sample_groups' in config:
        sg_indices = {}
        groups = sorted(config['sample_groups'].keys())
        unknown_samples = set()
        for sg, samplelist in config['sample_groups'].items():
            if sg not in sg_indices:
                sg_indices[sg] = []
            for sample in samplelist:
                if sample not in samples:
                    unknown_samples.add(sample)
                else:
                    sg_indices[sg].append(samples.index(sample))
        if unknown_samples:
            print("error:")
            print("defined samples:   " + ", ".join(sorted(samples)))
            print("undefined samples: " + ", ".join(sorted(unknown_samples)))
            exit("error: undefined samples: " + ", ".join(sorted(unknown_samples)))

    #calc mean quant and write out
    with open(output[0], "w") as handle:
        handle.write("#pg\tsample\tmin_seqlen\tmax_seqlen\traw_quant\traw_quant_sd\tquant\tquant_sd\n")
        for i in range(len(final_quant)):
            if 'sample_groups' in config:
                for g in groups:
                    g_raw_quant = [pg_raw_quant[i][x] for x in sg_indices[g]]
                    g_quant = [final_quant[i][x] for x in sg_indices[g]]
                    mean_raw_quant = np.mean(g_raw_quant)
                    mean_quant = np.mean(g_quant)
                    std_raw_quant = np.std(g_raw_quant)
                    std_quant = np.std(g_quant)
                    handle.write(f"{i}\t{g} (mean)\t-\t-\t{mean_raw_quant}\t{std_raw_quant}\t{mean_quant}\t{std_quant}\n")
            for j in range(len(samples)):
                handle.write(f"{i}\t{samples[j]}\t{min(pg_lens[i])}\t{max(pg_lens[i])}\t{pg_raw_quant[i][j]}\t-\t{final_quant[i][j]}\t-\n")
