"""Standalone architecture extracted from the original lidarx implementation."""
import torch
from torch import nn


def coarse_projection_features(x: torch.Tensor) -> torch.Tensor:
    """Return `projection_coarse_32x8x8` features for `[B, 1, D, H, W]` tensors."""
    if x.ndim != 5:
        msg = f"expected input with shape [batch, channel, depth, height, width], got {x.shape}"
        raise ValueError(msg)
    if x.shape[1] != 1:
        msg = f"expected one histogram channel, got {x.shape[1]}"
        raise ValueError(msg)

    x = x.squeeze(1)
    depth = x.shape[1]
    height = x.shape[2]
    width = x.shape[3]
    if depth % 32 != 0 or height % 8 != 0 or width % 8 != 0:
        msg = (
            "hist coarse projection requires depth divisible by 32 and "
            f"height/width divisible by 8, got {(depth, height, width)}"
        )
        raise ValueError(msg)

    global_features = torch.cat(
        [
            x.mean(dim=(1, 2, 3), keepdim=False).unsqueeze(1),
            x.std(dim=(1, 2, 3), unbiased=False, keepdim=False).unsqueeze(1),
            x.amin(dim=(1, 2, 3), keepdim=False).unsqueeze(1),
            x.amax(dim=(1, 2, 3), keepdim=False).unsqueeze(1),
        ],
        dim=1,
    )
    depth_features = torch.cat(
        [
            x.mean(dim=(2, 3)),
            x.std(dim=(2, 3), unbiased=False),
        ],
        dim=1,
    )
    row_features = x.mean(dim=(1, 3))
    col_features = x.mean(dim=(1, 2))
    coarse = (
        x.reshape(
            x.shape[0],
            32,
            depth // 32,
            8,
            height // 8,
            8,
            width // 8,
        )
        .mean(dim=(2, 4, 6))
        .reshape(x.shape[0], -1)
    )
    return torch.cat(
        [global_features, depth_features, row_features, col_features, coarse],
        dim=1,
    )


def _coarse_grid_features(
    x: torch.Tensor,
    *,
    depth_bins: int,
    row_bins: int,
    col_bins: int,
) -> torch.Tensor:
    depth = x.shape[1]
    height = x.shape[2]
    width = x.shape[3]
    if depth % depth_bins != 0 or height % row_bins != 0 or width % col_bins != 0:
        msg = (
            "hist coarse grid requires dimensions divisible by requested bins, "
            f"got volume {(depth, height, width)} and bins {(depth_bins, row_bins, col_bins)}"
        )
        raise ValueError(msg)

    return (
        x.reshape(
            x.shape[0],
            depth_bins,
            depth // depth_bins,
            row_bins,
            height // row_bins,
            col_bins,
            width // col_bins,
        )
        .mean(dim=(2, 4, 6))
        .reshape(x.shape[0], -1)
    )


def multiscale_coarse_projection_features(x: torch.Tensor) -> torch.Tensor:
    """Return coarse histogram features with added low-resolution context views."""
    base = coarse_projection_features(x)
    volume = x.squeeze(1)
    depth_context = _coarse_grid_features(volume, depth_bins=16, row_bins=8, col_bins=8)
    spatial_context = _coarse_grid_features(volume, depth_bins=32, row_bins=4, col_bins=4)
    return torch.cat([base, depth_context, spatial_context], dim=1)


class HistCoarseMultiScaleMLP(nn.Module):
    """Small supervised MLP over multi-scale coarse histogram projections.

    This keeps the strong non-backbone inductive bias of the coarse-only model,
    but gives the classifier both the best 32x8x8 block view and two coarser
    context views before a lightweight gated residual head.
    """

    def __init__(self, hidden_features: int = 128, dropout: float = 0.2) -> None:
        super().__init__()
        bottleneck_features = max(1, hidden_features // 2)
        self.input_norm = nn.LazyBatchNorm1d()
        self.input_projection = nn.Sequential(
            nn.LazyLinear(hidden_features),
            nn.SiLU(inplace=True),
            nn.Dropout(dropout),
        )
        self.residual = nn.Sequential(
            nn.Linear(hidden_features, hidden_features),
            nn.SiLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(hidden_features, hidden_features),
        )
        self.gate = nn.Sequential(
            nn.Linear(hidden_features, hidden_features),
            nn.Sigmoid(),
        )
        self.head = nn.Sequential(
            nn.LayerNorm(hidden_features),
            nn.SiLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(hidden_features, bottleneck_features),
            nn.SiLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(bottleneck_features, 1),
        )

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        """Return the learned 128-D representation immediately before the classifier."""
        features = self.input_norm(multiscale_coarse_projection_features(x))
        embedding = self.input_projection(features)
        return embedding + self.gate(embedding) * self.residual(embedding)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.head(self.forward_features(x))