# SPDX-License-Identifier: Apache-2.0
"""Weights of the three StreamPETR graphs (AutowareFoundation/camera_streampetr@v1.0), read the way the TT port
reads them.

Every file is parsed as data with the vendored ``ttaw.weights.OnnxWeights`` (C04): nothing is executed, and every
parameter is addressed through the graph node that consumes it (``/img_backbone/stage4/OSA4_2/concat/.../Conv``,
``/transformer/decoder/layers.3/attentions.1/attn/MatMul_1``, ...), so attributes (strides, pads) come from the
deployed graph itself. The exporter folded every BatchNorm into its convolution (S:4.1: "BN folded into conv weight +
bias"), so there is **no BN fold in the port**: conv weights and biases are used as stored. Linear layers are
``MatMul(x, W[in, out]) + Add(b)`` pairs (``linear``), the attention output projections ``Gemm`` with ``transB=1``
(``gemm`` returns ``W[in, out]`` like ``linear``). The head's six cls / reg branch and post-norm instances share their
initializers (the branch weights of the training code are shared by all decoder layers, S:4.3 step 8), which
:meth:`StreamPETRWeights.check_graph` verifies.

:meth:`StreamPETRWeights.check_graph` verifies that the files are the networks this port was written for (I/O names
and shapes, op histograms, the 118 convolutions mapped 1:1 to modules, MaxPool ceil mode, nearest Resize, LayerNorm
eps, the attention scale, TopK k, pc_range / position_range constants) and fails loudly otherwise.
"""
from __future__ import annotations

import collections
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, Tuple

import numpy as np

from .config import ONNX_FILES, StreamPETRConfig, default_config

__all__ = ["Conv", "StreamPETRWeights", "backbone_modules", "find_weights_dir", "OSA_STAGES"]

# VoVNet-99-eSE (vovnet.py:80-88): stage -> (OSA blocks, stage channels, concat output channels)
OSA_STAGES: Dict[int, Tuple[int, int, int]] = {2: (1, 128, 256), 3: (3, 160, 512), 4: (9, 192, 768),
                                               5: (3, 224, 1024)}

# op histograms of the deployed graphs (S:4.1-4.3; verify_onnx.py)
_HIST = {
    "backbone": {"Conv": 118, "Relu": 99, "Add": 29, "Concat": 16, "GlobalAveragePool": 16, "Clip": 16, "Div": 16,
                 "Mul": 16, "MaxPool": 3, "Squeeze": 2, "Resize": 1, "Reshape": 1},
    "position_embedding": {"Reshape": 11, "Unsqueeze": 6, "Gather": 6, "Slice": 6, "Div": 5, "Concat": 4, "Tile": 4,
                           "ScatterND": 4, "Squeeze": 3, "Expand": 3, "Mul": 3, "MatMul": 3, "Clip": 3, "Sub": 2,
                           "Add": 2, "Split": 1, "Abs": 1, "Max": 1, "Log": 1, "Relu": 1},
    "head": {"Add": 148, "MatMul": 128, "Reshape": 67, "Transpose": 58, "LayerNormalization": 40, "Slice": 37,
             "Unsqueeze": 36, "Relu": 36, "Concat": 27, "Mul": 24, "ScatterND": 19, "Gather": 14, "Softmax": 12,
             "Gemm": 12, "Sin": 10, "Cos": 10, "Sigmoid": 8, "Div": 6, "Tile": 5, "GatherElements": 5, "Clip": 3,
             "Squeeze": 2, "Sub": 2, "TopK": 2, "Flatten": 2, "Log": 1, "Cast": 1},
}
_IO = {
    "backbone": ((("img", (1, 5, 3, 480, 640)),), (("img_feats", (1, 5, 256, 30, 40)),)),
    "position_embedding": ((("img_metas_pad", (3,)), ("intrinsics", (1, 5, 4, 4)), ("img2lidar", (1, 5, 4, 4))),
                           (("pos_embed", (1, 6000, 256)), ("cone", (1, 6000, 8)))),
    "head": ((("x", (1, 5, 256, 30, 40)), ("pos_embed", (1, 6000, 256)), ("cone", (1, 6000, 8)),
              ("data_ego_pose", (1, 4, 4)), ("data_ego_pose_inv", (1, 4, 4)), ("pre_memory_embedding", (1, 1024, 256)),
              ("pre_memory_reference_point", (1, 1024, 3)), ("pre_memory_timestamp", (1, 1024, 1)),
              ("pre_memory_egopose", (1, 1024, 4, 4)), ("pre_memory_velo", (1, 1024, 2))),
             (("all_cls_scores", (1, 5, 5400)), ("all_bbox_preds", (1, 10, 5400)),
              ("post_memory_embedding", (1, 1280, 256)), ("post_memory_reference_point", (1, 1280, 3)),
              ("post_memory_timestamp", (1, 1280, 1)), ("post_memory_egopose", (1, 1280, 4, 4)),
              ("post_memory_velo", (1, 1280, 2)), ("reference_points", (1, 900, 3)), ("tgt", (1, 900, 256)),
              ("temp_memory", (1, 768, 256)), ("temp_pos", (1, 768, 256)), ("query_pos", (1, 900, 256)),
              ("query_pos_in", (1, 644, 256)), ("outs_dec", (6, 1, 900, 256)))),
}


def backbone_modules() -> Dict[str, List[str]]:
    """Conv node names of the backbone graph by stage, in graph order (118 = 3 + 7 + 21 + 63 + 21 + 3)."""
    out: Dict[str, List[str]] = {"stem": [f"/img_backbone/stem/stem_{i}/conv/Conv" for i in (1, 2, 3)]}
    for st, (nb, _, _) in OSA_STAGES.items():
        names = []
        for k in range(1, nb + 1):
            pre = f"/img_backbone/stage{st}/OSA{st}_{k}"
            names += [f"{pre}/layers.{i}/OSA{st}_{k}_{i}/conv/Conv" for i in range(5)]
            names += [f"{pre}/concat/OSA{st}_{k}_concat/conv/Conv", f"{pre}/ese/fc/Conv"]
        out[f"stage{st}"] = names
    out["neck"] = ["/img_neck/lateral_convs.0/conv/Conv", "/img_neck/lateral_convs.1/conv/Conv",
                   "/img_neck/fpn_convs.0/conv/Conv"]
    return out


@dataclass(frozen=True)
class Conv:
    """One convolution as the deployed graph runs it (float32 weights as stored; BN already folded by the exporter)."""

    node: str
    weight: np.ndarray            # [Cout, Cin, kh, kw]
    bias: Optional[np.ndarray]    # [Cout]
    stride: int
    padding: int


class StreamPETRWeights:
    """Parameters of the three deployed graphs, addressed by consuming node. ``weights_dir`` holds the HF files."""

    def __init__(self, weights_dir: Any, *, check: bool = True, verify_sha256: bool = False):
        from ..ttaw.weights import OnnxWeights

        self.dir = Path(weights_dir)
        self.onnx: Dict[str, Any] = {}
        for key, (name, sha) in ONNX_FILES.items():
            p = self.dir / name
            if not p.is_file():
                raise FileNotFoundError(f"{p} (the HF camera_streampetr v1.0 files are needed)")
            self.onnx[key] = OnnxWeights(p)
            if verify_sha256 and self.onnx[key].sha256 != sha:
                raise ValueError(f"{name}: sha256 {self.onnx[key].sha256} != pinned {sha}")
        self._convs: Dict[str, Conv] = {}
        if check:
            self.check_graph()

    @property
    def backbone(self):
        return self.onnx["backbone"]

    @property
    def pe(self):
        return self.onnx["position_embedding"]

    @property
    def head(self):
        return self.onnx["head"]

    def sha256(self) -> Dict[str, str]:
        return {ONNX_FILES[k][0]: w.sha256 for k, w in self.onnx.items()}

    # ---- addressing ---------------------------------------------------------------------------------------
    def conv(self, node: str) -> Conv:
        """A backbone convolution by node name (symmetric pads, group 1, dilation 1: checked)."""
        if node not in self._convs:
            p = self.backbone.conv(node)
            pads = tuple(p.pads)
            if len(set(pads)) != 1 or len(set(p.strides)) != 1 or set(p.dilations) != {1} or p.group != 1 \
                    or p.auto_pad not in ("NOTSET", ""):
                raise ValueError(f"{node}: unsupported attributes {p}")
            self._convs[node] = Conv(node, np.asarray(p.weight, np.float32),
                                     None if p.bias is None else np.asarray(p.bias, np.float32),
                                     int(p.strides[0]), int(pads[0]))
        return self._convs[node]

    def iter_convs(self) -> Iterator[Tuple[str, Conv]]:
        for names in backbone_modules().values():
            for n in names:
                yield n, self.conv(n)

    @staticmethod
    def _const_operand(w: Any, node: str) -> np.ndarray:
        ps = w.params(node)
        if len(ps) != 1:
            raise ValueError(f"{node}: expected one constant operand, got slots {sorted(ps)}")
        return next(iter(ps.values()))

    def linear(self, prefix: str, graph: str = "head") -> Tuple[np.ndarray, np.ndarray]:
        """``<prefix>/MatMul`` + ``<prefix>/Add`` -> (W [in, out], b [out]) float32."""
        w = self.onnx[graph]
        W = np.asarray(w.matmul_weight(prefix + "/MatMul"), np.float32)
        b = np.asarray(self._const_operand(w, prefix + "/Add"), np.float32)
        if W.ndim != 2 or b.shape != (W.shape[1],):
            raise ValueError(f"{prefix}: weight {W.shape} / bias {b.shape}")
        return W, b

    def attn_proj(self, prefix: str, which: str) -> Tuple[np.ndarray, np.ndarray]:
        """Q / K / V projection of ``<prefix>/attn`` (``MatMul`` / ``MatMul_1`` / ``MatMul_2`` + ``Add*``)."""
        suffix = {"q": "", "k": "_1", "v": "_2"}[which]
        w = self.head
        W = np.asarray(w.matmul_weight(f"{prefix}/attn/MatMul{suffix}"), np.float32)
        b = np.asarray(self._const_operand(w, f"{prefix}/attn/Add{suffix}"), np.float32)
        return W, b

    def gemm(self, node: str) -> Tuple[np.ndarray, np.ndarray]:
        """A ``Gemm`` (``transB=1``, alpha = beta = 1) -> (W [in, out], b [out]) float32."""
        g = self.head.gemm(node)
        if g.trans_a or not g.trans_b or g.alpha != 1.0 or g.beta != 1.0:
            raise ValueError(f"{node}: unexpected Gemm attributes {g}")
        return np.ascontiguousarray(np.asarray(g.weight, np.float32).T), np.asarray(g.bias, np.float32)

    def layer_norm(self, node: str) -> Tuple[np.ndarray, np.ndarray, float]:
        """``LayerNormalization`` -> (gamma, beta, eps) (axis -1 checked)."""
        n = self.head.node(node)
        if int(n.attrs.get("axis", -1)) != -1:
            raise ValueError(f"{node}: axis {n.attrs.get('axis')}")
        ps = self.head.params(node)
        return np.asarray(ps[1], np.float32), np.asarray(ps[2], np.float32), float(n.attrs.get("epsilon", 1e-5))

    def const(self, name: str, graph: str = "head") -> np.ndarray:
        return np.asarray(self.onnx[graph].array(name), np.float32)

    def learned_queries(self) -> Dict[str, np.ndarray]:
        """The constant-folded tensors of the 644 learned queries (S:4.3 "Folded constants"):
        ``reference_points`` [1, 644, 3] (``/Tile_output_0``), ``tgt`` [1, 644, 256] = MLN_mem(0, identity motion)
        (``/ego_pose_memory/Add_output_0``), ``query_pos`` [1, 644, 256] (``/Add_output_0``) and the debug output
        ``query_pos_in``."""
        return {"reference_points": self.const("/Tile_output_0"), "tgt": self.const("/ego_pose_memory/Add_output_0"),
                "query_pos": self.const("/Add_output_0"), "query_pos_in": self.const("query_pos_in")}

    def attention_scale(self) -> np.float32:
        """The ONNX ``/attn/Mul`` constant applied to Q after the projection (1/sqrt(32) in float32)."""
        vals = {float(self._const_operand(self.head, n.name)) for n in self.head.nodes("Mul", "*/attn/Mul")}
        if len(vals) != 1:
            raise ValueError(f"attention scales differ: {vals}")
        return np.float32(vals.pop())

    # ---- graph checks --------------------------------------------------------------------------------------
    def check_graph(self, cfg: Optional[StreamPETRConfig] = None) -> None:
        """The files are the deployed StreamPETR this port reproduces (module docstring); raises ``ValueError``."""
        cfg = cfg or default_config()
        problems: List[str] = []
        for key, w in self.onnx.items():
            hist = dict(collections.Counter(n.op_type for n in w.nodes()))
            if hist != _HIST[key]:
                problems.append(f"{key}: op histogram {hist} != {_HIST[key]}")
            ins, outs = _IO[key]
            if tuple(w.input_names) != tuple(n for n, _ in ins) or tuple(w.output_names) != tuple(n for n, _ in outs):
                problems.append(f"{key}: I/O names {w.input_names} -> {w.output_names}")
            if set(w.opset) != {"ai.onnx"}:
                problems.append(f"{key}: custom op domains {w.opset}")
        # backbone: every Conv maps 1:1 to a module; pools, resize, eSE hard-sigmoid
        convs = [n for names in backbone_modules().values() for n in names]
        if sorted(convs) != sorted(n.name for n in self.backbone.nodes("Conv")):
            problems.append("backbone: Conv nodes do not map 1:1 to the VoVNet-99-eSE + CPFPN modules")
        for n in self.backbone.nodes("MaxPool"):
            a = n.attrs
            if (int(a.get("ceil_mode", 0)), tuple(a["kernel_shape"]), tuple(a["strides"]), tuple(a.get("pads", ()))) \
                    != (1, (3, 3), (2, 2), (0, 0, 0, 0)):
                problems.append(f"backbone: {n.name} {a}")
        r = self.backbone.find_node("*", "Resize").attrs
        if (r.get("mode"), r.get("coordinate_transformation_mode"), r.get("nearest_mode")) != \
                ("nearest", "asymmetric", "floor"):
            problems.append(f"backbone: Resize {r}")
        clips = {tuple(float(v) for v in self.backbone.params(n.name).values()) for n in self.backbone.nodes("Clip")}
        if clips != {(0.0, 6.0)}:
            problems.append(f"backbone: eSE Clip bounds {clips}")
        divs = {float(next(iter(self.backbone.params(n.name).values()))) for n in self.backbone.nodes("Div")}
        if divs != {6.0}:
            problems.append(f"backbone: eSE Div constants {divs}")
        # position embedding: the position_range constants
        pr = cfg.position_range
        if not np.allclose(self.pe.params("/Sub")[1], np.float32(pr[:3]), rtol=0, atol=0) or not np.allclose(
                self.pe.params("/Div_3")[1], np.float32(np.subtract(pr[3:], pr[:3])), rtol=0, atol=0):
            problems.append("position_embedding: position_range constants differ from the configuration")
        # head: LN eps, attention scale, TopK k, pc_range, shared branches
        eps = {float(n.attrs.get("epsilon", 1e-5)) for n in self.head.nodes("LayerNormalization")}
        if eps != {float(np.float32(cfg.ln_eps))}:
            problems.append(f"head: LayerNorm eps {eps}")
        if self.attention_scale() != np.float32(1.0 / np.sqrt(cfg.head_dim)):
            problems.append(f"head: attention scale {self.attention_scale()}")
        ks = sorted(int(self.head.params(n.name)[1][0]) for n in self.head.nodes("TopK"))
        if ks != [1, cfg.num_propagated]:
            problems.append(f"head: TopK k {ks}")
        for n in self.head.nodes("TopK"):
            if int(n.attrs.get("largest", 1)) != 1 or int(n.attrs.get("sorted", 1)) != 1:
                problems.append(f"head: {n.name} {n.attrs}")
        pc = cfg.pc_range
        if not np.array_equal(self.head.params("/Sub")[1], np.float32(pc[:3])) or not np.array_equal(
                self.head.params("/Div_3")[1], np.float32(np.subtract(pc[3:], pc[:3]))):
            problems.append("head: pc_range constants (temporal alignment) differ from the configuration")
        if not np.array_equal(self.head.params("/Mul_69")[1], np.float32(np.subtract(pc[3:], pc[:3]))) or \
                not np.array_equal(self.head.params("/Add_8")[1], np.float32(pc[:3])):
            problems.append("head: pc_range constants (box decode) differ from the configuration")
        for branch in ("/cls_branches.0/cls_branches.0.0", "/reg_branches.0/reg_branches.0.0"):
            inits = {self.head.node(f"{branch}{'' if i == 0 else f'_{i}'}/MatMul").inputs[1]
                     for i in range(cfg.num_layers)}
            if len(inits) != 1:
                problems.append(f"head: {branch} is not shared by the {cfg.num_layers} layers")
        posts = {self.head.node(f"/transformer/decoder/post_norm{'' if i == 0 else f'_{i}'}/LayerNormalization")
                 .inputs[1:] for i in range(cfg.num_layers)}
        if len(posts) != 1:
            problems.append("head: the post_norm instances do not share their initializers")
        q = self.learned_queries()
        if q["reference_points"].shape != (1, cfg.num_query, 3) or q["tgt"].shape != (1, cfg.num_query, 256):
            problems.append("head: learned-query constants have unexpected shapes")
        if problems:
            raise ValueError("the ONNX files are not the StreamPETR v1.0 graphs this port reproduces:\n  "
                             + "\n  ".join(problems))


def find_weights_dir(weights_dir: Any = None) -> Optional[Path]:
    """``weights_dir`` if it holds the three ONNX files (``weights_dir`` may also name one of them), else None."""
    if weights_dir is None:
        return None
    p = Path(weights_dir)
    if p.is_file():
        p = p.parent
    return p if all((p / name).is_file() for name, _ in ONNX_FILES.values()) else None