#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""ST2 per-layer probe: every stacked [W_hi | W_lo] conv / transposed conv vs its two separate term convs, inside one
eager forward of the real graph on the shipped sample (EPI on), bit for bit per layer.

    bin/devrun -t 900 -- python -u code/scripts/diag/st2_probe.py

Each wrapped layer runs both forms on the same input; the stacked result continues the forward. Prints one line per
layer call (identical / number of differing elements / max abs diff) and the conv config of the layer.
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np


def main() -> int:
    import ttnn

    from tt_sceneseg.device import close_device, open_device
    from tt_sceneseg.host.preprocess import prepare_image
    from tt_sceneseg.reference import assets, config
    from tt_sceneseg.reference.weights import load_weights
    from tt_sceneseg.tt import layers as L
    from tt_sceneseg.tt.model import INPUT_SHAPE, TtSceneSeg
    from tt_sceneseg.ttaw.tensors import to_numpy

    code = Path(__file__).resolve().parents[2]
    wdir = assets.test_weights_dir()
    onnx = assets.find_in(wdir, config.ONNX_FILE) if wdir is not None else None
    if onnx is None:
        from huggingface_hub import snapshot_download
        from tt_sceneseg.api import SceneSeg

        d = snapshot_download(SceneSeg.DEFAULT_REPO, revision=SceneSeg.DEFAULT_REVISION,
                              allow_patterns=SceneSeg.ALLOW_PATTERNS)
        onnx = assets.find_in(Path(d), config.ONNX_FILE)
    dev = open_device(l1_small_size=98304)       # both forms of every layer prepared: more conv config tensors
    model = TtSceneSeg(load_weights(onnx), "bgr", epi=True, term_stack=True)
    stats = []

    def cmp(name, a, b):
        x, y = to_numpy(a), to_numpy(b)
        same = x.shape == y.shape and np.array_equal(x.view(np.uint8), y.view(np.uint8))
        nd = 0 if same else int((x != y).sum())
        md = 0.0 if same else float(np.abs(x.astype(np.float64) - y.astype(np.float64)).max())
        stats.append((name, same))
        print(f"{'SAME' if same else 'DIFF'} {name} {tuple(x.shape)} diff {nd} max {md:.3g}", flush=True)

    orig_conv = L.TermConv2d.__call__

    def conv_call(self, x, *, act_terms=1, output_dtype=None):
        if self.stacked is None:
            return orig_conv(self, x, act_terms=act_terms, output_dtype=output_dtype)
        xin = x.with_tensor(ttnn.clone(x.tensor))
        y1 = orig_conv(self, x, act_terms=act_terms, output_dtype=output_dtype)
        st, self.stacked = self.stacked, None
        y2 = orig_conv(self, xin, act_terms=act_terms, output_dtype=output_dtype)
        self.stacked = st
        cmp(f"conv {self.name} {x.height}x{x.width} {self.in_channels}->{self.out_channels} at={act_terms}",
            y1.tensor, y2.tensor)
        ttnn.deallocate(y2.tensor)
        ttnn.deallocate(xin.tensor)
        return y1

    orig_views = L.TermConvTranspose2d.term_views

    def views(self, x):
        if self.stacked is None:
            return orig_views(self, x)
        parts, cols, first = orig_views(self, x)
        st, self.stacked = self.stacked, None
        p2, _, _ = orig_views(self, x)
        self.stacked = st
        c = self.out_channels
        shp = [int(v) for v in parts[0].shape]
        for i, off in enumerate((0, c)):
            h = ttnn.slice(parts[0], [0, 0, 0, off], [shp[0], shp[1], shp[2], off + c])
            cmp(f"deconv {self.name} term {i} {x.height}x{x.width} {self.in_channels}->{c}", h, p2[i])
            ttnn.deallocate(h)
        for t in p2:
            ttnn.deallocate(t)
        return parts, cols, first

    L.TermConv2d.__call__ = conv_call
    L.TermConvTranspose2d.term_views = views
    canvas = prepare_image(str(code / "tt_sceneseg" / "samples" / "highway_normal_1.png")).device_input()
    import torch

    image = ttnn.from_torch(torch.from_numpy(np.ascontiguousarray(canvas.reshape(INPUT_SHAPE))), dtype=ttnn.uint8,
                            layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
    out = model.forward(image)
    print("class map", np.bincount(to_numpy(out["class_map"]).astype(np.int64).ravel(), minlength=3).tolist())
    n_diff = sum(not s for _, s in stats)
    print(f"{len(stats)} layer calls, {n_diff} differ")
    model.release()
    close_device(dev)
    return 0


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