# SPDX-License-Identifier: Apache-2.0
"""Weights of SceneSeg, read as DATA the way the TT port reads them.

:func:`load_weights` takes the deployed ``SceneSeg_FP32.onnx``: the vendored ``ttaw.weights.OnnxWeights`` addresses
every initializer by its consuming node (Conv / ConvTranspose / Gemm), and BatchNormalization is folded (fp64, one
rounding) exactly where a BN node follows a conv. The deployed graph has no BN node (PyTorch folded them at export,
eps 1e-5, S:sceneseg:249,374), so it folds nothing and returns the ONNX tensors unchanged: the numbers Autoware's
TensorRT engine was built from. The TT port calls the same function.

:func:`load_state_dict_weights` builds the same table from the upstream ``.pth`` (or its safetensors copy re-hosted
in ``weights/``, PLAN D10): it folds every torchvision ``Conv2dNormActivation`` BatchNorm with
``ttaw.weights.fold_bn_conv`` (fp64, one rounding). It is the independent cross-check of the ONNX path (tests:
equal to ~1e-6, S:sceneseg 6.2) and documents the BN fold the export did.

:data:`ARCH` is the expected layer list (100 convs, 5 transposed convs, 3 linears), written from the architecture
(torchvision ``efficientnet_b0().features`` with its stage table, SceneContext, SceneNeck, SceneSegHead;
S:sceneseg 4.2) rather than read from the graph: the loaders check every layer against it (kernel, stride, padding,
groups, channels, activation) and that nothing else is in the graph. Keys are the torch module paths of the conv /
linear (the ``.pth`` prefixes), e.g. ``Backbone.encoder.2.1.block.1.0`` (the depthwise conv of stage 2, block 1).
Importing this module needs neither onnx nor torch.
"""
from __future__ import annotations

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

import numpy as np

from . import config

__all__ = ["LayerSpec", "LayerWeights", "ARCH", "ARCH_BY_KEY", "WeightsError", "load_weights",
           "load_state_dict_weights", "mbconv_specs", "param_count", "iter_specs"]

ACTIVATIONS = ("silu", "gelu", "sigmoid", "none")


class WeightsError(ValueError):
    """The file does not hold the network this port reproduces."""


@dataclass(frozen=True)
class LayerSpec:
    """One weighted layer of the deployed graph.

    ``kind``: ``"conv"`` (weight ``[cout, cin / groups, k, k]``, padding ``k // 2``), ``"conv_transpose"`` (weight
    ``[cin, cout, k, k]``, stride ``k``, no padding) or ``"linear"`` (ONNX Gemm ``transB=1``: weight ``[cout, cin]``).
    ``act``: the activation that follows in the graph (``silu`` = x * sigmoid(x), ``gelu`` = the exact erf form,
    ``sigmoid``, ``none``). ``bn``: the ``.pth`` prefix of the BatchNorm folded into the conv (torchvision
    ``Conv2dNormActivation``), None if the layer has a bias of its own."""

    key: str
    node: str
    kind: str
    cin: int
    cout: int
    k: int = 1
    stride: int = 1
    groups: int = 1
    act: str = "none"
    bn: Optional[str] = None

    @property
    def pad(self) -> int:
        return self.k // 2 if self.kind == "conv" else 0

    @property
    def weight_shape(self) -> Tuple[int, ...]:
        if self.kind == "conv":
            return (self.cout, self.cin // self.groups, self.k, self.k)
        if self.kind == "conv_transpose":
            return (self.cin, self.cout, self.k, self.k)
        return (self.cout, self.cin)

    @property
    def depthwise(self) -> bool:
        return self.kind == "conv" and self.groups > 1 and self.groups == self.cin == self.cout


@dataclass(frozen=True)
class LayerWeights:
    """Float32 weight and bias of one layer, with its spec and source (ONNX node or ``.pth`` prefix)."""

    spec: LayerSpec
    weight: np.ndarray
    bias: np.ndarray
    source: str
    folded_bn: Optional[str] = None


# ----------------------------------------------------------------------------------------------- architecture

def _conv(key: str, node: str, cin: int, cout: int, k: int = 1, stride: int = 1, groups: int = 1, act: str = "none",
          bn: Optional[str] = None) -> LayerSpec:
    return LayerSpec(key, node, "conv", cin, cout, k, stride, groups, act, bn)


def mbconv_specs(stage: int, block: int, expand: int, k: int, stride: int, cin: int, cout: int) -> List[LayerSpec]:
    """torchvision ``MBConv`` (``efficientnet.py``): [expand 1x1 + BN + SiLU] (absent when expand == 1), depthwise kxk
    + BN + SiLU, SqueezeExcitation (fc1 1x1 + SiLU, fc2 1x1 + Sigmoid, squeeze = max(1, cin // 4) of the BLOCK input),
    project 1x1 + BN (no activation). Residual (StochasticDepth = identity at eval) when stride 1 and cin == cout."""
    cexp = cin * expand
    key = f"Backbone.encoder.{stage}.{block}.block"
    node = f"/Backbone/encoder.{stage}/encoder.{stage}.{block}/block"
    out: List[LayerSpec] = []
    i = 0
    if expand != 1:
        out.append(_conv(f"{key}.0.0", f"{node}/block.0/block.0.0/Conv", cin, cexp, 1, act="silu", bn=f"{key}.0.1"))
        i = 1
    out.append(_conv(f"{key}.{i}.0", f"{node}/block.{i}/block.{i}.0/Conv", cexp, cexp, k, stride, cexp, "silu",
                     f"{key}.{i}.1"))
    csq = max(1, cin // 4)
    out.append(_conv(f"{key}.{i + 1}.fc1", f"{node}/block.{i + 1}/fc1/Conv", cexp, csq, 1, act="silu"))
    out.append(_conv(f"{key}.{i + 1}.fc2", f"{node}/block.{i + 1}/fc2/Conv", csq, cexp, 1, act="sigmoid"))
    out.append(_conv(f"{key}.{i + 2}.0", f"{node}/block.{i + 2}/block.{i + 2}.0/Conv", cexp, cout, 1,
                     bn=f"{key}.{i + 2}.1"))
    return out


def _build_arch() -> List[LayerSpec]:
    a: List[LayerSpec] = [_conv("Backbone.encoder.0.0", "/Backbone/encoder.0/encoder.0.0/Conv", 3,
                                config.STEM_CHANNELS, 3, 2, act="silu", bn="Backbone.encoder.0.1")]
    for s, (e, k, stride, cin, cout, n) in enumerate(config.B0_SETTINGS, start=1):
        for b in range(n):
            a += mbconv_specs(s, b, e, k, stride if b == 0 else 1, cin if b == 0 else cout, cout)
    a.append(_conv("Backbone.encoder.8.0", "/Backbone/encoder.8/encoder.8.0/Conv", config.B0_SETTINGS[-1][4],
                   config.HEAD_CHANNELS, 1, act="silu", bn="Backbone.encoder.8.1"))
    mlp = config.CONTEXT_MLP
    for i in range(3):
        a.append(LayerSpec(f"SceneContext.context_layer_{i}", f"/SceneContext/context_layer_{i}/Gemm", "linear",
                           mlp[i], mlp[i + 1], act="gelu" if i < 2 else "sigmoid"))
    ch = config.CONTEXT_CONVS
    for i in range(4):
        a.append(_conv(f"SceneContext.context_layer_{i + 3}", f"/SceneContext/context_layer_{i + 3}/Conv", ch[i],
                       ch[i + 1], 3, act="gelu"))
    for mod, blocks in (("SceneNeck", config.NECK_BLOCKS), ("SceneSegHead", config.HEAD_BLOCKS)):
        for up, cup, skip, cskip, decodes in blocks:
            a.append(LayerSpec(f"{mod}.{up}", f"/{mod}/{up}/ConvTranspose", "conv_transpose", cup, cup, 2, 2))
            if skip is not None:
                a.append(_conv(f"{mod}.{skip}", f"/{mod}/{skip}/Conv", cskip, cup, 1))
            for name, ci, co in decodes:
                a.append(_conv(f"{mod}.{name}", f"/{mod}/{name}/Conv", ci, co, 3, act="gelu"))
    name, ci, co = config.FINAL_CONV
    a.append(_conv(f"SceneSegHead.{name}", f"/SceneSegHead/{name}/Conv", ci, co, 3))
    return a


ARCH: Tuple[LayerSpec, ...] = tuple(_build_arch())
ARCH_BY_KEY: Dict[str, LayerSpec] = {s.key: s for s in ARCH}
assert len(ARCH_BY_KEY) == len(ARCH) == 108
assert sum(s.kind == "conv" for s in ARCH) == 100 and sum(s.kind == "conv_transpose" for s in ARCH) == 5
assert sum(s.depthwise for s in ARCH) == 16 and sum(s.bn is not None for s in ARCH) == 1 + 16 + 15 + 16 + 1


def iter_specs(prefix: str = "") -> Iterator[LayerSpec]:
    """The :data:`ARCH` entries under a key prefix, in graph order."""
    return (s for s in ARCH if s.key.startswith(prefix))


def param_count(weights: Dict[str, LayerWeights]) -> int:
    return int(sum(lw.weight.size + lw.bias.size for lw in weights.values()))


# ------------------------------------------------------------------------------------------------- ONNX path

def _activation(w, node) -> str:
    """The activation pattern that consumes ``node``'s output in the exported graph: SiLU = Sigmoid + Mul, exact
    GELU = Div (by sqrt 2) + Mul, Sigmoid alone, else none (Add / Conv / graph output)."""
    ops = sorted(c.op_type for c in w.consumers(node.outputs[0]))
    if ops == ["Mul", "Sigmoid"]:
        return "silu"
    if ops == ["Div", "Mul"]:
        return "gelu"
    if ops == ["Sigmoid"]:
        return "sigmoid"
    return "none"


def load_weights(onnx_path: Union[str, Path], *, check_sha256: bool = True) -> Dict[str, LayerWeights]:
    """Every weighted layer of ``SceneSeg_FP32.onnx`` as float32 :class:`LayerWeights`, keyed like :data:`ARCH`
    and checked against it. Raises :class:`WeightsError` if the file is not the expected network."""
    from ..ttaw.weights import OnnxWeights

    w = OnnxWeights(onnx_path)
    if check_sha256 and w.sha256 != config.ONNX_SHA256:
        raise WeightsError(f"{onnx_path}: sha256 {w.sha256} is not {config.ONNX_FILE} ({config.ONNX_SHA256}); pass "
                           "check_sha256=False for another export of the same graph")
    weighted = {n.name: n for n in w.nodes() if n.op_type in ("Conv", "ConvTranspose", "Gemm")}
    expected = {s.node for s in ARCH}
    missing = sorted(expected - set(weighted))
    extra = sorted(set(weighted) - expected)
    if missing or extra:
        raise WeightsError(f"layer layout differs from ARCH: missing {missing[:4]}, unexpected {extra[:4]}")
    out: Dict[str, LayerWeights] = {}
    for spec in ARCH:
        node = weighted[spec.node]
        folded = None
        if spec.kind == "linear":
            g = w.gemm(node.name)
            if g.trans_a or not g.trans_b or g.alpha != 1.0 or g.beta != 1.0:
                raise WeightsError(f"{node.name}: Gemm attributes {g}")
            weight, bias = g.weight, g.bias
        else:
            p = w.conv(node.name)
            if spec.kind == "conv":
                expect = dict(kernel_shape=(spec.k, spec.k), strides=(spec.stride, spec.stride), pads=(spec.pad,) * 4,
                              dilations=(1, 1), group=spec.groups)
            else:
                expect = dict(kernel_shape=(spec.k, spec.k), strides=(spec.stride, spec.stride), pads=(0,) * 4,
                              dilations=(1, 1), group=1)
            got = dict(kernel_shape=tuple(p.kernel_shape), strides=tuple(p.strides), pads=tuple(p.pads or (0,) * 4),
                       dilations=tuple(p.dilations or (1, 1)), group=p.group)
            bad_pad = p.auto_pad not in ("", "NOTSET", None) or tuple(p.output_padding or ()) not in ((), (0, 0))
            if got != expect or bad_pad:
                raise WeightsError(f"{node.name}: attributes {got} (auto_pad {p.auto_pad}), expected {expect}")
            weight, bias = p.weight, p.bias
            bn = [c for c in w.consumers(node.outputs[0]) if c.op_type == "BatchNormalization"]
            if bn:                                     # not in the deployed export; folded like the TT port would
                weight, bias = w.fold_conv_bn(node.name, bn[0].name)
                folded = bn[0].name
                node = bn[0]
        if tuple(weight.shape) != spec.weight_shape:
            raise WeightsError(f"{spec.node}: weight {tuple(weight.shape)}, expected {spec.weight_shape}")
        if bias is None:
            bias = np.zeros(spec.cout, np.float32)
        if tuple(np.shape(bias)) != (spec.cout,):
            raise WeightsError(f"{spec.node}: bias {np.shape(bias)}, expected ({spec.cout},)")
        act = _activation(w, node)
        if act != spec.act:
            raise WeightsError(f"{spec.node}: activation {act}, ARCH says {spec.act}")
        out[spec.key] = LayerWeights(spec, np.ascontiguousarray(weight, np.float32),
                                     np.ascontiguousarray(bias, np.float32), spec.node, folded)
    return out


# ------------------------------------------------------------------------------------- .pth / safetensors path

def load_state_dict_weights(path: Union[str, Path], *, eps: float = config.BN_EPS) -> Dict[str, LayerWeights]:
    """The same table from the upstream state_dict (``.pth`` via ``torch.load(weights_only=True)`` or the re-hosted
    ``.safetensors``): BatchNorms folded in fp64 with one rounding (``ttaw.weights.fold_bn_conv``). Raises
    :class:`WeightsError` on a missing or unexpected tensor (``num_batches_tracked`` is ignored)."""
    from ..ttaw.weights import fold_bn_conv, load_state_dict

    sd = load_state_dict(path)
    used = set()
    out: Dict[str, LayerWeights] = {}
    for spec in ARCH:
        try:
            weight = np.asarray(sd[f"{spec.key}.weight"])
            used.add(f"{spec.key}.weight")
            if spec.bn is None:
                bias = np.asarray(sd[f"{spec.key}.bias"])
                used.add(f"{spec.key}.bias")
                folded = None
            else:
                bn = sd.bn(spec.bn, eps)
                used.update(f"{spec.bn}.{n}" for n in ("weight", "bias", "running_mean", "running_var"))
                weight, bias = fold_bn_conv(weight, None, bn)
                folded = spec.bn
        except KeyError as e:
            raise WeightsError(f"{path}: {e}") from None
        if tuple(weight.shape) != spec.weight_shape:
            raise WeightsError(f"{spec.key}: weight {tuple(weight.shape)}, expected {spec.weight_shape}")
        out[spec.key] = LayerWeights(spec, np.ascontiguousarray(weight, np.float32),
                                     np.ascontiguousarray(bias, np.float32), spec.key, folded)
    extra = sorted(k for k in sd if k not in used and not k.endswith("num_batches_tracked"))
    if extra:
        raise WeightsError(f"{path}: tensors not in ARCH: {extra[:4]}")
    return out