CNN.py

"""
WaveSystem Set Parser V5 for ICWDS
==================================
A lightweight physics-seeded set parser for wave directional spectra.

Key changes from v4.8:
- dual proposal source: external physics proposals + learned seed proposals;
- multi-round relation message passing;
- slot set prediction with per-slot existence probabilities;
- no count-soft -> round -> top-N hard deletion;
- separate core and support mask heads;
- auxiliary wind-sea and wave-age heads for real weak supervision.
"""
from __future__ import annotations

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

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


@dataclass
class WaveSystemSetParserV5Config:
    n_freqs: int = 47
    n_dirs: int = 72
    n_slots: int = 6
    p_phys_max: int = 24
    learned_seed_k: int = 8
    prop_feat_dim: int = 22
    width: int = 32
    depth: int = 4
    node_dim: int = 64
    slot_dim: int = 64
    msg_rounds: int = 3
    pair_feat_dim: int = 8
    seed_sigma_f: float = 0.085
    seed_sigma_t: float = 0.16  # fraction of pi
    prior_gain_core: float = 3.2
    prior_gain_support: float = 2.0
    min_exist_bias: float = -0.8

    @property
    def p_total(self) -> int:
        return self.p_phys_max + self.learned_seed_k

    def to_dict(self):
        return asdict(self)


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

    def forward(self, x):
        return self.act(self.norm(self.pw(self.dw(x))))


class PhysicsAwareModule(nn.Module):
    def __init__(self, ch: int, n_dirs: int):
        super().__init__()
        dirs = torch.linspace(0, 2 * math.pi, n_dirs + 1)[:n_dirs]
        self.n_dirs = n_dirs
        self.register_buffer("cos_d", torch.cos(dirs).view(1, 1, 1, n_dirs))
        self.register_buffer("sin_d", torch.sin(dirs).view(1, 1, 1, n_dirs))
        self.fuse = nn.Conv2d(ch + 3, ch, 1, bias=False)
        self.norm = nn.GroupNorm(min(8, ch), ch)
        self.act = nn.GELU()

    def _phys_features(self, E):
        eps = 1e-8
        En = torch.nan_to_num(E, nan=0.0, posinf=1.0, neginf=0.0).clamp_min(0)
        En = En / (En.amax(dim=(2, 3), keepdim=True) + eps)
        nf = En.shape[2]
        rev_cumsum = torch.flip(torch.cumsum(torch.flip(En, dims=[2]), dim=2), dims=[2])
        col_sum = En.sum(dim=2, keepdim=True) + eps
        hf_tail = rev_cumsum / col_sum
        cx = (En * self.cos_d).sum(dim=3, keepdim=True)
        cy = (En * self.sin_d).sum(dim=3, keepdim=True)
        row_sum = En.sum(dim=3, keepdim=True) + eps
        dir_conc = torch.sqrt(cx ** 2 + cy ** 2) / row_sum
        dir_conc = dir_conc.expand(-1, -1, -1, self.n_dirs)
        fcoord = torch.linspace(0, 1, nf, device=E.device, dtype=E.dtype).view(1, 1, nf, 1)
        fc = (En * fcoord).sum(dim=2, keepdim=True) / col_sum
        spread = torch.sqrt((En * (fcoord - fc) ** 2).sum(dim=2, keepdim=True) / col_sum)
        spread = spread.expand(-1, -1, nf, -1)
        return torch.cat([hf_tail, dir_conc, spread], dim=1)

    def forward(self, h, E):
        return self.act(self.norm(self.fuse(torch.cat([h, self._phys_features(E)], dim=1))))


class RelationBlock(nn.Module):
    def __init__(self, dim: int, pair_dim: int):
        super().__init__()
        self.q = nn.Linear(dim, dim, bias=False)
        self.k = nn.Linear(dim, dim, bias=False)
        self.v = nn.Linear(dim, dim, bias=False)
        self.edge_bias = nn.Sequential(
            nn.Linear(pair_dim, dim // 2), nn.GELU(), nn.Linear(dim // 2, 1)
        )
        self.update = nn.Sequential(
            nn.Linear(dim * 2, dim * 2), nn.GELU(), nn.Linear(dim * 2, dim)
        )
        self.norm = nn.LayerNorm(dim)

    def forward(self, node, pair_feat, valid):
        # node [B,P,D], pair_feat [B,P,P,F], valid [B,P]
        D = node.shape[-1]
        q = self.q(node)
        k = self.k(node)
        v = self.v(node)
        logits = torch.einsum("bpd,bqd->bpq", q, k) / math.sqrt(D)
        logits = logits + self.edge_bias(pair_feat).squeeze(-1)
        pair_valid = valid[:, :, None].bool() & valid[:, None, :].bool()
        logits = logits.masked_fill(~pair_valid, -30.0)
        attn = torch.softmax(logits, dim=-1)
        attn = torch.nan_to_num(attn, nan=0.0)
        msg = torch.einsum("bpq,bqd->bpd", attn, v)
        upd = self.update(torch.cat([node, msg], dim=-1))
        out = self.norm(node + upd)
        return out * valid[:, :, None]


class WaveSystemSetParserV5(nn.Module):
    def __init__(self, cfg: Optional[WaveSystemSetParserV5Config] = None):
        super().__init__()
        self.cfg = cfg or WaveSystemSetParserV5Config()
        c = self.cfg

        # Names intentionally match the v4.8 backbone for partial warm-start.
        self.stem = nn.Conv2d(4, c.width, 3, padding=1)
        self.stem_norm = nn.GroupNorm(min(8, c.width), c.width)
        self.physics = PhysicsAwareModule(c.width, c.n_dirs)
        self.blocks = nn.ModuleList([DepthwiseSeparable(c.width, c.width) for _ in range(c.depth)])
        self.global_pool = nn.AdaptiveAvgPool2d(1)

        self.seed_head = nn.Sequential(
            nn.Conv2d(c.width, c.width, 3, padding=1), nn.GELU(), nn.Conv2d(c.width, 1, 1)
        )
        self.global_support_head = nn.Sequential(
            nn.Conv2d(c.width, c.width, 3, padding=1), nn.GELU(), nn.Conv2d(c.width, 1, 1)
        )

        self.node_mlp = nn.Sequential(
            nn.Linear(c.width + c.prop_feat_dim, c.node_dim), nn.GELU(),
            nn.Linear(c.node_dim, c.node_dim), nn.GELU(),
        )
        self.node_keep = nn.Linear(c.node_dim, 1)
        self.relation = nn.ModuleList([
            RelationBlock(c.node_dim, c.pair_feat_dim) for _ in range(c.msg_rounds)
        ])
        self.edge_head = nn.Sequential(
            nn.Linear(4 * c.node_dim + c.pair_feat_dim, c.node_dim), nn.GELU(),
            nn.Linear(c.node_dim, c.node_dim), nn.GELU(), nn.Linear(c.node_dim, 1)
        )

        self.slot_queries = nn.Parameter(torch.randn(c.n_slots, c.slot_dim) * 0.02)
        self.slot_q = nn.Linear(c.slot_dim, c.node_dim, bias=False)
        self.slot_k = nn.Linear(c.node_dim, c.node_dim, bias=False)
        self.slot_v = nn.Linear(c.node_dim, c.slot_dim, bias=False)
        self.slot_update = nn.Sequential(
            nn.Linear(c.slot_dim * 2, c.slot_dim * 2), nn.GELU(), nn.Linear(c.slot_dim * 2, c.slot_dim)
        )
        self.slot_norm = nn.LayerNorm(c.slot_dim)
        self.exist_head = nn.Linear(c.slot_dim, 1)
        nn.init.constant_(self.exist_head.bias, c.min_exist_bias)
        self.count_head = nn.Sequential(
            nn.Linear(c.width + c.slot_dim, c.slot_dim), nn.GELU(), nn.Linear(c.slot_dim, c.n_slots + 1)
        )

        self.pixel_embed = nn.Conv2d(c.width, c.slot_dim, 1, bias=False)
        self.core_slot_proj = nn.Linear(c.slot_dim, c.slot_dim, bias=False)
        self.support_slot_proj = nn.Linear(c.slot_dim, c.slot_dim, bias=False)
        self.bg_head = nn.Sequential(
            nn.Conv2d(c.width, c.width, 3, padding=1), nn.GELU(), nn.Conv2d(c.width, 1, 1)
        )

        self.windsea_head = nn.Sequential(
            nn.Conv2d(c.width, c.width // 2, 3, padding=1), nn.GELU(), nn.Conv2d(c.width // 2, 1, 1)
        )
        self.wave_age_head = nn.Sequential(
            nn.Conv2d(c.width, c.width // 2, 3, padding=1), nn.GELU(), nn.Conv2d(c.width // 2, 1, 1)
        )

        self.register_buffer("fgrid", torch.linspace(0, 1, c.n_freqs).view(1, 1, c.n_freqs, 1))
        ang = torch.linspace(0, 2 * math.pi, c.n_dirs + 1)[:c.n_dirs]
        self.register_buffer("tgrid", ang.view(1, 1, 1, c.n_dirs))
        self.register_buffer("sin_grid", torch.sin(ang).view(1, 1, 1, c.n_dirs))
        self.register_buffer("cos_grid", torch.cos(ang).view(1, 1, 1, c.n_dirs))
        self._coord_cache = None

    def _coord_channels(self, B, device, dtype):
        if self._coord_cache is None:
            nf, nd = self.cfg.n_freqs, self.cfg.n_dirs
            fcoord = torch.linspace(0, 1, nf).view(1, 1, nf, 1).expand(1, 1, nf, nd)
            ang = torch.linspace(0, 2 * math.pi, nd + 1)[:nd].view(1, 1, 1, nd).expand(1, 1, nf, nd)
            self._coord_cache = torch.cat([fcoord, torch.sin(ang), torch.cos(ang)], dim=1)
        return self._coord_cache.to(device=device, dtype=dtype).expand(B, -1, -1, -1)

    @staticmethod
    def _safe_softmax(x, dim=-1):
        x = torch.nan_to_num(x, nan=0.0, posinf=30.0, neginf=-30.0).clamp(-30, 30)
        x = x - x.max(dim=dim, keepdim=True).values.detach()
        return torch.nan_to_num(torch.softmax(x, dim=dim), nan=0.0)

    def _periodic_max_pool(self, x, kernel=3):
        pad = kernel // 2
        xp = torch.cat([x[..., -pad:], x, x[..., :pad]], dim=-1)
        yp = F.max_pool2d(xp, kernel_size=kernel, stride=1, padding=(pad, 0))
        return yp

    def _learned_proposals(self, E01, seed_logits, support_logits) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        c = self.cfg
        B, _, H, W = seed_logits.shape
        seed = torch.sigmoid(seed_logits)
        support = torch.sigmoid(support_logits)
        pooled = self._periodic_max_pool(seed, 3)
        peaks = seed * (seed >= pooled - 1e-6).float()
        vals, idx = torch.topk(peaks.flatten(1), k=c.learned_seed_k, dim=1)
        rr = (idx // W).float() / max(H - 1, 1)
        cc = idx % W
        theta = cc.float() / W * (2 * math.pi)

        fg = self.fgrid.to(E01.dtype)
        tg = self.tgrid.to(E01.dtype)
        df = (fg - rr[:, :, None, None]) / c.seed_sigma_f
        dt = torch.atan2(
            torch.sin(tg - theta[:, :, None, None]),
            torch.cos(tg - theta[:, :, None, None]),
        ) / (math.pi * c.seed_sigma_t)
        gauss = torch.exp(-0.5 * (df * df + dt * dt))
        masks = (gauss * (0.15 + 0.85 * support)).clamp(0, 1)
        valid = (vals > 0.03).float()
        masks = masks * valid[:, :, None, None]

        total = E01.sum(dim=(2, 3)).clamp_min(1e-6)
        mass = torch.einsum("bchw,bphw->bp", E01, masks) / total
        peak = vals.clamp(0, 1)
        area = masks.mean(dim=(2, 3))
        sf = torch.full_like(mass, c.seed_sigma_f)
        st = torch.full_like(mass, c.seed_sigma_t)
        zeros = torch.zeros_like(mass)
        ones = torch.ones_like(mass)
        sin_t = torch.sin(theta)
        cos_t = torch.cos(theta)
        quality = peak
        feats = torch.stack([
            mass, peak, area, rr, sin_t, cos_t, sf, st,
            area, rr, torch.full_like(mass, c.learned_seed_k / float(c.p_total)), quality,
            peak, zeros, ones * 0.5, mass, ones, ones, zeros, peak, ones, zeros,
        ], dim=-1)
        return masks, feats, valid

    def _node_pool(self, h, masks, valid):
        denom = masks.flatten(2).sum(dim=2).clamp_min(1.0)
        pooled = torch.einsum("bchw,bphw->bpc", h, masks) / denom[:, :, None]
        return pooled * valid[:, :, None]

    def _pair_features(self, feats):
        fi = feats[:, :, 3][:, :, None]; fj = feats[:, :, 3][:, None, :]
        d_f = (fi - fj).abs()
        si = feats[:, :, 4][:, :, None]; ci = feats[:, :, 5][:, :, None]
        sj = feats[:, :, 4][:, None, :]; cj = feats[:, :, 5][:, None, :]
        dot = (si * sj + ci * cj).clamp(-1 + 1e-5, 1 - 1e-5)
        d_t = torch.acos(dot) / math.pi
        mi = feats[:, :, 0][:, :, None]; mj = feats[:, :, 0][:, None, :]
        d_m = (mi - mj).abs()
        sfi = feats[:, :, 6][:, :, None]; sfj = feats[:, :, 6][:, None, :]
        sti = feats[:, :, 7][:, :, None]; stj = feats[:, :, 7][:, None, :]
        d_sf = (sfi - sfj).abs(); d_st = (sti - stj).abs()
        prom = torch.minimum(feats[:, :, 12][:, :, None], feats[:, :, 12][:, None, :])
        stripe = torch.maximum(feats[:, :, 13][:, :, None], feats[:, :, 13][:, None, :])
        qual = torch.minimum(feats[:, :, 11][:, :, None], feats[:, :, 11][:, None, :])
        return torch.stack([d_f, d_t, d_m, d_sf, d_st, prom, stripe, qual], dim=-1)

    def _slot_reasoning(self, node, node_keep, valid):
        B = node.shape[0]
        slots = self.slot_queries[None].expand(B, -1, -1)
        key = self.slot_k(node)
        value = self.slot_v(node)
        for _ in range(2):
            q = self.slot_q(slots)
            score = torch.einsum("bkd,bpd->bkp", q, key) / math.sqrt(key.shape[-1])
            score = score + torch.log(node_keep[:, None, :].clamp_min(1e-4))
            score = score.masked_fill(valid[:, None, :] <= 0, -30.0)
            attn = self._safe_softmax(score, dim=-1)
            ctx = torch.einsum("bkp,bpd->bkd", attn, value)
            slots = self.slot_norm(slots + self.slot_update(torch.cat([slots, ctx], dim=-1)))
        return slots

    def forward(self, x, prop_masks, prop_feats, prop_valid) -> Dict[str, torch.Tensor]:
        c = self.cfg
        B, _, H, W = x.shape
        prop_masks = torch.nan_to_num(prop_masks.float(), nan=0.0).clamp(0, 1)
        prop_feats = torch.nan_to_num(prop_feats.float(), nan=0.0, posinf=5.0, neginf=-5.0).clamp(-5, 5)
        prop_valid = prop_valid.float().clamp(0, 1)

        h_in = torch.cat([x, self._coord_channels(B, x.device, x.dtype)], dim=1)
        h = F.gelu(self.stem_norm(self.stem(h_in)))
        E01 = ((x + 1.0) * 0.5).clamp(0, 1)
        h = h + self.physics(h, E01)
        for blk in self.blocks:
            h = h + blk(h)

        seed_logits = self.seed_head(h)
        global_support_logits = self.global_support_head(h)
        learned_masks, learned_feats, learned_valid = self._learned_proposals(E01, seed_logits, global_support_logits)

        masks = torch.cat([prop_masks, learned_masks], dim=1)
        feats = torch.cat([prop_feats, learned_feats], dim=1)
        valid = torch.cat([prop_valid, learned_valid], dim=1)

        node_pool = self._node_pool(h, masks, valid)
        node = self.node_mlp(torch.cat([node_pool, feats], dim=-1)) * valid[:, :, None]
        pair_feat = self._pair_features(feats)
        for block in self.relation:
            node = block(node, pair_feat, valid)

        node_keep_logit = self.node_keep(node).squeeze(-1).masked_fill(valid <= 0, -20.0)
        node_keep = torch.sigmoid(node_keep_logit) * valid

        # Pair merge probability is diagnostic/training supervision; set prediction itself is not hard-merged.
        P = node.shape[1]
        ni = node[:, :, None, :].expand(B, P, P, -1)
        nj = node[:, None, :, :].expand(B, P, P, -1)
        edge_in = torch.cat([ni, nj, (ni - nj).abs(), ni * nj, pair_feat], dim=-1)
        edge_logits = self.edge_head(edge_in).squeeze(-1)
        pair_valid = valid[:, :, None].bool() & valid[:, None, :].bool()
        eye = torch.eye(P, device=x.device, dtype=torch.bool)[None]
        edge_valid = pair_valid & (~eye)
        edge_logits = edge_logits.masked_fill(~edge_valid, 0.0)

        slots = self._slot_reasoning(node, node_keep, valid)
        exist_logit = self.exist_head(slots).squeeze(-1)
        exist_prob = torch.sigmoid(exist_logit)

        slot_key = self.slot_q(slots)
        node_slot_logits = torch.einsum("bpd,bkd->bpk", node, slot_key) / math.sqrt(node.shape[-1])
        node_slot_logits = node_slot_logits.masked_fill(valid[:, :, None] <= 0, -20.0)
        node_slot = self._safe_softmax(node_slot_logits, dim=-1) * valid[:, :, None]
        assign = node_slot * node_keep[:, :, None]
        prior_signal = torch.einsum("bpk,bphw->bkhw", assign, masks)
        prior_signal = prior_signal / prior_signal.amax(dim=(2, 3), keepdim=True).clamp_min(1e-6)
        prior_signal = prior_signal.clamp(0, 1)
        prior_expand = F.max_pool2d(prior_signal, 5, stride=1, padding=2)

        pix = self.pixel_embed(h)
        core_vec = self.core_slot_proj(slots)
        support_vec = self.support_slot_proj(slots)
        core_logits = torch.einsum("bkd,bdhw->bkhw", core_vec, pix) / math.sqrt(c.slot_dim)
        support_logits = torch.einsum("bkd,bdhw->bkhw", support_vec, pix) / math.sqrt(c.slot_dim)
        core_logits = core_logits + c.prior_gain_core * (prior_signal - 0.35)
        support_logits = support_logits + c.prior_gain_support * (prior_expand - 0.30) + 1.25 * (torch.sigmoid(global_support_logits) - 0.5)
        core_prob = torch.sigmoid(core_logits)
        support_prob = torch.sigmoid(support_logits)
        mask_prob = (support_prob * (0.35 + 0.65 * core_prob)).clamp(0, 1)

        slot_score = mask_prob * exist_prob[:, :, None, None]
        bg_logit = self.bg_head(h)
        bg_score = torch.sigmoid(bg_logit)
        scores = torch.cat([slot_score, bg_score], dim=1).clamp_min(1e-6)
        prob = scores / scores.sum(dim=1, keepdim=True).clamp_min(1e-6)

        global_feat = self.global_pool(h).flatten(1)
        count_logits = self.count_head(torch.cat([global_feat, slots.mean(dim=1)], dim=-1))

        return {
            "prob": prob,
            "mask_prob": mask_prob,
            "core_prob": core_prob,
            "support_prob": support_prob,
            "exist_logit": exist_logit,
            "exist_prob": exist_prob,
            "count_logits": count_logits,
            "seed_logits": seed_logits,
            "global_support_logits": global_support_logits,
            "windsea_logits": self.windsea_head(h),
            "wave_age_pred": torch.sigmoid(self.wave_age_head(h)),
            "node_keep_logit": node_keep_logit,
            "node_slot_logits": node_slot_logits,
            "edge_logits": edge_logits,
            "edge_valid": edge_valid,
            "all_prop_masks": masks,
            "all_prop_valid": valid,
            "prior_signal": prior_signal,
        }


def transfer_v48_backbone(model: WaveSystemSetParserV5, state_dict: Dict[str, torch.Tensor]) -> Dict[str, float]:
    """Load shape-compatible v4.8 backbone tensors into V5."""
    target = model.state_dict()
    prefixes = ("stem.", "stem_norm.", "physics.", "blocks.")
    matched = {}
    for k, v in state_dict.items():
        kk = k
        for pfx in ("module.", "model.", "cnn.", "seg_model.", "seg."):
            if kk.startswith(pfx):
                kk = kk[len(pfx):]
        if kk.startswith(prefixes) and kk in target and tuple(v.shape) == tuple(target[kk].shape):
            matched[kk] = v
    model.load_state_dict(matched, strict=False)
    total_backbone = sum(v.numel() for k, v in target.items() if k.startswith(prefixes))
    loaded = sum(target[k].numel() for k in matched)
    return {
        "matched_tensors": len(matched),
        "backbone_numel_coverage": loaded / max(total_backbone, 1),
    }

# -----------------------------------------------------------------------------
# Archive/runtime helpers
# -----------------------------------------------------------------------------
def build_model(config: Optional[WaveSystemSetParserV5Config] = None) -> WaveSystemSetParserV5:
    """Construct the exact V5 segmentation architecture used by CNN.pt."""
    return WaveSystemSetParserV5(config or WaveSystemSetParserV5Config())


def load_cnn_checkpoint(
    checkpoint_path: str,
    map_location: str | torch.device = "cpu",
    strict: bool = True,
) -> WaveSystemSetParserV5:
    """
    Load the archived CNN.pt checkpoint.

    CNN.pt is expected to be the renamed V5 phase-C complete checkpoint:
      WaveSystemSetParserV5/checkpoints/phase_C/complete.pt
    """
    ck = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
    cfg_dict = ck.get("config", {}) if isinstance(ck, dict) else {}
    allowed = set(WaveSystemSetParserV5Config.__dataclass_fields__)
    cfg = WaveSystemSetParserV5Config(**{k: v for k, v in cfg_dict.items() if k in allowed})
    model = WaveSystemSetParserV5(cfg)

    if isinstance(ck, dict):
        state = ck.get("model", ck.get("state_dict", ck))
    else:
        state = ck
    model.load_state_dict(state, strict=strict)
    return model


__all__ = [
    "WaveSystemSetParserV5Config",
    "WaveSystemSetParserV5",
    "transfer_v48_backbone",
    "build_model",
    "load_cnn_checkpoint",
]