#!/usr/bin/env python3
"""Sensor map from every heating pattern at once (main + scan sessions), with a heat-spread model.

Observed: each sensor's response to each pattern = its mean change over t in [6, 9] s minus the die mean.
Predicted: for a sensor at cell p, sum over heated cells q of exp(-0.77 |p - q|) (the lab's model: 0.52 C at
the source, 0.24 one tile away, 0.10 two away), minus the mean over the 35 sensor cells.
Score(sensor, cell) = correlation of observed and predicted over the patterns; an optimal one-to-one
assignment (Hungarian) places the 35 sensors. Writes map2.json.
"""
import json
import math
import sys

sys.path.insert(0, __file__.rsplit("/", 1)[0])
from sensormap import CELLS, GRID, hungarian  # noqa: E402

DECAY = 0.77


def corr(a, b):
    n = len(a)
    ma, mb = sum(a) / n, sum(b) / n
    va = sum((x - ma) ** 2 for x in a)
    vb = sum((y - mb) ** 2 for y in b)
    if va == 0 or vb == 0:
        return 0.0
    return sum((x - ma) * (y - mb) for x, y in zip(a, b)) / math.sqrt(va * vb)


def main(data_paths, out_json):
    runs, times = [], None
    for p in data_paths.split(","):
        D = json.load(open(p))
        times = D["times"]
        runs += D["runs"]
    win = [i for i, t in enumerate(times) if 6.0 <= t <= 9.0]
    pos_of_shire = {}
    for r, row in enumerate(GRID):
        for c, name in enumerate(row):
            if name.startswith("S"):
                pos_of_shire[int(name[1:])] = (r, c)
    pats = [r for r in runs if r["per_shire"] == 32]  # the single core heats nothing measurable
    obs = []  # obs[k][s]
    pred_cell = []  # pred_cell[k][cell]
    for R in pats:
        fr = R["frames"]
        m = [sum(fr[i][s] for i in win) / len(win) for s in range(35)]
        die = sum(m[:34]) / 34
        obs.append([x - die for x in m])
        heat = [pos_of_shire[s] for s in R["shires"]]
        f = [sum(math.exp(-DECAY * math.hypot(r - hr, c - hc)) for hr, hc in heat) for (r, c, _a) in CELLS]
        mf = sum(f) / len(f)
        pred_cell.append([x - mf for x in f])
    K = len(pats)
    score = [[corr([obs[k][s] for k in range(K)], [pred_cell[k][j] for k in range(K)]) for j in range(len(CELLS))]
             for s in range(35)]
    assign = hungarian([[-score[s][j] for j in range(len(CELLS))] for s in range(35)])
    ident = {a: j for j, (_r, _c, a) in enumerate(CELLS)}
    per, agree, near = [], 0, 0
    for s in range(35):
        j = assign[s]
        r, c, assumed = CELLS[j]
        ar, ac, _ = CELLS[ident[s]]
        d = abs(r - ar) + abs(c - ac)
        agree += d == 0
        near += d <= 1
        ranked = sorted(range(len(CELLS)), key=lambda x: -score[s][x])
        per.append({"sensor": s, "assumed": [ar, ac], "placed": [r, c], "agrees": d == 0,
                    "score_assumed": round(score[s][ident[s]], 3), "score_placed": round(score[s][j], 3),
                    "rank_of_assumed": ranked.index(ident[s]) + 1, "best_cell": list(CELLS[ranked[0]][:2])})
    tot_id = sum(score[s][ident[s]] for s in range(35))
    tot_opt = sum(score[s][assign[s]] for s in range(35))
    # null: the identity's total against random relabelings of the sensors
    import random
    rnd = random.Random(7)
    null = []
    for _ in range(2000):
        perm = list(range(35))
        rnd.shuffle(perm)
        null.append(sum(score[s][ident[perm[s]]] for s in range(35)))
    p_val = sum(1 for x in null if x >= tot_id) / len(null)
    grid_emp = [["P" if GRID[r][c] == "P" else "" for c in range(6)] for r in range(6)]
    for p in per:
        s = p["sensor"]
        grid_emp[p["placed"][0]][p["placed"][1]] = "IO" if s == 34 else ("M%d" % s if s >= 32 else "S%d" % s)
    out = {"patterns": [r["name"] for r in pats], "agree": agree, "near": near, "n": 35,
           "total_identity": round(tot_id, 2), "total_optimal": round(tot_opt, 2),
           "null_mean": round(sum(null) / len(null), 2), "null_p": p_val,
           "median_rank_assumed": sorted(p["rank_of_assumed"] for p in per)[17],
           "assumed_top3": sum(1 for p in per if p["rank_of_assumed"] <= 3),
           "grid_empirical": grid_emp, "grid_assumed": GRID, "per_sensor": per}
    json.dump(out, open(out_json, "w"))
    print("patterns %d; identity total corr %.2f vs optimal %.2f; random relabelings mean %.2f, p=%.4f"
          % (K, tot_id, tot_opt, out["null_mean"], p_val))
    print("optimal placement agrees with the assumed cell for %d/35, within one cell for %d/35; "
          "assumed cell ranks in each sensor's top 3 for %d/35 (median rank %d of 35)"
          % (agree, near, out["assumed_top3"], out["median_rank_assumed"]))


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])