"""
V5-Aware Analytic-Centered Structured Semantic VAE (AC-SSVAE)
================================================================
Stage-2 feature extractor for ICWDS.

Design goals
------------
1. Consume WaveSystemSetParserV5 soft wave-system outputs instead of hard labels.
2. Use an analytic physical descriptor as the prior center, and learn only bounded
   semantic corrections.
3. Canonicalize each wave system before shape encoding so the free latent does not
   waste capacity on location, direction, scale, or energy.
4. Keep a small causal shape latent. The residual decoder is constructed as
       Delta(z, s) = F(z, s) - F(0, s)
   so z=0 exactly returns the semantic base.
5. Expose deterministic decoding APIs for quantization, intervention, and later
   30-byte packet design.

Internal semantic vector (9D)
-----------------------------
[log_energy,
 peak_frequency_01,
 sin_peak_direction,
 cos_peak_direction,
 frequency_spread_01,
 direction_spread_over_pi,
 f_theta_correlation,
 frequency_skew,
 direction_skew]

The packet-facing direction is still one circular quantity. sin/cos are only the
neural internal representation.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, asdict
from typing import Dict, Optional, Tuple

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


@dataclass
class V5AwareACSSVAEConfig:
    n_freqs: int = 47
    n_dirs: int = 72
    n_slots: int = 6
    semantic_dim: int = 9
    shape_dim: int = 4
    context_dim: int = 24

    canonical_h: int = 32
    canonical_w: int = 48
    canonical_f_extent: float = 3.0
    canonical_t_extent: float = 3.0

    encoder_width: int = 48
    hidden_dim: int = 160
    decoder_width: int = 64

    # Bounded semantic correction magnitudes around analytic center.
    # [logE, f, sin/cos correction, spreads, rho, skews]
    delta_loge: float = 0.35
    delta_f: float = 0.08
    delta_angle_rad: float = math.radians(18.0)
    delta_log_spread: float = 0.45
    delta_rho_logit: float = 0.75
    delta_skew: float = 0.90

    min_spread_f: float = 0.010
    max_spread_f: float = 0.45
    min_spread_t: float = 0.015
    max_spread_t: float = 0.95
    max_abs_rho: float = 0.92
    max_abs_skew: float = 3.0

    shape_log_residual_scale: float = 2.25
    energy_log_den: float = math.log1p(47 * 72)
    eps: float = 1e-6

    def to_dict(self):
        return asdict(self)


class ConvNormAct(nn.Module):
    def __init__(self, ci: int, co: int, stride: int = 1):
        super().__init__()
        self.conv = nn.Conv2d(ci, co, 3, stride=stride, padding=1, bias=False)
        self.norm = nn.GroupNorm(min(8, co), co)
        self.act = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.act(self.norm(self.conv(x)))


class ResBlock(nn.Module):
    def __init__(self, ch: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(ch, ch, 3, padding=1, bias=False),
            nn.GroupNorm(min(8, ch), ch),
            nn.GELU(),
            nn.Conv2d(ch, ch, 3, padding=1, bias=False),
            nn.GroupNorm(min(8, ch), ch),
        )
        self.act = nn.GELU()

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


def _wrap_angle(x: torch.Tensor) -> torch.Tensor:
    return torch.atan2(torch.sin(x), torch.cos(x))


def _atanh_safe(x: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
    x = x.clamp(-1 + eps, 1 - eps)
    return 0.5 * (torch.log1p(x) - torch.log1p(-x))


def reparameterize(mu: torch.Tensor, logvar: torch.Tensor, stochastic: bool) -> torch.Tensor:
    if not stochastic:
        return mu
    std = torch.exp(0.5 * logvar)
    return mu + std * torch.randn_like(std)


@torch.no_grad()
def analytic_semantics(
    energy: torch.Tensor,
    mask_prob: torch.Tensor,
    core_prob: Optional[torch.Tensor] = None,
    eps: float = 1e-6,
    energy_log_den: Optional[float] = None,
) -> torch.Tensor:
    """Compute stable 9D physical descriptors from soft wave-system masks.

    Parameters
    ----------
    energy: [B,1,H,W] or [B,H,W], non-negative display-domain energy.
    mask_prob: [B,K,H,W].
    core_prob: optional [B,K,H,W]. Core weighting sharpens peak estimation.
    """
    if energy.ndim == 3:
        energy = energy[:, None]
    B, K, H, W = mask_prob.shape
    dtype, device = energy.dtype, energy.device
    if energy_log_den is None:
        energy_log_den = math.log1p(H * W)

    E = torch.nan_to_num(energy, nan=0.0, posinf=1.0, neginf=0.0).clamp_min(0)
    M = torch.nan_to_num(mask_prob, nan=0.0).clamp(0, 1)
    part = E * M
    mass = part.sum(dim=(-2, -1)).clamp_min(eps)
    loge = torch.log1p(mass) / float(energy_log_den)

    fgrid = torch.linspace(0, 1, H, device=device, dtype=dtype).view(1, 1, H, 1)
    theta = torch.linspace(0, 2 * math.pi, W + 1, device=device, dtype=dtype)[:W]
    tgrid = theta.view(1, 1, 1, W)

    peak_src = part
    if core_prob is not None:
        core = torch.nan_to_num(core_prob, nan=0.0).clamp(0, 1)
        peak_src = part * (0.25 + 0.75 * core)

    # Smooth soft-peak estimator. It is more stable than hard argmax on noisy spectra.
    pnorm = peak_src / peak_src.amax(dim=(-2, -1), keepdim=True).clamp_min(eps)
    peak_w = torch.softmax((12.0 * pnorm).flatten(2), dim=-1).view(B, K, H, W)
    fp = (peak_w * fgrid).sum(dim=(-2, -1))
    sx = (peak_w * torch.sin(tgrid)).sum(dim=(-2, -1))
    cx = (peak_w * torch.cos(tgrid)).sum(dim=(-2, -1))
    theta_p = torch.atan2(sx, cx)
    sinp = torch.sin(theta_p)
    cosp = torch.cos(theta_p)

    w = part / mass[:, :, None, None]
    df = fgrid - fp[:, :, None, None]
    dt = _wrap_angle(tgrid - theta_p[:, :, None, None]) / math.pi

    sf = torch.sqrt((w * df.square()).sum(dim=(-2, -1)).clamp_min(eps))
    st = torch.sqrt((w * dt.square()).sum(dim=(-2, -1)).clamp_min(eps))

    zf = df / sf[:, :, None, None].clamp_min(1e-3)
    zt = dt / st[:, :, None, None].clamp_min(1e-3)
    rho = (w * zf * zt).sum(dim=(-2, -1)).clamp(-0.95, 0.95)
    skew_f = (w * zf.pow(3)).sum(dim=(-2, -1)).clamp(-4.0, 4.0)
    skew_t = (w * zt.pow(3)).sum(dim=(-2, -1)).clamp(-4.0, 4.0)

    return torch.stack([
        loge, fp, sinp, cosp, sf, st, rho, skew_f, skew_t
    ], dim=-1)


class Canonicalizer(nn.Module):
    """Differentiably center/scale each slot in frequency-direction coordinates."""
    def __init__(self, cfg: V5AwareACSSVAEConfig):
        super().__init__()
        self.cfg = cfg
        u = torch.linspace(-cfg.canonical_f_extent, cfg.canonical_f_extent, cfg.canonical_h)
        v = torch.linspace(-cfg.canonical_t_extent, cfg.canonical_t_extent, cfg.canonical_w)
        self.register_buffer("u", u.view(1, 1, cfg.canonical_h, 1))
        self.register_buffer("v", v.view(1, 1, 1, cfg.canonical_w))

    def forward(
        self,
        energy: torch.Tensor,
        mask_prob: torch.Tensor,
        core_prob: torch.Tensor,
        support_prob: torch.Tensor,
        semantics: torch.Tensor,
    ) -> torch.Tensor:
        if energy.ndim == 3:
            energy = energy[:, None]
        B, K, H, W = mask_prob.shape
        eps = self.cfg.eps

        E = energy[:, None].expand(B, K, 1, H, W)
        M = mask_prob[:, :, None]
        C = core_prob[:, :, None]
        S = support_prob[:, :, None]
        part = E * M
        cpart = E * C
        spart = E * S

        peak = part.amax(dim=(-2, -1), keepdim=True).clamp_min(eps)
        channels = torch.cat([
            part / peak,
            cpart / peak,
            spart / peak,
            M,
            C,
            S,
        ], dim=2).reshape(B * K, 6, H, W)

        sem = semantics.reshape(B * K, -1)
        fp = sem[:, 1].clamp(0, 1)
        theta = torch.atan2(sem[:, 2], sem[:, 3]) % (2 * math.pi)
        sf = sem[:, 4].clamp(self.cfg.min_spread_f, self.cfg.max_spread_f)
        st = sem[:, 5].clamp(self.cfg.min_spread_t, self.cfg.max_spread_t)

        uf = self.u.to(channels.dtype)
        vt = self.v.to(channels.dtype)
        f_sample = fp[:, None, None, None] + sf[:, None, None, None] * uf
        theta_sample = theta[:, None, None, None] + (math.pi * st[:, None, None, None]) * vt
        theta_frac = torch.remainder(theta_sample, 2 * math.pi) / (2 * math.pi)

        # Tile direction three times and sample from the middle copy.
        tiled = torch.cat([channels, channels, channels], dim=-1)
        y = 2.0 * f_sample - 1.0
        xpix = (1.0 + theta_frac) * W - 0.5
        x = 2.0 * xpix / max(3 * W - 1, 1) - 1.0
        grid = torch.stack([
            x.expand(-1, -1, self.cfg.canonical_h, self.cfg.canonical_w).squeeze(1),
            y.expand(-1, -1, self.cfg.canonical_h, self.cfg.canonical_w).squeeze(1),
        ], dim=-1)
        out = F.grid_sample(
            tiled,
            grid,
            mode="bilinear",
            padding_mode="zeros",
            align_corners=True,
        )
        return out.reshape(B, K, 6, self.cfg.canonical_h, self.cfg.canonical_w)


class SlotEncoder(nn.Module):
    def __init__(self, cfg: V5AwareACSSVAEConfig):
        super().__init__()
        w = cfg.encoder_width
        self.image_net = nn.Sequential(
            ConvNormAct(6, w, 1),
            ResBlock(w),
            ConvNormAct(w, w * 2, 2),
            ResBlock(w * 2),
            ConvNormAct(w * 2, w * 3, 2),
            ResBlock(w * 3),
            ConvNormAct(w * 3, w * 4, 2),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )
        self.context_net = nn.Sequential(
            nn.Linear(cfg.context_dim + cfg.semantic_dim + 1, cfg.hidden_dim // 2),
            nn.LayerNorm(cfg.hidden_dim // 2),
            nn.GELU(),
            nn.Linear(cfg.hidden_dim // 2, cfg.hidden_dim // 2),
            nn.GELU(),
        )
        image_dim = w * 4
        self.fuse = nn.Sequential(
            nn.Linear(image_dim + cfg.hidden_dim // 2, cfg.hidden_dim),
            nn.LayerNorm(cfg.hidden_dim),
            nn.GELU(),
            nn.Linear(cfg.hidden_dim, cfg.hidden_dim),
            nn.GELU(),
        )
        self.sem_mu = nn.Linear(cfg.hidden_dim, cfg.semantic_dim)
        self.sem_logvar = nn.Linear(cfg.hidden_dim, cfg.semantic_dim)
        self.shape_mu = nn.Linear(cfg.hidden_dim, cfg.shape_dim)
        self.shape_logvar = nn.Linear(cfg.hidden_dim, cfg.shape_dim)

        # Start from the analytic center and nearly-zero free shape code.
        nn.init.zeros_(self.sem_mu.weight)
        nn.init.zeros_(self.sem_mu.bias)
        nn.init.constant_(self.sem_logvar.bias, -4.0)
        nn.init.normal_(self.shape_mu.weight, std=0.01)
        nn.init.zeros_(self.shape_mu.bias)
        nn.init.constant_(self.shape_logvar.bias, -3.0)

    def forward(
        self,
        canonical: torch.Tensor,
        analytic_sem: torch.Tensor,
        context: torch.Tensor,
        exist_prob: torch.Tensor,
    ) -> Dict[str, torch.Tensor]:
        B, K = canonical.shape[:2]
        img = self.image_net(canonical.reshape(B * K, *canonical.shape[2:]))
        ctx_in = torch.cat([
            context.reshape(B * K, -1),
            analytic_sem.reshape(B * K, -1),
            exist_prob.reshape(B * K, 1),
        ], dim=-1)
        ctx = self.context_net(ctx_in)
        h = self.fuse(torch.cat([img, ctx], dim=-1))
        sem_mu = self.sem_mu(h).reshape(B, K, -1)
        sem_logvar = self.sem_logvar(h).clamp(-8.0, 3.0).reshape(B, K, -1)
        shape_mu = self.shape_mu(h).reshape(B, K, -1)
        shape_logvar = self.shape_logvar(h).clamp(-8.0, 3.0).reshape(B, K, -1)
        return {
            "sem_delta_mu": sem_mu,
            "sem_delta_logvar": sem_logvar,
            "shape_mu": shape_mu,
            "shape_logvar": shape_logvar,
        }


class SemanticRenderer(nn.Module):
    """Parameter-free generalized elliptical renderer from corrected semantics."""
    def __init__(self, cfg: V5AwareACSSVAEConfig):
        super().__init__()
        self.cfg = cfg
        f = torch.linspace(0, 1, cfg.n_freqs).view(1, 1, cfg.n_freqs, 1)
        t = torch.linspace(0, 2 * math.pi, cfg.n_dirs + 1)[:cfg.n_dirs].view(1, 1, 1, cfg.n_dirs)
        self.register_buffer("fgrid", f)
        self.register_buffer("tgrid", t)

    def forward(self, sem: torch.Tensor) -> torch.Tensor:
        cfg = self.cfg
        eps = cfg.eps
        loge, fp = sem[..., 0], sem[..., 1]
        theta = torch.atan2(sem[..., 2], sem[..., 3])
        sf = sem[..., 4].clamp(cfg.min_spread_f, cfg.max_spread_f)
        st = sem[..., 5].clamp(cfg.min_spread_t, cfg.max_spread_t)
        rho = sem[..., 6].clamp(-cfg.max_abs_rho, cfg.max_abs_rho)
        skew_f = sem[..., 7].clamp(-cfg.max_abs_skew, cfg.max_abs_skew)
        skew_t = sem[..., 8].clamp(-cfg.max_abs_skew, cfg.max_abs_skew)

        df = (self.fgrid - fp[:, :, None, None]) / sf[:, :, None, None].clamp_min(1e-3)
        dt = _wrap_angle(self.tgrid - theta[:, :, None, None])
        dt = dt / (math.pi * st[:, :, None, None].clamp_min(1e-3))
        den = (1.0 - rho.square()).clamp_min(0.08)
        q = (df.square() + dt.square() - 2.0 * rho[:, :, None, None] * df * dt) / den[:, :, None, None]
        shape = torch.exp(-0.5 * q.clamp_max(60.0))
        asym = torch.exp(
            0.28 * skew_f[:, :, None, None] * torch.tanh(df)
            + 0.28 * skew_t[:, :, None, None] * torch.tanh(dt)
        ).clamp(0.15, 6.0)
        shape = shape * asym
        shape = shape / shape.sum(dim=(-2, -1), keepdim=True).clamp_min(eps)
        mass = torch.expm1(loge * cfg.energy_log_den).clamp_min(0.0)
        return shape * mass[:, :, None, None]


class ShapeResidualNet(nn.Module):
    def __init__(self, cfg: V5AwareACSSVAEConfig):
        super().__init__()
        self.cfg = cfg
        w = cfg.decoder_width
        self.fc = nn.Sequential(
            nn.Linear(cfg.semantic_dim + cfg.shape_dim, cfg.hidden_dim),
            nn.GELU(),
            nn.Linear(cfg.hidden_dim, w * 6 * 9),
            nn.GELU(),
        )
        self.net = nn.Sequential(
            ConvNormAct(w, w, 1),
            ResBlock(w),
            nn.Upsample(size=(12, 18), mode="bilinear", align_corners=False),
            ConvNormAct(w, w, 1),
            ResBlock(w),
            nn.Upsample(size=(24, 36), mode="bilinear", align_corners=False),
            ConvNormAct(w, w // 2, 1),
            ResBlock(w // 2),
            nn.Upsample(size=(48, 72), mode="bilinear", align_corners=False),
            ConvNormAct(w // 2, w // 2, 1),
            nn.Conv2d(w // 2, 1, 3, padding=1),
        )

    def field(self, sem: torch.Tensor, zshape: torch.Tensor) -> torch.Tensor:
        B, K = sem.shape[:2]
        x = torch.cat([sem, zshape], dim=-1).reshape(B * K, -1)
        h = self.fc(x).view(B * K, self.cfg.decoder_width, 6, 9)
        out = self.net(h)[:, :, : self.cfg.n_freqs, : self.cfg.n_dirs]
        return out.reshape(B, K, self.cfg.n_freqs, self.cfg.n_dirs)

    def forward(self, sem: torch.Tensor, zshape: torch.Tensor) -> torch.Tensor:
        raw = self.field(sem, zshape)
        zero = self.field(sem, torch.zeros_like(zshape))
        return self.cfg.shape_log_residual_scale * torch.tanh(raw - zero)


class V5AwareStructuredSemanticVAE(nn.Module):
    def __init__(self, cfg: Optional[V5AwareACSSVAEConfig] = None):
        super().__init__()
        self.cfg = cfg or V5AwareACSSVAEConfig()
        self.canonicalizer = Canonicalizer(self.cfg)
        self.encoder = SlotEncoder(self.cfg)
        self.semantic_renderer = SemanticRenderer(self.cfg)
        self.shape_decoder = ShapeResidualNet(self.cfg)

    def apply_semantic_delta(self, base: torch.Tensor, delta: torch.Tensor) -> torch.Tensor:
        c = self.cfg
        out = base.clone()
        out[..., 0] = (base[..., 0] + c.delta_loge * torch.tanh(delta[..., 0])).clamp_min(0.0)
        out[..., 1] = (base[..., 1] + c.delta_f * torch.tanh(delta[..., 1])).clamp(0.0, 1.0)

        theta0 = torch.atan2(base[..., 2], base[..., 3])
        # Use both sin/cos correction channels to form one stable tangent-angle correction.
        dtheta = c.delta_angle_rad * torch.tanh(0.7071 * (delta[..., 2] - delta[..., 3]))
        theta = theta0 + dtheta
        out[..., 2] = torch.sin(theta)
        out[..., 3] = torch.cos(theta)

        out[..., 4] = (
            base[..., 4].clamp_min(c.min_spread_f)
            * torch.exp(c.delta_log_spread * torch.tanh(delta[..., 4]))
        ).clamp(c.min_spread_f, c.max_spread_f)
        out[..., 5] = (
            base[..., 5].clamp_min(c.min_spread_t)
            * torch.exp(c.delta_log_spread * torch.tanh(delta[..., 5]))
        ).clamp(c.min_spread_t, c.max_spread_t)
        out[..., 6] = torch.tanh(
            _atanh_safe(base[..., 6], c.eps)
            + c.delta_rho_logit * torch.tanh(delta[..., 6])
        ).clamp(-c.max_abs_rho, c.max_abs_rho)
        out[..., 7] = (
            base[..., 7] + c.delta_skew * torch.tanh(delta[..., 7])
        ).clamp(-c.max_abs_skew, c.max_abs_skew)
        out[..., 8] = (
            base[..., 8] + c.delta_skew * torch.tanh(delta[..., 8])
        ).clamp(-c.max_abs_skew, c.max_abs_skew)
        return out

    def decode_from_codes(
        self,
        semantics: torch.Tensor,
        shape_code: torch.Tensor,
        exist_prob: Optional[torch.Tensor] = None,
    ) -> Dict[str, torch.Tensor]:
        base = self.semantic_renderer(semantics)
        residual_log = self.shape_decoder(semantics, shape_code)
        # Multiplicative residual in log domain, followed by exact energy renormalization.
        part_hat = (base + self.cfg.eps) * torch.exp(residual_log)
        target_mass = torch.expm1(semantics[..., 0] * self.cfg.energy_log_den).clamp_min(0.0)
        part_hat = part_hat * (
            target_mass[:, :, None, None]
            / part_hat.sum(dim=(-2, -1), keepdim=True).clamp_min(self.cfg.eps)
        )
        if exist_prob is not None:
            part_hat = part_hat * exist_prob[:, :, None, None]
            base = base * exist_prob[:, :, None, None]
        return {
            "semantic_base": base,
            "shape_residual_log": residual_log,
            "part_hat": part_hat,
        }

    def forward(
        self,
        energy: torch.Tensor,
        mask_prob: torch.Tensor,
        core_prob: torch.Tensor,
        support_prob: torch.Tensor,
        exist_prob: torch.Tensor,
        slot_context: torch.Tensor,
        stochastic: bool = True,
        shape_code_override: Optional[torch.Tensor] = None,
        semantic_delta_override: Optional[torch.Tensor] = None,
    ) -> Dict[str, torch.Tensor]:
        analytic = analytic_semantics(
            energy,
            mask_prob,
            core_prob=core_prob,
            eps=self.cfg.eps,
            energy_log_den=self.cfg.energy_log_den,
        )
        canonical = self.canonicalizer(
            energy,
            mask_prob,
            core_prob,
            support_prob,
            analytic,
        )
        enc = self.encoder(canonical, analytic, slot_context, exist_prob)
        sem_delta = (
            semantic_delta_override
            if semantic_delta_override is not None
            else reparameterize(enc["sem_delta_mu"], enc["sem_delta_logvar"], stochastic)
        )
        zshape = (
            shape_code_override
            if shape_code_override is not None
            else reparameterize(enc["shape_mu"], enc["shape_logvar"], stochastic)
        )
        corrected = self.apply_semantic_delta(analytic, sem_delta)
        dec = self.decode_from_codes(corrected, zshape, exist_prob=exist_prob)
        return {
            "analytic_semantics": analytic,
            "corrected_semantics": corrected,
            "semantic_delta": sem_delta,
            "shape_code": zshape,
            "canonical": canonical,
            **enc,
            **dec,
        }

    @property
    def token_dim_per_wave(self) -> int:
        return self.cfg.semantic_dim + self.cfg.shape_dim

    def num_params(self) -> int:
        return sum(p.numel() for p in self.parameters())


def kl_standard_normal(mu: torch.Tensor, logvar: torch.Tensor, weight: Optional[torch.Tensor] = None) -> torch.Tensor:
    kl = -0.5 * (1.0 + logvar - mu.square() - logvar.exp()).sum(dim=-1)
    if weight is None:
        return kl.mean()
    w = weight.float()
    return (kl * w).sum() / w.sum().clamp_min(1.0)


def covariance_penalty(x: torch.Tensor, weight: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Off-diagonal covariance penalty for [B,K,D] latent tensors."""
    D = x.shape[-1]
    flat = x.reshape(-1, D)
    if weight is not None:
        w = weight.reshape(-1).float()
        keep = w > 0.25
        flat = flat[keep]
    if flat.shape[0] < max(4, D):
        return x.new_zeros(())
    flat = flat - flat.mean(dim=0, keepdim=True)
    cov = flat.T @ flat / max(flat.shape[0] - 1, 1)
    off = cov - torch.diag(torch.diag(cov))
    return off.square().mean()


def semantic_distance(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Per-slot interpretable semantic distance with circular direction handling."""
    d = []
    d.append((pred[..., 0] - target[..., 0]).abs())
    d.append(2.0 * (pred[..., 1] - target[..., 1]).abs())
    tp = torch.atan2(pred[..., 2], pred[..., 3])
    tt = torch.atan2(target[..., 2], target[..., 3])
    d.append(_wrap_angle(tp - tt).abs() / math.pi)
    d.append(2.0 * (pred[..., 4] - target[..., 4]).abs())
    d.append(1.5 * (pred[..., 5] - target[..., 5]).abs())
    d.append(0.5 * (pred[..., 6] - target[..., 6]).abs())
    d.append(0.15 * (pred[..., 7] - target[..., 7]).abs())
    d.append(0.15 * (pred[..., 8] - target[..., 8]).abs())
    return torch.stack(d, dim=-1).mean(dim=-1)


def marginal_l1(pred: torch.Tensor, target: torch.Tensor, eps: float = 1e-6) -> Tuple[torch.Tensor, torch.Tensor]:
    p = pred.clamp_min(0)
    t = target.clamp_min(0)
    pf = p.sum(dim=-1); tf = t.sum(dim=-1)
    pt = p.sum(dim=-2); tt = t.sum(dim=-2)
    pf = pf / pf.sum(dim=-1, keepdim=True).clamp_min(eps)
    tf = tf / tf.sum(dim=-1, keepdim=True).clamp_min(eps)
    pt = pt / pt.sum(dim=-1, keepdim=True).clamp_min(eps)
    tt = tt / tt.sum(dim=-1, keepdim=True).clamp_min(eps)
    return (pf - tf).abs().mean(), (pt - tt).abs().mean()


def reconstruction_terms(pred: torch.Tensor, target: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> Dict[str, torch.Tensor]:
    w = weight[:, :, None, None].float()
    denom = w.sum().clamp_min(1.0)
    mse_slot = (pred - target).square().mean(dim=(-2, -1))
    log_slot = (
        torch.log1p(20.0 * pred.clamp_min(0))
        - torch.log1p(20.0 * target.clamp_min(0))
    ).abs().mean(dim=(-2, -1))
    mass_p = pred.sum(dim=(-2, -1))
    mass_t = target.sum(dim=(-2, -1))
    energy_slot = (mass_p - mass_t).abs() / mass_t.clamp_min(eps)
    mse = (mse_slot * weight).sum() / denom
    log_l1 = (log_slot * weight).sum() / denom
    energy = (energy_slot * weight).sum() / denom
    freq, direction = marginal_l1(pred * w, target * w, eps=eps)
    return {
        "mse": mse,
        "log_l1": log_l1,
        "energy": energy,
        "freq": freq,
        "direction": direction,
    }


__all__ = [
    "V5AwareACSSVAEConfig",
    "V5AwareStructuredSemanticVAE",
    "analytic_semantics",
    "kl_standard_normal",
    "covariance_penalty",
    "semantic_distance",
    "reconstruction_terms",
    "marginal_l1",
]