"""
WaveSystemGraphParser v4.7b protected real+synthetic core-locked: core-attached support topology graph parser for ICWDS
===================================================================
Purpose
-------
This model is the next-stage frontend after WaveSystemParser v3.x.
It does not ask a CNN to draw wave-system masks from scratch. Instead:

    E(f,theta)
      -> physical peak/basin proposals generated outside the model
      -> node/edge graph reasoning over proposals
      -> learned merge / keep / slot assembly
      -> light CNN boundary refinement

This treats watershed-like basins as over-segmentation proposals, not labels.
The learnable part decides which candidates are physical wave systems, which
should be merged, and which should be sent to the downstream self-pruning VAE.

v4.5 is designed for a double-layer physical teacher: peak-core proposals define
system identity/count, while valley-constrained support proposals recover the full
energetic wave-system footprint. Stripe-like bands can be attached as tails but
are not allowed to become independent systems without a peak core.

All operations are lightweight and Colab-friendly. No Transformer blocks.
"""
import math
from dataclasses import dataclass, asdict
from typing import Optional, Dict

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


@dataclass
class GraphParserV47Config:
    n_freqs: int = 47
    n_dirs: int = 72
    n_slots: int = 6
    bg_index: int = 6
    p_max: int = 18
    prop_feat_dim: int = 22
    width: int = 32
    depth: int = 4
    node_dim: int = 48
    edge_dim: int = 64
    pair_feat_dim: int = 8
    use_coord: bool = True
    use_physics: bool = True
    rank_temp: float = 0.12
    count_min: float = 1.0
    count_max: float = 6.0
    bg_prior_bias: float = 0.38
    bg_energy_suppress: float = 14.5
    bg_energy_gamma: float = 0.70
    proposal_logit_gain: float = 5.25

    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__()
        self.n_dirs = n_dirs
        dirs = torch.linspace(0, 2 * math.pi, n_dirs + 1)[: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 WaveSystemGraphParserV47(nn.Module):
    """Proposal graph parser.

    forward inputs
    --------------
    x:            [B,1,47,72], normalized to [-1,1]
    prop_masks:   [B,P,47,72], binary/soft physical basin proposals
    prop_feats:   [B,P,F], proposal features produced by the training script
    prop_valid:   [B,P], 1 if proposal exists

    outputs include final probability masks, node slot assignments, edge logits,
    count prediction, and diagnostics.
    """
    def __init__(self, cfg: Optional[GraphParserV47Config] = None):
        super().__init__()
        self.cfg = cfg or GraphParserV47Config()
        c = self.cfg
        in_ch = 1 + (3 if c.use_coord else 0)
        self.stem = nn.Conv2d(in_ch, c.width, 3, padding=1)
        self.stem_norm = nn.GroupNorm(min(8, c.width), c.width)
        self.physics = PhysicsAwareModule(c.width, c.n_dirs) if c.use_physics else None
        self.blocks = nn.ModuleList([DepthwiseSeparable(c.width, c.width) for _ in range(c.depth)])
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.residual_head = nn.Conv2d(c.width, c.n_slots + 1, 1)
        nn.init.zeros_(self.residual_head.weight)
        nn.init.zeros_(self.residual_head.bias)

        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.node_slot = nn.Linear(c.node_dim, c.n_slots)
        self.count_head = nn.Sequential(
            nn.Linear(c.width + c.node_dim, c.width), nn.GELU(), nn.Linear(c.width, c.n_slots)
        )
        # Edge head uses node_i, node_j, absolute difference, product, plus handcrafted physical pair features.
        self.edge_head = nn.Sequential(
            nn.Linear(4 * c.node_dim + c.pair_feat_dim, c.edge_dim), nn.GELU(),
            nn.Linear(c.edge_dim, c.edge_dim), nn.GELU(), nn.Linear(c.edge_dim, 1)
        )
        self._coord_cache = None


    @staticmethod
    def _clean_tensor(x, fill=0.0, lo=-30.0, hi=30.0):
        return torch.nan_to_num(x, nan=fill, posinf=hi, neginf=lo).clamp(lo, hi)

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

    def _coord_channels(self, B, device, dtype):
        c = self.cfg
        if self._coord_cache is None:
            nf, nd = c.n_freqs, c.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)

    def _rank_prune_presence(self, slot_mass, count_soft):
        c = self.cfg
        B, K = slot_mass.shape
        _, order = torch.sort(slot_mass, dim=1, descending=True)
        rank_pos = torch.arange(1, K + 1, device=slot_mass.device, dtype=slot_mass.dtype).view(1, K)
        cs = count_soft.clamp(c.count_min, c.count_max).view(B, 1)
        gate_sorted = torch.sigmoid((cs + 0.5 - rank_pos) / c.rank_temp)
        gate = torch.zeros_like(gate_sorted).scatter(1, order, gate_sorted)
        return gate.clamp(0, 1)

    def _node_pool(self, h, prop_masks, prop_valid):
        # h: [B,C,H,W], prop_masks [B,P,H,W]
        B, C, H, W = h.shape
        P = prop_masks.shape[1]
        denom = prop_masks.flatten(2).sum(dim=2).clamp_min(1.0)  # [B,P]
        pooled = torch.einsum("bchw,bphw->bpc", h, prop_masks) / denom[:, :, None]
        pooled = pooled * prop_valid[:, :, None]
        return pooled

    def _edge_logits(self, node, prop_feats, prop_valid):
        B, P, D = node.shape
        ni = node[:, :, None, :].expand(B, P, P, D)
        nj = node[:, None, :, :].expand(B, P, P, D)
        # Pair physical features from proposal features: distance in f/theta and mass contrast.
        # feat layout is defined in training script: mass, peak, area, mu_f, sin_t, cos_t, ...
        fi = prop_feats[:, :, 3][:, :, None]
        fj = prop_feats[:, :, 3][:, None, :]
        d_f = (fi - fj).abs()
        si = prop_feats[:, :, 4][:, :, None]; ci = prop_feats[:, :, 5][:, :, None]
        sj = prop_feats[:, :, 4][:, None, :]; cj = prop_feats[:, :, 5][:, None, :]
        dot = (si * sj + ci * cj).clamp(-1.0 + 1e-5, 1.0 - 1e-5)
        d_t = torch.acos(dot) / math.pi
        mi = prop_feats[:, :, 0][:, :, None]
        mj = prop_feats[:, :, 0][:, None, :]
        d_m = (mi - mj).abs()
        sf_i = prop_feats[:, :, 6][:, :, None]; sf_j = prop_feats[:, :, 6][:, None, :]
        st_i = prop_feats[:, :, 7][:, :, None]; st_j = prop_feats[:, :, 7][:, None, :]
        d_sf = (sf_i - sf_j).abs()
        d_st = (st_i - st_j).abs()
        prom_i = prop_feats[:, :, 12][:, :, None]; prom_j = prop_feats[:, :, 12][:, None, :]
        prom_min = torch.minimum(prom_i, prom_j)
        stripe_i = prop_feats[:, :, 13][:, :, None]; stripe_j = prop_feats[:, :, 13][:, None, :]
        stripe_max = torch.maximum(stripe_i, stripe_j)
        qual_i = prop_feats[:, :, 11][:, :, None]; qual_j = prop_feats[:, :, 11][:, None, :]
        qual_min = torch.minimum(qual_i, qual_j)
        pair_phys = torch.stack([d_f, d_t, d_m, d_sf, d_st, prom_min, stripe_max, qual_min], dim=-1)
        inp = torch.cat([ni, nj, (ni - nj).abs(), ni * nj, pair_phys], dim=-1)
        e = self.edge_head(inp).squeeze(-1)
        valid_pair = (prop_valid[:, :, None] * prop_valid[:, None, :]).bool()
        eye = torch.eye(P, device=node.device, dtype=torch.bool).view(1, P, P)
        e = e.masked_fill(~valid_pair | eye, 0.0)
        return e, valid_pair & (~eye)

    def forward(self, x, prop_masks, prop_feats, prop_valid, prior_w=2.5, residual_w=0.0):
        c = self.cfg
        B, _, H, W = x.shape
        prop_masks = torch.nan_to_num(prop_masks.float(), nan=0.0, posinf=0.0, neginf=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.0, 5.0)
        prop_valid = prop_valid.float().clamp(0, 1)
        h_in = torch.cat([x, self._coord_channels(B, x.device, x.dtype)], dim=1) if c.use_coord else x
        h = self._clean_tensor(F.gelu(self.stem_norm(self.stem(h_in))), lo=-20.0, hi=20.0)
        E01 = torch.nan_to_num((x + 1.0) * 0.5, nan=0.0, posinf=1.0, neginf=0.0).clamp(0, 1)
        if self.physics is not None:
            h = self._clean_tensor(h + self.physics(h, E01), lo=-20.0, hi=20.0)
        for blk in self.blocks:
            h = self._clean_tensor(h + blk(h), lo=-20.0, hi=20.0)
        global_feat = self.global_pool(h).flatten(1)
        node_pool = self._node_pool(h, prop_masks, prop_valid)
        node_in = torch.cat([node_pool, prop_feats], dim=-1)
        node = self._clean_tensor(self.node_mlp(node_in), lo=-20.0, hi=20.0) * prop_valid[:, :, None]
        node_keep_logit = self._clean_tensor(self.node_keep(node).squeeze(-1), lo=-20.0, hi=20.0).masked_fill(prop_valid <= 0, -20.0)
        node_keep = torch.sigmoid(node_keep_logit) * prop_valid
        node_slot_logits = self._clean_tensor(self.node_slot(node), lo=-20.0, hi=20.0).masked_fill(prop_valid[:, :, None] <= 0, -20.0)
        node_slot = self._safe_softmax(node_slot_logits, dim=-1) * prop_valid[:, :, None]
        node_context = (node * node_keep[:, :, None]).sum(dim=1) / node_keep.sum(dim=1, keepdim=True).clamp_min(1.0)
        count_logits = self._clean_tensor(self.count_head(torch.cat([global_feat, node_context], dim=-1)), lo=-20.0, hi=20.0)
        count_probs = torch.sigmoid(count_logits)
        count_soft = count_probs.sum(dim=1).clamp(c.count_min, c.count_max)

        # Proposal graph edge logits.
        edge_logits, edge_valid = self._edge_logits(node, prop_feats, prop_valid)
        edge_logits = self._clean_tensor(edge_logits, lo=-20.0, hi=20.0)
        # Slot priors from proposal assembly.
        assign = node_slot * node_keep[:, :, None]
        prior_signal = torch.einsum("bpk,bphw->bkhw", assign, prop_masks)
        # Normalize each slot prior but preserve zero slots.
        prior_signal = prior_signal / prior_signal.amax(dim=(2, 3), keepdim=True).clamp_min(1e-6)
        prior_signal = torch.nan_to_num(prior_signal, nan=0.0, posinf=1.0, neginf=0.0).clamp(0.0, 1.0)
        denom = E01.sum(dim=(2, 3)).clamp_min(1e-6)
        slot_mass = (prior_signal * E01).sum(dim=(2, 3)) / denom
        presence = self._rank_prune_presence(slot_mass, count_soft)
        prior_signal = prior_signal * presence[:, :, None, None]

        En = E01 / E01.amax(dim=(2, 3), keepdim=True).clamp_min(1e-6)
        bg_prior_logit = c.bg_prior_bias - c.bg_energy_suppress * En.pow(c.bg_energy_gamma)
        residual_logits = self.residual_head(h)
        signal_logits = prior_w * (c.proposal_logit_gain * torch.log(prior_signal.clamp_min(1e-6))) + residual_w * residual_logits[:, :c.n_slots]
        bg_logits = prior_w * bg_prior_logit + residual_w * residual_logits[:, c.n_slots:c.n_slots + 1]
        logits = self._clean_tensor(torch.cat([signal_logits, bg_logits], dim=1), lo=-60.0, hi=60.0)
        prob_raw = self._safe_softmax(logits, dim=1)
        prob = prob_raw / prob_raw.sum(dim=1, keepdim=True).clamp_min(1e-6)

        prior_logits = self._clean_tensor(torch.cat([c.proposal_logit_gain * torch.log(prior_signal.clamp_min(1e-6)), bg_prior_logit], dim=1), lo=-60.0, hi=60.0)
        prior_prob = self._safe_softmax(prior_logits, dim=1)
        return {
            "logits": logits, "prob": prob, "prob_raw": prob_raw,
            "prior_signal": prior_signal, "prior_logits": prior_logits, "prior_prob": prior_prob,
            "residual_logits": residual_logits,
            "node": node, "node_keep_logit": node_keep_logit, "node_keep": node_keep,
            "node_slot_logits": node_slot_logits, "node_slot": node_slot,
            "edge_logits": edge_logits, "edge_valid": edge_valid,
            "count_logits": count_logits, "count_probs": count_probs, "count_soft": count_soft,
            "slot_mass": slot_mass, "presence": presence,
        }

    @torch.no_grad()
    def prior_foreground_metrics(self, out, E01, high_quantile=0.80):
        E = E01[:, 0] if E01.dim() == 4 else E01
        flat = E.flatten(1)
        thr = torch.quantile(flat, high_quantile, dim=1).view(-1, 1, 1)
        mask_hi = E >= thr
        denom = mask_hi.float().sum().clamp_min(1.0)
        prior_prob = out.get("prior_prob", torch.softmax(out["prior_logits"], dim=1))
        fg_prior = prior_prob[:, :self.cfg.n_slots].sum(dim=1)
        bg_prior = prior_prob[:, self.cfg.bg_index]
        return {"prior_fg_hi": (fg_prior * mask_hi).sum() / denom,
                "prior_bg_hi": (bg_prior * mask_hi).sum() / denom}

    @torch.no_grad()
    def dominance_metrics(self, out, E01, high_quantile=0.80):
        final = out["prob"].argmax(dim=1)
        prior = out["prior_logits"].argmax(dim=1)
        resid = out["residual_logits"].argmax(dim=1)
        E = E01[:, 0] if E01.dim() == 4 else E01
        flat = E.flatten(1)
        thr = torch.quantile(flat, high_quantile, dim=1).view(-1, 1, 1)
        mask = E >= thr
        denom = mask.float().sum().clamp_min(1.0)
        return {"prior_agree": ((final == prior) & mask).float().sum() / denom,
                "residual_agree": ((final == resid) & mask).float().sum() / denom}

    def graph_regularizers(self, out, prop_masks, prop_valid, E01):
        # Slot compactness and smoothness proxies for final masks.
        c = self.cfg
        prob = out["prob"][:, :c.n_slots]
        E = E01[:, 0] if E01.dim() == 4 else E01
        nf, nd = c.n_freqs, c.n_dirs
        f = torch.linspace(0, 1, nf, device=E.device, dtype=E.dtype).view(1, 1, nf, 1)
        theta = torch.linspace(0, 2 * math.pi, nd + 1, device=E.device, dtype=E.dtype)[:nd].view(1, 1, 1, nd)
        w = prob * E[:, None]
        mass = w.sum(dim=(2, 3)).clamp_min(1e-8)
        mu_f = (w * f).sum(dim=(2, 3)) / mass
        cx = (w * torch.cos(theta)).sum(dim=(2, 3)) / mass
        cy = (w * torch.sin(theta)).sum(dim=(2, 3)) / mass
        mu_t = torch.atan2(cy, cx)
        df2 = (f - mu_f[:, :, None, None]) ** 2
        dt = torch.atan2(torch.sin(theta - mu_t[:, :, None, None]), torch.cos(theta - mu_t[:, :, None, None])) / math.pi
        radius = ((w * (df2 + dt ** 2)).sum(dim=(2, 3)) / mass).mean()
        tv = (prob[:, :, 1:, :] - prob[:, :, :-1, :]).abs().mean() + (prob[:, :, :, 1:] - prob[:, :, :, :-1]).abs().mean()
        # Encourage node slot assignments to be confident only for valid proposals.
        ns = out["node_slot"].clamp_min(1e-8)
        ent = -(ns * ns.log()).sum(dim=-1)
        ent = (ent * prop_valid).sum() / prop_valid.sum().clamp_min(1.0)
        return {"slot_radius": radius, "slot_tv": tv, "node_slot_entropy": ent}

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


# Backward-compatible alias
WaveSystemGraphParserV4 = WaveSystemGraphParserV47

# v4.8 stable aliases: keep the proven v47b architecture/checkpoint compatibility.
# The v4.8 fixes are implemented in the training objective and deterministic
# energy/noise-aware post-processing, not by changing tensor shapes.
GraphParserV48Config = GraphParserV47Config
WaveSystemGraphParserV48 = WaveSystemGraphParserV47