#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""CPU emulation of the TT graph's numerics (no device): the real ``tt_sceneseg.tt`` graph executed by the fake ttnn
of ``common/tests/host`` (bf16-rounded operands and outputs, float64 products and sums), one and two weight terms, on
a few frames, against the fp32 CPU reference goldens: class-map agreement, foreground IoU, logits PCC. The evidence
behind the weight-terms default (PORT_LOG.md M1), next to the device runs.

    source bin/tt-env.sh; cd bundles/sceneseg-p150
    PYTHONPATH=$PWD/code:$PYTHONPATH TT_VISIBLE_DEVICES=none OMP_NUM_THREADS=4 python code/scripts/diag/ss_precision_emulation.py \
        --frames tut_highway_rain_2,ps019_f40_front_camera --terms 1,2 --json ../../logs/sceneseg/m1_precision_emulation.json
"""
from __future__ import annotations

import argparse
import importlib.util
import json
import sys
import time
from pathlib import Path

import numpy as np

BUNDLE = Path(__file__).resolve().parents[3]
WS = BUNDLE.parents[1]
sys.path.insert(0, str(BUNDLE / "code"))


def install_fake():
    """The fake ttnn (+ its CNN ops and a row sum) as ``ttnn``; never the real one."""
    host = WS / "common" / "tests" / "host"
    spec = importlib.util.spec_from_file_location("fake_ttnn", host / "fake_ttnn.py")
    fake = importlib.util.module_from_spec(spec)
    sys.modules["fake_ttnn"] = fake
    spec.loader.exec_module(fake)
    ttnn = fake.install()
    spec = importlib.util.spec_from_file_location("fake_ttnn_cnn", host / "fake_ttnn_cnn.py")
    cnn = importlib.util.module_from_spec(spec)
    sys.modules["fake_ttnn_cnn"] = cnn
    spec.loader.exec_module(cnn)
    cnn.install(ttnn)

    def sum_(t, dim=None, keepdim=False, **_):
        return ttnn._op("sum", [t], lambda x: x.sum(axis=dim, keepdims=keepdim), t.dtype, t.layout, (dim, keepdim))

    ttnn.sum = sum_
    return ttnn


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--frames", default="tut_highway_rain_2,ps019_f40_front_camera")
    ap.add_argument("--terms", default="1,2")
    ap.add_argument("--order", default="bgr")
    ap.add_argument("--json", default="")
    a = ap.parse_args()
    ttnn = install_fake()
    if not hasattr(ttnn, "subtract"):          # the activation term split (tt/layers.split_activation)
        ttnn.subtract = lambda x, y, **_: ttnn._op("subtract", [x, y], lambda u, v: u - v, x.dtype, x.layout)
    import torch

    torch.set_num_threads(4)
    from tt_sceneseg.host.preprocess import prepare_bgr
    from tt_sceneseg.reference import assets
    from tt_sceneseg.reference.weights import load_weights
    from tt_sceneseg.tests import sceneseg_eval as ev
    from tt_sceneseg.tt.model import TtSceneSeg
    from tt_sceneseg.ttaw.device import open_device
    from tt_sceneseg.ttaw.metrics import pcc

    dev = open_device(0, dispatch="eth", num_command_queues=1)
    weights = load_weights(assets.onnx_path())
    gold = WS / "research" / "sceneseg" / "goldens" / "public"
    inputs = WS / "research" / "sceneseg" / "public_data" / "inputs"
    results = {}
    for terms in (int(t) for t in a.terms.split(",")):
        model = TtSceneSeg(weights, a.order, weight_terms=terms)
        for fid in a.frames.split(","):
            t0 = time.perf_counter()
            bgr = np.ascontiguousarray(ev.load_rgb(inputs / fid / "source_bgr.png")[:, :, ::-1])
            canvas = prepare_bgr(bgr, channel_order=a.order).device_input()
            image = ttnn.from_torch(torch.from_numpy(canvas), dtype=ttnn.uint8, layout=ttnn.ROW_MAJOR_LAYOUT,
                                    device=dev)
            out = model.forward(image, logits=True)
            cmap = ttnn.to_torch(out["class_map"]).numpy().reshape(320, 640)
            logits = ttnn.to_torch(out["logits"]).float().numpy()
            with np.load(gold / f"{fid}.npz") as z:
                ref_logits, ref_cls = z[f"logits_{a.order}"], z[f"cls_{a.order}"]
            rec = {"logits_pcc": pcc(logits, ref_logits), "agreement": ev.argmax_agreement(cmap, ref_cls),
                   "foreground_iou": ev.foreground_iou(cmap, ref_cls),
                   "fg_fraction": float((cmap == 1).mean()), "fg_fraction_fp32": float((ref_cls == 1).mean()),
                   "seconds": round(time.perf_counter() - t0, 1)}
            results[f"terms{terms}/{fid}"] = rec
            print(f"terms={terms} {fid}: {json.dumps(rec)}", flush=True)
        model.release()
    if a.json:
        Path(a.json).write_text(json.dumps({"order": a.order, "results": results}, indent=1) + "\n")
    return 0


if __name__ == "__main__":
    sys.exit(main())