#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device diagnostic (M1b follow-up): where does the device lose accuracy outside the point island?

    cd bundles/frnet-p150 && bin/devrun -t 1200 -- env PYTHONPATH=code python code/scripts/diag_conv_precision.py \\
        --out /home/ubuntu/experiments/tt-models/logs/frnet/m1b/diag_conv_precision.json

Teacher-forced units on the OT128 sample's golden taps (torch fp32 reference): each unit's device output is compared
with the golden output (``dev``) and with a CPU emulation of the device's rounding points with exact fp32 weights
(``emu``: bf16 inputs, conv outputs rounded to bf16 before / after each activation, bf16 adds); ``dev_vs_emu``
rel-L2 measures the device's arithmetic error beyond the emulated bf16 roundings (0 for a perfect fp32 conv engine).
Units: the stem (3 convs + HardSwish), every stage's BasicBlocks (stage input -> ``layer<l>``), every attention
(``pixel_fusion<l>`` -> ``attention<l>``), the ``addcmul`` (-> ``stage<l>``), ``fuse2``, the softmax, the first conv
of each stage alone (no activation), and the element-wise ops (HardSwish, sigmoid) on bf16 inputs. Configurations:
``FRNET_CONV_TERMS`` 1 (float32 conv weights: the device's TF32-like path) and 2 (C17 two-term bf16 weights).
Eager runs (no trace): a diagnostic, not a gate.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve()
sys.path.insert(0, str(HERE.parents[1]))


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--terms", nargs="*", type=int, default=[1, 2])
    args = ap.parse_args()

    import torch
    import torch.nn.functional as F

    torch.set_num_threads(4)
    from tt_frnet.device import device_session
    from tt_frnet.reference.config import FRNetConfig
    from tt_frnet.reference.weights import load_weights
    from tt_frnet.tests._data import golden_dir, weights_dir
    from tt_frnet.tt.layers import free
    from tt_frnet.tt.model import KNOBS, LOGIT_PAD, TtFRNet
    from tt_frnet.ttaw.metrics import error_stats, pcc
    from tt_frnet.ttaw.ops.conv import FeatureMap, fmap_from_numpy, fmap_to_numpy
    from tt_frnet.ttaw.tensors import round_to_bf16

    variant = "ot128"
    wdir = weights_dir()
    cfg = FRNetConfig.from_model_dir(wdir, variant)
    weights = load_weights(wdir, variant)
    taps = np.load(golden_dir() / variant / "pandaset_019_40.taps.npz", mmap_mode="r")

    def g(name):
        return np.asarray(taps[name], np.float32)

    def bf(a):
        return round_to_bf16(np.asarray(a, np.float32))

    def t(a):
        return torch.from_numpy(np.ascontiguousarray(a, np.float32))

    # ---- CPU emulation of the device rounding points (exact fp32 weights)
    def emu_conv(x, layer, act=True):
        y = F.conv2d(t(x), t(layer.weight), t(layer.bias), stride=layer.stride, padding=layer.padding).numpy()
        y = bf(y)
        if act and layer.activation == "hardswish":
            y = bf(F.hardswish(t(y)).numpy())
        elif act and layer.activation == "sigmoid":
            y = bf(torch.sigmoid(t(y)).numpy())
        return y

    def emu_blocks(x, st):
        for c1, c2, ds in st.blocks:
            y = emu_conv(emu_conv(x, c1), c2)
            idt = emu_conv(x, ds) if ds is not None else x
            x = bf(F.hardswish(t(bf(y + idt))).numpy())
        return x

    def stats(dev, emu, gold):
        d = error_stats(dev, gold)
        e = error_stats(emu, gold)
        de = error_stats(dev, emu)
        return {"pcc_dev": d["pcc"], "pcc_emu": e["pcc"], "rel_l2_dev": d["rel_l2"], "rel_l2_emu": e["rel_l2"],
                "rel_l2_dev_vs_emu": de["rel_l2"], "max_abs_dev_vs_emu": de["max_abs"]}

    results = {"variant": variant, "frame": "pandaset_019_40", "configs": {}}
    stage_inputs = ["backbone.fusion_stem", "backbone.stage1", "backbone.stage2", "backbone.stage3"]
    emu_cache = {}
    with device_session() as device:
        for terms in args.terms:
            t0 = time.time()
            knobs = KNOBS.read(CONV_TERMS=terms)
            net = TtFRNet(device, weights, cfg, knobs=knobs)
            res = {}

            def up(nchw):
                return fmap_from_numpy(bf(nchw), device, dtype="bfloat16")

            def report(name, dev, emu, gold):
                if callable(emu):                   # the CPU emulation is the same for every configuration
                    if name not in emu_cache:
                        emu_cache[name] = emu()
                    emu = emu_cache[name]
                res[name] = stats(dev, emu, gold)
                r = res[name]
                print(f"[terms={terms}] {name:28s} pcc dev {r['pcc_dev']:.6f} emu {r['pcc_emu']:.6f} | rel_l2 dev "
                      f"{r['rel_l2_dev']:.3e} emu {r['rel_l2_emu']:.3e} dev-vs-emu {r['rel_l2_dev_vs_emu']:.3e}",
                      flush=True)

            # stem
            fr = g("backbone.frustum")
            fr32 = np.concatenate([fr, np.zeros_like(fr)], 1)
            x = up(fr32)
            for conv in net.stem:
                y = conv(x)
                free(x)
                x = y
            dev = fmap_to_numpy(x)
            free(x)
            def emu_stem():
                e = bf(fr)
                for c in weights.stem:
                    e = emu_conv(e, c)
                return e
            report("stem", dev, emu_stem, g("backbone.stem"))
            # stages
            for st_w, st, src in zip(weights.stages, net.stages, stage_inputs):
                l = st["index"]
                xin = g(src)
                # first conv alone (pre-activation)
                c1 = st["blocks"][0][0]
                x = up(xin)
                y = c1.conv(x)
                dev = fmap_to_numpy(y)
                free(y)
                lay = st_w.blocks[0][0]
                exact = F.conv2d(t(bf(xin)), t(lay.weight), t(lay.bias), stride=lay.stride,
                                 padding=lay.padding).numpy()
                report(f"layer{l}.block0.conv1(noact)", dev, bf(exact), exact)
                # blocks
                for c1_, c2_, ds_ in st["blocks"]:
                    y1 = c1_(x)
                    y = c2_(y1)
                    free(y1)
                    idt = ds_(x) if ds_ is not None else x
                    import ttnn

                    s = ttnn.add(y.tensor, idt.tensor, memory_config=ttnn.DRAM_MEMORY_CONFIG)
                    out = y.with_tensor(ttnn.hardswish(s))
                    free(s, y)
                    if ds_ is not None:
                        free(idt)
                    free(x)
                    x = out
                dev = fmap_to_numpy(x)
                free(x)
                report(f"layer{l}", dev, lambda: emu_blocks(bf(xin), st_w), g(f"backbone.layer{l}"))
                # attention
                pfz = g(f"backbone.pixel_fusion{l}")
                x = up(pfz)
                a0 = st["attention"][0](x)
                a1 = st["attention"][1](a0)
                free(a0)
                dev = fmap_to_numpy(a1)
                free(a1)
                report(f"attention{l}", dev,
                       lambda: emu_conv(emu_conv(bf(pfz), st_w.attention[0]), st_w.attention[1]),
                       g(f"backbone.attention{l}"))
                # addcmul
                lay_g, att_g = g(f"backbone.layer{l}"), g(f"backbone.attention{l}")
                xl, att = up(lay_g), up(att_g)
                out = ttnn.addcmul(xl.tensor, x.tensor, att.tensor, memory_config=ttnn.DRAM_MEMORY_CONFIG)
                dev = fmap_to_numpy(x.with_tensor(out))
                free(out, xl, att, x)
                e = bf(bf(lay_g) + bf(pfz) * bf(att_g))
                report(f"stage{l}(addcmul)", dev, e, g(f"backbone.stage{l}"))
            # fuse2
            f1 = g("decoder.fuse1")
            x = up(f1)
            y = net.fuse2(x)
            dev = fmap_to_numpy(y)
            free(x, y)
            report("fuse2", dev, lambda: emu_conv(bf(f1), weights.fuse_layer2), g("decoder.fuse2"))
            # element-wise ops on bf16 inputs (synthetic: the stage-1 map scaled to a wide range)
            import torch as _t
            import ttnn

            src = g("backbone.layer1")[:, :, :32, :] * 3.0
            xin = up(src)
            for opname, fn, ref in (("hardswish", ttnn.hardswish, F.hardswish), ("sigmoid", ttnn.sigmoid,
                                                                                 torch.sigmoid)):
                yo = xin.with_tensor(fn(xin.tensor))
                dev = fmap_to_numpy(yo)
                free(yo)
                exact = ref(t(bf(src))).numpy()
                report(f"op.{opname}(bf16)", dev, bf(exact), exact)
            free(xin)
            # softmax (fp32 logits, padded with -1e30)
            lg = g("head.logits")[:65536]
            pad = np.full((lg.shape[0], LOGIT_PAD), -1e30, np.float32)
            pad[:, : lg.shape[1]] = lg
            lt = ttnn.from_torch(_t.from_numpy(pad.reshape(1, 1, *pad.shape)), dtype=ttnn.float32,
                                 layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG)
            pr = ttnn.softmax(lt, dim=-1, compute_kernel_config=net._softmax_cfg, numeric_stable=True)
            dev = ttnn.to_torch(pr).float().numpy().reshape(-1, LOGIT_PAD)[:, : lg.shape[1]]
            free(lt, pr)
            exact = torch.softmax(t(lg), -1).numpy()
            report("softmax(fp32)", dev, exact, g("pred_probs")[:65536])
            net.release()
            res["seconds"] = round(time.time() - t0, 1)
            results["configs"][f"conv_terms={terms}"] = res
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(results, indent=1))
    print("wrote", args.out)
    return 0


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