#!/usr/bin/env python3
"""Empirical sensor map from the column and row scans (analyze.py's data.json for the scan session).

Each sensor's response to a pattern = its mean change over t in [6, 9] s (the end of heating) minus the
whole-die mean change. A sensor's column is where its column responses peak, its row where its row
responses peak; an optimal one-to-one assignment (Hungarian) places all 35 sensors on the 35 cells that
have one. Writes map.json.
"""
import json
import sys

GRID = [["S0", "S8", "S3", "M32", "IO", "P"],
        ["S24", "S16", "S4", "S12", "S20", "S28"],
        ["S9", "S1", "S13", "S21", "S29", "S5"],
        ["S25", "S17", "S14", "S22", "S30", "S6"],
        ["S2", "S10", "S18", "S26", "S15", "S7"],
        ["S11", "S19", "S27", "M33", "S23", "S31"]]
CELLS = []  # (row, col, assumed sensor index); the PCIe cell has no sensor; IO/PCIe are an unknown pair
for r, row in enumerate(GRID):
    for c, name in enumerate(row):
        if name == "P":
            continue
        idx = 34 if name == "IO" else int(name[1:])
        CELLS.append((r, c, idx))


def hungarian(cost):
    """Minimum-cost assignment for a square matrix (list of lists). Returns row -> col."""
    n = len(cost)
    INF = float("inf")
    u, v, p, way = [0.0] * (n + 1), [0.0] * (n + 1), [0] * (n + 1), [0] * (n + 1)
    for i in range(1, n + 1):
        p[0] = i
        j0 = 0
        minv = [INF] * (n + 1)
        used = [False] * (n + 1)
        while True:
            used[j0] = True
            i0, delta, j1 = p[j0], INF, 0
            for j in range(1, n + 1):
                if not used[j]:
                    cur = cost[i0 - 1][j - 1] - u[i0] - v[j]
                    if cur < minv[j]:
                        minv[j], way[j] = cur, j0
                    if minv[j] < delta:
                        delta, j1 = minv[j], j
            for j in range(n + 1):
                if used[j]:
                    u[p[j]] += delta
                    v[j] -= delta
                else:
                    minv[j] -= delta
            j0 = j1
            if p[j0] == 0:
                break
        while True:
            j1 = way[j0]
            p[j0] = p[j1]
            j0 = j1
            if j0 == 0:
                break
    res = [0] * n
    for j in range(1, n + 1):
        res[p[j] - 1] = j - 1
    return res


def main(scan_json, out_json):
    D = json.load(open(scan_json))
    times = D["times"]
    win = [i for i, t in enumerate(times) if 6.0 <= t <= 9.0]
    runs = {r["name"]: r for r in D["runs"]}

    def response(name):
        fr = runs[name]["frames"]
        mean = [sum(fr[i][s] for i in win) / len(win) for s in range(35)]
        die = sum(mean[:34]) / 34
        return [m - die for m in mean], [round(m, 3) for m in mean]

    col_names = ["col-c%d" % k for k in range(1, 7)]
    row_names = ["row-r%d" % k for k in range(6)]
    RC = {n: response(n) for n in col_names + row_names if n in runs}
    colR = [RC[n][0] for n in col_names]
    rowR = [RC[n][0] for n in row_names]
    sensors = list(range(35))
    # score(sensor, cell) = column response + row response at that cell; maximize -> minimize negative
    cost = [[-(colR[c][s] + rowR[r][s]) for (r, c, _a) in CELLS] for s in sensors]
    assign = hungarian(cost)
    placed, agree = {}, 0
    for s in sensors:
        r, c, assumed = CELLS[assign[s]]
        placed[s] = (r, c)
        if assumed == s:
            agree += 1
    # per-sensor argmax column / row, and their margins
    per = []
    for s in sensors:
        cols = [colR[c][s] for c in range(6)]
        rows = [rowR[r][s] for r in range(6)]
        bc = max(range(6), key=lambda c: cols[c])
        br = max(range(6), key=lambda r: rows[r])
        sc = sorted(cols, reverse=True)
        sr = sorted(rows, reverse=True)
        a_r, a_c = next((r, c) for (r, c, a) in CELLS if a == s)
        per.append({"sensor": s, "assumed": [a_r, a_c], "placed": list(placed[s]), "best_col": bc, "best_row": br,
                    "col_margin": round(sc[0] - sc[1], 3), "row_margin": round(sr[0] - sr[1], 3),
                    "col_resp": [round(x, 3) for x in cols], "row_resp": [round(x, 3) for x in rows],
                    "agrees": [a_r, a_c] == list(placed[s])})
    near = sum(1 for p in per if abs(p["placed"][0] - p["assumed"][0]) + abs(p["placed"][1] - p["assumed"][1]) <= 1)
    col_ok = sum(1 for p in per if p["best_col"] == p["assumed"][1])
    row_ok = sum(1 for p in per if p["best_row"] == p["assumed"][0])
    grid_emp = [["P" if GRID[r][c] == "P" else "" for c in range(6)] for r in range(6)]
    for s, (r, c) in placed.items():
        grid_emp[r][c] = "IO" if s == 34 else ("M%d" % s if s >= 32 else "S%d" % s)
    out = {"grid_assumed": GRID, "grid_empirical": grid_emp, "agree": agree, "near": near, "n": 35,
           "best_col_ok": col_ok, "best_row_ok": row_ok, "per_sensor": per,
           "col_maps": {n: RC[n][1] for n in col_names}, "row_maps": {n: RC[n][1] for n in row_names}}
    json.dump(out, open(out_json, "w"))
    print("sensors placed at their assumed cell: %d/35; within one cell: %d/35; best column right %d/35, best row right %d/35"
          % (agree, near, col_ok, row_ok))
    for p in per:
        if not p["agrees"]:
            print("  sensor %d: assumed r%dc%d, placed r%dc%d (best col c%d margin %.2f, best row r%d margin %.2f)"
                  % (p["sensor"], p["assumed"][0], p["assumed"][1] + 1, p["placed"][0], p["placed"][1] + 1,
                     p["best_col"] + 1, p["col_margin"], p["best_row"], p["row_margin"]))


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