"""Compact CNN-encoder + Lite-DPT depth model.

The high-resolution stages are pure convolutional; attention, if present in the
encoder, is confined to the stride-32 bottleneck.  The decoder is a lightweight
top-down FPN with depthwise-separable residual blocks and an ASPP bottleneck.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import torch
import torch.nn.functional as F
from torch import nn

try:
    import timm
except ImportError as exc:  # pragma: no cover
    raise ImportError("install timm: uv add timm") from exc


@dataclass
class CompactDepthOutput:
    predicted_depth: torch.Tensor
    aux: dict[str, torch.Tensor] | None = None


class ConvNormAct(nn.Sequential):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        *,
        stride: int = 1,
        groups: int = 1,
        dilation: int = 1,
        activation: bool = True,
    ) -> None:
        padding = dilation * (kernel_size // 2)
        layers: list[nn.Module] = [
            nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size,
                stride=stride,
                padding=padding,
                groups=groups,
                dilation=dilation,
                bias=False,
            ),
            nn.BatchNorm2d(out_channels),
        ]
        if activation:
            layers.append(nn.ReLU(inplace=True))
        super().__init__(*layers)


class DWResidualBlock(nn.Module):
    """Inverted-bottleneck depthwise block: 1x1 expand, 3x3 depthwise, 1x1 project."""

    def __init__(self, channels: int, expansion: int = 2) -> None:
        super().__init__()
        hidden = channels * expansion
        self.expand = ConvNormAct(channels, hidden, 1)
        self.depthwise = ConvNormAct(hidden, hidden, 3, groups=hidden)
        self.project = nn.Sequential(
            nn.Conv2d(hidden, channels, 1, bias=False),
            nn.BatchNorm2d(channels),
        )
        self.act = nn.ReLU(inplace=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.act(x + self.project(self.depthwise(self.expand(x))))


class ASPP(nn.Module):
    """Lightweight ASPP with a reduced channel width for NPU/CPU friendliness."""

    def __init__(self, channels: int, hidden: int = 64, rates: tuple[int, ...] = (2, 4, 6)) -> None:
        super().__init__()
        self.reduce = ConvNormAct(channels, hidden, 1)
        self.branches = nn.ModuleList(
            [ConvNormAct(hidden, hidden, 3, dilation=rate) for rate in rates]
        )
        self.global_branch = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(hidden, hidden, 1, bias=False),
            nn.ReLU(inplace=True),
        )
        self.project = ConvNormAct(hidden * (len(rates) + 2), channels, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.reduce(x)
        size = x.shape[-2:]
        features = [x]
        features.extend(branch(x) for branch in self.branches)
        global_feature = self.global_branch(x)
        features.append(F.interpolate(global_feature, size=size, mode="bilinear", align_corners=False))
        return self.project(torch.cat(features, dim=1))


class EfficientViTEncoder(nn.Module):
    """EfficientViT with explicit stride-4/8/16 taps plus the stride-32 stage."""

    def __init__(self, model_name: str = "efficientvit_m2.r224_in1k", pretrained: bool = True) -> None:
        super().__init__()
        self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
        self.out_channels = (32, 64, 128, 192)

    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
        pe = self.model.patch_embed
        x = pe.conv1(x)
        x = pe.relu1(x)
        x = pe.conv2(x)
        x = pe.relu2(x)
        f4 = x
        x = pe.conv3(x)
        x = pe.relu3(x)
        f8 = x
        x = pe.conv4(x)
        f16 = x
        x = self.model.stages[0](x)
        x = self.model.stages[1](x)
        f32 = x
        return [f4, f8, f16, f32]


class TimmFeaturesEncoder(nn.Module):
    """Wrapper around a timm features_only model."""

    def __init__(
        self,
        model_name: str,
        *,
        out_indices: tuple[int, ...],
        pretrained: bool = True,
    ) -> None:
        super().__init__()
        self.model = timm.create_model(
            model_name,
            pretrained=pretrained,
            features_only=True,
            out_indices=out_indices,
        )
        info = self.model.feature_info
        self.out_channels = tuple(info[i]["num_chs"] for i in out_indices)

    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
        return list(self.model(x))


class LiteDPT(nn.Module):
    """Top-down fusion decoder with depthwise residual refinement."""

    def __init__(
        self,
        in_channels: tuple[int, int, int, int],
        channels: tuple[int, int, int, int] = (80, 112, 144, 176),
        *,
        blocks_per_level: int = 2,
        aspp_hidden: int = 64,
        deep_supervision: bool = False,
    ) -> None:
        super().__init__()
        c4, c8, c16, c32 = channels
        self.channels = channels
        self.deep_supervision = deep_supervision
        self.lateral4 = ConvNormAct(in_channels[0], c4, 1)
        self.lateral8 = ConvNormAct(in_channels[1], c8, 1)
        self.lateral16 = ConvNormAct(in_channels[2], c16, 1)
        self.lateral32 = ConvNormAct(in_channels[3], c32, 1)

        self.aspp = ASPP(c32, hidden=aspp_hidden)

        self.fuse32_to16 = ConvNormAct(c32, c16, 1)
        self.fuse16_to8 = ConvNormAct(c16, c8, 1)
        self.fuse8_to4 = ConvNormAct(c8, c4, 1)

        self.refine32 = nn.Sequential(*[DWResidualBlock(c32) for _ in range(blocks_per_level)])
        self.refine16 = nn.Sequential(*[DWResidualBlock(c16) for _ in range(blocks_per_level)])
        self.refine8 = nn.Sequential(*[DWResidualBlock(c8) for _ in range(blocks_per_level)])
        self.refine4 = nn.Sequential(*[DWResidualBlock(c4) for _ in range(blocks_per_level)])

        self.head = nn.Sequential(
            ConvNormAct(c4, c4, 3, groups=c4),
            nn.Conv2d(c4, 1, 1),
        )
        if deep_supervision:
            self.aux_heads = nn.ModuleDict(
                {
                    "16": nn.Conv2d(c16, 1, 1),
                    "8": nn.Conv2d(c8, 1, 1),
                }
            )
        else:
            self.aux_heads = None

    def forward(
        self, features: list[torch.Tensor], *, return_aux: bool = False
    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
        f4, f8, f16, f32 = features
        x4 = self.lateral4(f4)
        x8 = self.lateral8(f8)
        x16 = self.lateral16(f16)
        x32 = self.aspp(self.lateral32(f32))
        x32 = self.refine32(x32)

        x16 = x16 + F.interpolate(
            self.fuse32_to16(x32), size=x16.shape[-2:], mode="bilinear", align_corners=False
        )
        x16 = self.refine16(x16)
        aux16 = self.aux_heads["16"](x16) if self.aux_heads is not None else None

        x8 = x8 + F.interpolate(
            self.fuse16_to8(x16), size=x8.shape[-2:], mode="bilinear", align_corners=False
        )
        x8 = self.refine8(x8)
        aux8 = self.aux_heads["8"](x8) if self.aux_heads is not None else None

        x4 = x4 + F.interpolate(
            self.fuse8_to4(x8), size=x4.shape[-2:], mode="bilinear", align_corners=False
        )
        x4 = self.refine4(x4)

        depth = F.softplus(self.head(x4))
        if return_aux and aux16 is not None and aux8 is not None:
            return depth, {
                "16": F.softplus(aux16).squeeze(1),
                "8": F.softplus(aux8).squeeze(1),
            }
        return depth


class CompactDepthModel(nn.Module):
    """Encoder + Lite-DPT depth model with padded, full-resolution output."""

    ENCODERS = {
        "efficientvit_m1": ("efficientvit_m1.r224_in1k", "efficientvit"),
        "efficientvit_m2": ("efficientvit_m2.r224_in1k", "efficientvit"),
        "fastvit_t8": ("fastvit_t8.apple_in1k", "timm"),
        "mobilenetv4_conv_s": ("mobilenetv4_conv_small.e2400_r224_in1k", "timm"),
        "repvit_m0_9": ("repvit_m0_9.dist_300e_in1k", "timm"),
    }
    TIMM_OUT_INDICES = {
        "fastvit_t8": (0, 1, 2, 3),
        "mobilenetv4_conv_s": (1, 2, 3, 4),
        "repvit_m0_9": (0, 1, 2, 3),
    }

    def __init__(
        self,
        encoder: str = "efficientvit_m2",
        *,
        pretrained: bool = True,
        channels: tuple[int, int, int, int] = (80, 112, 144, 176),
        blocks_per_level: int = 2,
        aspp_hidden: int = 64,
        deep_supervision: bool = False,
    ) -> None:
        super().__init__()
        if encoder not in self.ENCODERS:
            raise ValueError(f"unknown encoder {encoder!r}; choices: {sorted(self.ENCODERS)}")
        model_name, kind = self.ENCODERS[encoder]
        if kind == "efficientvit":
            self.encoder = EfficientViTEncoder(model_name, pretrained=pretrained)
        else:
            self.encoder = TimmFeaturesEncoder(
                model_name,
                out_indices=self.TIMM_OUT_INDICES[encoder],
                pretrained=pretrained,
            )
        self.encoder_name = encoder
        self.decoder = LiteDPT(
            self.encoder.out_channels,
            channels=channels,
            blocks_per_level=blocks_per_level,
            aspp_hidden=aspp_hidden,
            deep_supervision=deep_supervision,
        )
        self.channels = channels
        self.blocks_per_level = blocks_per_level
        self.aspp_hidden = aspp_hidden
        self.deep_supervision = deep_supervision

    def forward(
        self, pixel_values: torch.Tensor, *, return_aux: bool = False
    ) -> CompactDepthOutput:
        height, width = pixel_values.shape[-2:]
        pad_h = (-height) % 32
        pad_w = (-width) % 32
        x = pixel_values
        if pad_h or pad_w:
            x = F.pad(x, (0, pad_w, 0, pad_h), mode="replicate")
        features = self.encoder(x)
        decoder_out = self.decoder(features, return_aux=return_aux)
        if isinstance(decoder_out, tuple):
            depth, aux = decoder_out
        else:
            depth, aux = decoder_out, None
        depth = F.interpolate(depth, size=(height, width), mode="bilinear", align_corners=False)
        return CompactDepthOutput(predicted_depth=depth.squeeze(1), aux=aux)

    def config(self) -> dict[str, Any]:
        return {
            "encoder": self.encoder_name,
            "channels": list(self.channels),
            "blocks_per_level": self.blocks_per_level,
            "aspp_hidden": self.aspp_hidden,
            "deep_supervision": self.deep_supervision,
        }


def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())


def save_compact(model: CompactDepthModel, path, **extra: Any) -> None:
    torch.save(
        {
            "state_dict": model.state_dict(),
            "config": model.config(),
            "params": count_parameters(model),
            **extra,
        },
        path,
    )


def load_compact(
    path, *, map_location="cpu", deep_supervision: bool | None = None
) -> CompactDepthModel:
    checkpoint = torch.load(path, map_location=map_location, weights_only=False)
    config = checkpoint["config"]
    use_deep_supervision = (
        config.get("deep_supervision", False)
        if deep_supervision is None
        else deep_supervision
    )
    model = CompactDepthModel(
        config["encoder"],
        pretrained=False,
        channels=tuple(config["channels"]),
        blocks_per_level=config.get("blocks_per_level", 2),
        aspp_hidden=config.get("aspp_hidden", 64),
        deep_supervision=use_deep_supervision,
    )
    try:
        model.load_state_dict(checkpoint["state_dict"])
    except RuntimeError:
        # Adding auxiliary heads changes the state-dict keys. The main weights
        # still match; the new heads start randomly and are trained.
        model.load_state_dict(checkpoint["state_dict"], strict=False)
    return model