#!/usr/bin/env python
# build_score_app_data.py — data layer for the interactive results page.
# WXS and RNA are each scored as INDEPENDENT standalone groups; the 3 genomic
# assays (WGS/INDELS/META) form their own 7 combos. No cross-group mixes.
# Two k-modes: FIXED (k1 in {21,31,51}) and FREE (all k1 in the sweep). For each
# (entity, k-mode, scale-filter, score-flavor) we take the BEST-ACHIEVABLE score
# (argmax over the allowed config subset) and record the winning config.
# Emits score_app_data.json (embedded into the HTML) + composite_full_all_k.tsv.
import os, json
import numpy as np
import pandas as pd
import build_composite_score as b

HERE = os.path.dirname(os.path.abspath(__file__))
GENOMIC = ["WGS_SIM", "INDELS_SIM", "METAGENOMIC_SIM"]
WXS = "WXS_SIM"
RNA = "RNA_SIM"
FLAVORS = ["score_equal", "score_seq_weighted",
           "score_seq_change", "score_depth", "score_mapping", "score_breadth"]
SCALES = [1000, 10000, 100000]


def compute_full():
    """Composite over ALL configs (no k1 restriction) for the 7 genomic combos + WXS + RNA."""
    df = b.load(restrict_k1=False)
    atomic = b.build_atomic(df)
    combos = b.all_combos(GENOMIC) + [(WXS,)] + [(RNA,)]
    comp = b.build_composite(atomic, combos=combos)
    def grp(label):
        if label == b.SHORT[RNA]: return "rna"
        if label == b.SHORT[WXS]: return "wxs"
        return "genomic"
    comp["group"] = comp["combo"].map(grp)
    return comp, atomic


def best_records(comp):
    """For each (combo, k-mode, scale-filter, flavor): best-achievable score + config."""
    recs = []
    for combo, sub in comp.groupby("combo"):
        group = sub["group"].iloc[0]; n_assays = int(sub["n_assays"].iloc[0])
        for kmode, k1_allowed in (("fixed", b.K1S), ("free", None)):
            km = sub if k1_allowed is None else sub[sub["k1"].isin(k1_allowed)]
            for scale in ["all"] + SCALES:
                ss = km if scale == "all" else km[km["scale"] == scale]
                for fl in FLAVORS:
                    v = ss[fl].dropna()
                    if v.empty:
                        continue
                    i = v.idxmax(); row = ss.loc[i]
                    recs.append(dict(
                        group=group, combo=combo, n_assays=n_assays,
                        kmode=kmode, scale=str(scale), flavor=fl,
                        score=round(float(row[fl]), 4),
                        k1=int(row["k1"]), k2_ext=int(row["k2_ext"]), win_scale=int(row["scale"])))
    return recs


def main():
    comp, atomic = compute_full()
    recs = best_records(comp)
    combo_order = (comp[["combo", "n_assays", "group"]].drop_duplicates()
                   .sort_values(["group", "n_assays", "combo"]))
    data = {
        "meta": {
            "reference": "Canis familiaris CanFam3.1 (GCF_000002285.3)",
            "n_configs_full": int(comp[["k1", "k2_ext", "scale"]].drop_duplicates().shape[0]),
            "k1_all": sorted(int(x) for x in comp["k1"].unique()),
            "k1_fixed": b.K1S, "k2_ext": sorted(int(x) for x in comp["k2_ext"].unique()),
            "scales": SCALES, "weights_seq_change": b.WEIGHTS["seq_change"],
        },
        "flavors": FLAVORS,
        "genomic_combos": combo_order[combo_order.group == "genomic"]["combo"].tolist(),
        "wxs_combo": b.SHORT[WXS],
        "rna_combo": b.SHORT[RNA],
        "records": recs,
        "full_columns": ["group", "combo", "n_assays", "k1", "k2_ext", "scale"] + b.SCORE_COLS,
        "full_rows": comp[["group", "combo", "n_assays", "k1", "k2_ext", "scale"] + b.SCORE_COLS]
                 .round(4).values.tolist(),
    }
    with open(os.path.join(HERE, "score_app_data.json"), "w") as f:
        json.dump(data, f, separators=(",", ":"))
    comp.round(4).to_csv(os.path.join(HERE, "composite_full_all_k.tsv"), sep="\t", index=False)
    print(f"wrote score_app_data.json  ({len(recs)} best-records, {len(data['full_rows'])} full rows)")
    print("genomic combos:", len(data["genomic_combos"]), "| wxs:", data["wxs_combo"],
          "| rna:", data["rna_combo"], "| full configs:", data["meta"]["n_configs_full"])


if __name__ == "__main__":
    main()