# SPDX-License-Identifier: Apache-2.0
"""Bring-up check of ``tt.l1res`` (OPT round 4): the real encoder weights, the round-3 encoder vs the L1-resident one,
per bottleneck of layers 3-4, every intermediate (conv1, conv2, conv3, identity, output) compared bit for bit.
Both encoders get the same layer-2 output (the round-3 encoder's).

    bin/devrun -t 900 -- python -u code/scripts/debug_l1res.py
"""
from __future__ import annotations

import sys

import numpy as np


def main():
    import torch
    import ttnn

    from tt_bevdet.device import close_device, open_device
    from tt_bevdet.reference.config import default_config
    from tt_bevdet.reference.weights import BEVDetWeights
    from tt_bevdet.tests import paths
    from tt_bevdet.tt.conv_tune import apply_tune
    from tt_bevdet.tt.image_encoder import TtImageEncoder
    from tt_bevdet.tt.l1res import L1ResBlock, block_mc, ensure
    from tt_bevdet.tt.params import DEFAULT_PRECISION, BEVDetParams
    from tt_bevdet.ttaw.ops.conv import is_dram

    def bits(t):
        if t.memory_config().is_sharded():
            t = ttnn.to_memory_config(t, ttnn.DRAM_MEMORY_CONFIG)
        return ttnn.to_torch(ttnn.from_device(t)).to(torch.bfloat16).view(torch.int16).numpy().reshape(
            -1, int(t.shape[-1]))

    def cmp(tag, a, b):
        x, y = bits(a), bits(b)
        print(f"  {tag}: {int(np.count_nonzero(x != y))} differ; l1 mc "
              f"{'dram' if is_dram(b) else b.memory_config().shard_spec.shape if b.memory_config().is_sharded() else 'l1'}",
              flush=True)

    cfg = default_config()
    params = BEVDetParams(BEVDetWeights(paths.weights_onnx()), cfg)
    dev = open_device()
    try:
        kw = dict(precision=DEFAULT_PRECISION, stem_in="kin", conv_l1="all", backbone_res="next12", mm1x1=True)
        ref = TtImageEncoder(params, **kw)
        l1 = TtImageEncoder(params, l1res=True, **kw)
        apply_tune(ref.convs())
        apply_tune(l1.convs())
        n, (h, w) = cfg.n_cams, cfg.input_hw
        rng = np.random.default_rng(7)
        crop = rng.integers(0, 256, (n, h, w, 3), dtype=np.uint8)
        din = ttnn.from_torch(torch.from_numpy(crop.reshape(1, 1, n * h, w * 3)), dtype=ttnn.uint8,
                              layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        x = ref.kin_stem(din)
        for blk in ref.stages.stages[0] + ref.stages.stages[1]:
            x = blk(x)
        xr = xl = x
        gy = int(dev.compute_with_storage_grid_size().y)
        for s in (2, 3):
            for rb, lb in zip(ref.stages.stages[s], l1.stages.stages[s]):
                b = rb
                lbb: L1ResBlock = lb
                print(f"{b.name}: in ref {'dram' if is_dram(xr.tensor) else 'l1'} / l1 "
                      f"{'dram' if is_dram(xl.tensor) else 'l1'}", flush=True)
                y_r = b.conv1(xr)
                if is_dram(xl.tensor):
                    y_l = lbb.blk.conv1(xl)
                else:
                    y_l = lbb._mm(ensure(xl, block_mc(xl.rows, xl.channels, 8, gy))[0], lbb.w1,
                                  b.conv3.in_channels, True, gy)
                cmp("conv1", y_r.tensor, y_l.tensor)
                z_r = b.conv2(y_r)
                getattr(lbb.blk.conv2, "conv").slicing = "l1" if is_dram(y_l.tensor) else "auto"
                z_l = lbb.blk.conv2(y_l)
                cmp("conv2", z_r.tensor, z_l.tensor)
                zs, mv = ensure(z_l, block_mc(z_l.rows, z_l.channels, 8, gy))
                print("  conv2 moved:", mv, flush=True)
                y3_r = b.conv3(z_r)
                y3_l = lbb._mm(zs, lbb.w3, b.conv3.out_channels, False, gy)
                cmp("conv3", y3_r.tensor, y3_l.tensor)
                if b.downsample is not None:
                    d_r = b.downsample(xr)
                    xd = xl if is_dram(xl.tensor) else xl.with_tensor(
                        ttnn.to_memory_config(xl.tensor, ttnn.DRAM_MEMORY_CONFIG))    # as L1ResBlock does
                    d_l = lbb.blk.downsample(xd)
                    cmp("downsample", d_r.tensor, d_l.tensor)
                xr = b(xr)
                xl = lb(xl)
                cmp("out", xr.tensor, xl.tensor)
    finally:
        close_device(dev)
    return 0


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