multi_mechanism_model.py

"""
Multi-Mechanism Normalizing Flow for Joint Mechanism Identification
and Bayesian Parameter Inference from Multi-Scan-Rate CV Signals.

Architecture:
    1. MultiScanEncoder: per-CV CNN + Set Transformer (SAB + PMA) -> context vector
    2. MechanismClassifier: MLP head -> p(mechanism | x)
    3. Per-mechanism ConditionalFlow heads: p(theta_m | x, mechanism=m)

The model performs two-level inference:
    Level 1: Mechanism probabilities p(m | x) from the classifier
    Level 2: Parameter posteriors p(theta_m | x, m) from mechanism-specific flows

Training loss = classification CE + mechanism-specific NLL (weighted by true label).
"""

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

from flow_model import (
    SignalEncoder,
    ActNorm,
    ConditionalSplineCoupling,
    ConditionalAffineCoupling,
    MECHANISM_LIST,
    MECHANISM_PARAMS,
)


class MechanismClassifier(nn.Module):
    """MLP classifier: context vector -> mechanism probabilities."""
    def __init__(self, d_context=128, n_mechanisms=4, hidden_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_context, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim // 2, n_mechanisms),
        )

    def forward(self, context):
        """Returns raw logits [B, n_mechanisms]."""
        return self.net(context)


class SAB(nn.Module):
    """Self-Attention Block (Set Transformer, Lee et al. 2019).

    Applies multi-head self-attention + feed-forward with residual
    connections and layer norm over a set of elements.
    """
    def __init__(self, d, n_heads=4):
        super().__init__()
        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(d)
        self.ffn = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
        self.norm2 = nn.LayerNorm(d)

    def forward(self, h, key_padding_mask=None):
        """
        Args:
            h: [B, N, d]
            key_padding_mask: [B, N] True = ignore this position
        """
        h2, _ = self.attn(h, h, h, key_padding_mask=key_padding_mask)
        h = self.norm1(h + h2)
        h = self.norm2(h + self.ffn(h))
        return h


class PMA(nn.Module):
    """Pooling by Multi-head Attention (Set Transformer, Lee et al. 2019).

    Uses learnable seed vectors as queries that attend to the input set,
    producing a fixed-size output regardless of set cardinality.
    """
    def __init__(self, d, n_heads=4, n_seeds=1):
        super().__init__()
        self.seed = nn.Parameter(torch.randn(1, n_seeds, d))
        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
        self.norm = nn.LayerNorm(d)

    def forward(self, h, key_padding_mask=None):
        """
        Args:
            h: [B, N, d]
            key_padding_mask: [B, N] True = ignore this position
        Returns:
            [B, n_seeds, d]
        """
        S = self.seed.expand(h.size(0), -1, -1)
        out, _ = self.attn(S, h, h, key_padding_mask=key_padding_mask)
        return self.norm(out + S)


class MultiScanEncoder(nn.Module):
    """
    Encode a set of multi-scan-rate CVs into a single context vector.

    Architecture (Set Transformer, default):
        1. Shared per-CV CNN encoder -> per-CV embedding
        2. Augment with [log10(sigma), log10(peak_flux)]
        3. SAB: self-attention across scan rates (cross-CV interaction)
        4. PMA: attention-based pooling to single vector
        5. rho MLP: project to final context

    With aggregation='mean_pool', steps 3-4 are replaced by masked mean
    pooling (no learned cross-CV interaction).

    Input:  x [B, N_sigma, 3, T], scan_mask [B, N_sigma, T],
            sigmas [B, N_sigma], flux_scales [B, N_sigma]
    Output: context [B, d_context]
    """
    def __init__(self, in_channels=3, d_model=128, d_context=128, n_heads=4,
                 aggregation='set_transformer'):
        super().__init__()
        self.aggregation = aggregation
        self.per_cv_encoder = SignalEncoder(
            in_channels=in_channels, d_model=d_model, d_context=d_context,
        )
        self.cv_augment = nn.Sequential(
            nn.Linear(d_context + 2, d_context),
            nn.GELU(),
        )
        if aggregation == 'set_transformer':
            self.sab = SAB(d_context, n_heads=n_heads)
            self.pma = PMA(d_context, n_heads=n_heads, n_seeds=1)
        elif aggregation == 'deepsets':
            self.phi = nn.Sequential(
                nn.Linear(d_context, d_context),
                nn.GELU(),
            )
        elif aggregation == 'mean_pool':
            pass
        else:
            raise ValueError(f"Unknown aggregation: {aggregation!r}")
        self.rho = nn.Sequential(
            nn.Linear(d_context, d_context),
            nn.GELU(),
            nn.Linear(d_context, d_context),
        )

    def forward(self, x, scan_mask=None, sigmas=None, flux_scales=None):
        """
        Args:
            x: [B, N_sigma, 3, T] multi-scan CV waveforms
            scan_mask: [B, N_sigma, T] valid timestep mask
            sigmas: [B, N_sigma] log10 scan rates
            flux_scales: [B, N_sigma] log10(peak_flux) per CV

        Returns:
            context: [B, d_context]
        """
        B, N, C, T = x.shape

        x_flat = x.reshape(B * N, C, T)
        mask_flat = scan_mask.reshape(B * N, T) if scan_mask is not None else None

        h_flat = self.per_cv_encoder(x_flat, mask=mask_flat)  # [B*N, d_context]
        h = h_flat.reshape(B, N, -1)  # [B, N, d_context]

        if sigmas is None:
            sigmas = torch.zeros(B, N, device=x.device)
        if flux_scales is None:
            flux_scales = torch.zeros(B, N, device=x.device)

        aug_features = torch.stack([sigmas, flux_scales], dim=-1)  # [B, N, 2]
        h = self.cv_augment(torch.cat([h, aug_features], dim=-1))  # [B, N, d_context]

        # Build key_padding_mask: True where CV is padded (invalid)
        if scan_mask is not None:
            cv_invalid = ~scan_mask.any(dim=-1)  # [B, N] True = padded
        else:
            cv_invalid = None

        if self.aggregation == 'set_transformer':
            h = self.sab(h, key_padding_mask=cv_invalid)
            h = self.pma(h, key_padding_mask=cv_invalid)  # [B, 1, d_context]
            h = h.squeeze(1)  # [B, d_context]
        elif self.aggregation in ('deepsets', 'mean_pool'):
            if self.aggregation == 'deepsets':
                h = self.phi(h)
            if cv_invalid is not None:
                cv_valid = (~cv_invalid).unsqueeze(-1).float()  # [B, N, 1]
                h = (h * cv_valid).sum(dim=1) / cv_valid.sum(dim=1).clamp(min=1)
            else:
                h = h.mean(dim=1)  # [B, d_context]

        context = self.rho(h)
        return context


class MechanismFlow(nn.Module):
    """
    Single-mechanism conditional flow: p(theta_m | context).

    Lightweight flow head operating on the context vector from the shared encoder.
    """
    def __init__(
        self,
        theta_dim,
        d_context=128,
        n_coupling_layers=6,
        hidden_dim=96,
        coupling_type='spline',
        n_bins=8,
        tail_bound=5.0,
    ):
        super().__init__()
        self.theta_dim = theta_dim
        self.coupling_type = coupling_type
        self.tail_bound = tail_bound

        self.flows = nn.ModuleList()
        for i in range(n_coupling_layers):
            mask_type = 'even' if i % 2 == 0 else 'odd'
            self.flows.append(ActNorm(theta_dim))
            if coupling_type == 'spline':
                self.flows.append(
                    ConditionalSplineCoupling(
                        dim=theta_dim,
                        d_context=d_context,
                        hidden_dim=hidden_dim,
                        mask_type=mask_type,
                        n_bins=n_bins,
                        tail_bound=tail_bound,
                    )
                )
            else:
                self.flows.append(
                    ConditionalAffineCoupling(
                        dim=theta_dim,
                        d_context=d_context,
                        hidden_dim=hidden_dim,
                        mask_type=mask_type,
                    )
                )

        self.register_buffer('theta_mean', torch.zeros(theta_dim))
        self.register_buffer('theta_std', torch.ones(theta_dim))

    def set_theta_stats(self, mean, std):
        self.theta_mean.copy_(torch.as_tensor(mean, dtype=torch.float32))
        self.theta_std.copy_(torch.as_tensor(std, dtype=torch.float32))

    def normalize_theta(self, theta):
        return (theta - self.theta_mean) / self.theta_std

    def denormalize_theta(self, theta_norm):
        return theta_norm * self.theta_std + self.theta_mean

    def forward_flow(self, z, context):
        total_log_det = torch.zeros(z.shape[0], device=z.device)
        h = z
        for layer in self.flows:
            if isinstance(layer, ActNorm):
                h, ld = layer(h)
                total_log_det += ld
            else:
                h, ld = layer(h, context)
                total_log_det += ld
        return h, total_log_det

    def inverse_flow(self, theta_norm, context):
        total_log_det = torch.zeros(theta_norm.shape[0], device=theta_norm.device)
        h = theta_norm
        for layer in reversed(self.flows):
            if isinstance(layer, ActNorm):
                h = layer.inverse(h)
                total_log_det -= layer.log_scale.sum()
            else:
                h, ld = layer.inverse(h, context)
                total_log_det += ld
        return h, total_log_det

    def log_prob(self, theta, context):
        """Compute log p(theta | context) for this mechanism's parameters."""
        theta_norm = self.normalize_theta(theta)
        if self.coupling_type == 'spline':
            theta_norm = theta_norm.clamp(-self.tail_bound, self.tail_bound)
        z, log_det = self.inverse_flow(theta_norm, context)
        log_pz = -0.5 * (z ** 2 + math.log(2 * math.pi)).sum(dim=-1)
        log_det_norm = -torch.log(self.theta_std).sum()
        log_p = log_pz + log_det + log_det_norm
        return log_p.clamp(min=-50.0, max=50.0)

    @torch.no_grad()
    def sample(self, context, n_samples=100, temperature=1.0):
        """Sample theta from this mechanism's posterior.

        Args:
            temperature: Posterior inflation factor. Can be:
                - scalar: uniform scaling for all parameters
                - 1-D tensor of shape [theta_dim]: per-parameter scaling
                T > 1 broadens the posterior by scaling samples around
                their per-example mean in theta-space.
        """
        B = context.shape[0]
        context_rep = context.unsqueeze(1).expand(-1, n_samples, -1).reshape(B * n_samples, -1)
        z = torch.randn(B * n_samples, self.theta_dim, device=context.device)
        theta_norm, _ = self.forward_flow(z, context_rep)
        theta = self.denormalize_theta(theta_norm)
        theta = theta.reshape(B, n_samples, self.theta_dim)

        if isinstance(temperature, torch.Tensor):
            T = temperature.to(theta.device).reshape(1, 1, -1)
            mu = theta.mean(dim=1, keepdim=True)
            theta = mu + T * (theta - mu)
        elif temperature != 1.0:
            mu = theta.mean(dim=1, keepdim=True)
            theta = mu + temperature * (theta - mu)

        return theta

    def sample_with_grad(self, context, n_samples=64):
        """Sample theta with gradients enabled (for calibration loss).

        Uses the reparameterization trick: z ~ N(0,I) is fixed noise,
        gradients flow through the flow's forward transform.
        """
        B = context.shape[0]
        context_rep = context.unsqueeze(1).expand(-1, n_samples, -1).reshape(B * n_samples, -1)
        z = torch.randn(B * n_samples, self.theta_dim, device=context.device)
        theta_norm, _ = self.forward_flow(z, context_rep)
        theta = self.denormalize_theta(theta_norm)
        return theta.reshape(B, n_samples, self.theta_dim)


SUMMARY_DIM = 21  # 3 scan rates * 6 features + 3 log10(sigma)


class SummaryProjection(nn.Module):
    """Project hand-crafted summary statistics to context space.

    Replaces the Set Transformer encoder for the Summary-ECFlow ablation,
    keeping everything else (classifier, flow heads, training) identical.
    """
    def __init__(self, summary_dim=SUMMARY_DIM, d_context=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(summary_dim, d_context),
            nn.GELU(),
            nn.Linear(d_context, d_context),
            nn.GELU(),
            nn.Linear(d_context, d_context),
        )

    def forward(self, summary):
        """summary: [B, summary_dim] -> [B, d_context]"""
        return self.net(summary)


class MultiMechanismFlow(nn.Module):
    """
    Joint mechanism identification and parameter inference model.

    Combines:
        - Multi-scan-rate signal encoder (Set Transformer over per-CV embeddings)
        - Mechanism classifier
        - Per-mechanism normalizing flow heads

    If use_summary_features=True, replaces the signal encoder with a simple
    MLP projection from hand-crafted summary statistics (21-dim) to context
    space, keeping all other components identical. This isolates the effect
    of the input representation (full signal vs. summary statistics).
    """
    def __init__(
        self,
        d_context=128,
        d_model=128,
        n_coupling_layers=6,
        hidden_dim=96,
        coupling_type='spline',
        n_bins=8,
        tail_bound=5.0,
        aggregation='set_transformer',
        use_summary_features=False,
    ):
        super().__init__()
        self.n_mechanisms = len(MECHANISM_LIST)
        self.mechanism_list = MECHANISM_LIST
        self.d_context = d_context
        self.use_summary_features = use_summary_features

        if use_summary_features:
            self.summary_proj = SummaryProjection(
                summary_dim=SUMMARY_DIM, d_context=d_context,
            )
            self.encoder = None
        else:
            self.encoder = MultiScanEncoder(
                in_channels=3, d_model=d_model, d_context=d_context,
                aggregation=aggregation,
            )
            self.summary_proj = None

        self.classifier = MechanismClassifier(
            d_context=d_context,
            n_mechanisms=self.n_mechanisms,
            hidden_dim=hidden_dim,
        )

        # Mechanisms whose parameters are not identifiable from the signal
        # (e.g. Nernst: E0_offset=0, dA=1 are constants, dB is unidentifiable).
        # These get a flow head for architecture uniformity but are excluded
        # from the NLL loss during training.
        self.skip_nll_mechanisms: set = set()

        self.flow_heads = nn.ModuleDict()
        for mech in MECHANISM_LIST:
            theta_dim = MECHANISM_PARAMS[mech]['dim']
            self.flow_heads[mech] = MechanismFlow(
                theta_dim=theta_dim,
                d_context=d_context,
                n_coupling_layers=n_coupling_layers,
                hidden_dim=hidden_dim,
                coupling_type=coupling_type,
                n_bins=n_bins,
                tail_bound=tail_bound,
            )

    def set_theta_stats(self, mechanism, mean, std):
        """Set normalization stats for a specific mechanism's flow head."""
        self.flow_heads[mechanism].set_theta_stats(mean, std)

    def encode_signal(self, x, scan_mask=None, sigmas=None, flux_scales=None,
                       summary=None):
        """
        Encode input into a context vector.

        In full-signal mode: uses Set Transformer encoder on x.
        In summary mode: projects 21-dim hand-crafted features to context space.

        Args:
            x: [B, N_sigma, 3, T] (ignored in summary mode)
            scan_mask: [B, N_sigma, T]
            sigmas: [B, N_sigma] log10 scan rates
            flux_scales: [B, N_sigma] log10(std(flux)) per CV
            summary: [B, 21] hand-crafted summary statistics (summary mode only)
        """
        if self.use_summary_features:
            assert summary is not None, "summary features required in summary mode"
            return self.summary_proj(summary)
        return self.encoder(x, scan_mask=scan_mask, sigmas=sigmas,
                            flux_scales=flux_scales)

    def forward(self, x, mechanism_ids, mech_theta, mech_theta_mask=None,
                scan_mask=None, sigmas=None, flux_scales=None, summary=None):
        """
        Compute classification logits and per-sample NLL for the true mechanism.

        Returns:
            dict with 'logits' [B, n_mechanisms] and 'nll' [B]
        """
        context = self.encode_signal(x, scan_mask=scan_mask, sigmas=sigmas,
                                     flux_scales=flux_scales, summary=summary)
        logits = self.classifier(context)

        nll = torch.zeros(x.shape[0], device=x.device)

        for m_idx, mech in enumerate(MECHANISM_LIST):
            sel = (mechanism_ids == m_idx)
            if not sel.any():
                continue
            if mech in self.skip_nll_mechanisms:
                continue

            theta_dim = MECHANISM_PARAMS[mech]['dim']
            ctx_m = context[sel]
            theta_m = mech_theta[sel, :theta_dim]

            log_p = self.flow_heads[mech].log_prob(theta_m, ctx_m)
            bad = ~torch.isfinite(log_p)
            if bad.any():
                log_p = torch.where(bad, torch.full_like(log_p, -10.0).detach(), log_p)
            nll[sel] = -log_p

        return {'logits': logits, 'nll': nll}

    def forward_with_calibration(self, x, mechanism_ids, mech_theta,
                                 mech_theta_mask=None, scan_mask=None,
                                 sigmas=None, flux_scales=None,
                                 cal_n_samples=64, cal_levels=(0.5, 0.9),
                                 cal_beta=20.0, summary=None):
        """
        Forward pass with additional calibration loss.

        Extends forward() by drawing posterior samples and computing a
        differentiable coverage penalty that encourages the flow's credible
        intervals to match their nominal levels.

        Returns:
            dict with 'logits', 'nll', and 'cal_loss' (scalar)
        """
        context = self.encode_signal(x, scan_mask=scan_mask, sigmas=sigmas,
                                     flux_scales=flux_scales, summary=summary)
        logits = self.classifier(context)

        nll = torch.zeros(x.shape[0], device=x.device)
        cal_losses = []

        for m_idx, mech in enumerate(MECHANISM_LIST):
            sel = (mechanism_ids == m_idx)
            if not sel.any():
                continue
            if mech in self.skip_nll_mechanisms:
                continue

            theta_dim = MECHANISM_PARAMS[mech]['dim']
            ctx_m = context[sel]
            theta_m = mech_theta[sel, :theta_dim]

            # Standard NLL
            log_p = self.flow_heads[mech].log_prob(theta_m, ctx_m)
            bad = ~torch.isfinite(log_p)
            if bad.any():
                log_p = torch.where(bad, torch.full_like(log_p, -10.0).detach(), log_p)
            nll[sel] = -log_p

            # Calibration: draw samples WITH gradients
            if ctx_m.shape[0] < 4:
                continue
            samples = self.flow_heads[mech].sample_with_grad(
                ctx_m, n_samples=cal_n_samples,
            )  # [B_m, K, D]

            # Inverse-spread weights: collapsed parameters (small posterior std)
            # get more calibration pressure than well-spread parameters.
            with torch.no_grad():
                param_std = samples.std(dim=1).clamp(min=1e-4)  # [B_m, D]
                inv_spread_w = 1.0 / param_std                  # [B_m, D]
                inv_spread_w = inv_spread_w / inv_spread_w.mean()

            for level in cal_levels:
                alpha = (1.0 - level) / 2.0
                lower = torch.quantile(samples, alpha, dim=1)       # [B_m, D]
                upper = torch.quantile(samples, 1 - alpha, dim=1)   # [B_m, D]

                inside = (
                    torch.sigmoid(cal_beta * (theta_m - lower))
                    * torch.sigmoid(cal_beta * (upper - theta_m))
                )  # [B_m, D]

                per_sample_loss = (inside - level).pow(2)  # [B_m, D]
                cal_losses.append((per_sample_loss * inv_spread_w).mean())

        if cal_losses:
            cal_loss = torch.stack(cal_losses).mean()
        else:
            cal_loss = torch.tensor(0.0, device=x.device)

        return {'logits': logits, 'nll': nll, 'cal_loss': cal_loss}

    @torch.no_grad()
    def predict(self, x, scan_mask=None, sigmas=None, flux_scales=None,
                n_samples=200, top_k=None, temperature=1.0,
                temperature_map=None, summary=None):
        """
        Full inference: classify mechanism, then sample parameters.

        Args:
            x: [B, N_sigma, 3, T] multi-scan input signals
            scan_mask: [B, N_sigma, T]
            sigmas: [B, N_sigma]
            flux_scales: [B, N_sigma]
            n_samples: posterior samples per mechanism
            top_k: if set, only sample from top-k most likely mechanisms
            temperature: scalar fallback (>1 broadens posteriors)
            temperature_map: dict mapping mechanism name -> list of
                per-parameter temperatures. Overrides scalar temperature
                for mechanisms present in the map.
            summary: [B, 21] hand-crafted summary stats (summary mode only)

        Returns:
            dict with mechanism_probs, mechanism_pred, samples, stats
        """
        context = self.encode_signal(x, scan_mask=scan_mask, sigmas=sigmas,
                                     flux_scales=flux_scales, summary=summary)
        logits = self.classifier(context)
        probs = F.softmax(logits, dim=-1)
        pred = probs.argmax(dim=-1)

        probs_clamped = probs.clamp(min=1e-7, max=1 - 1e-7)
        xdB = 10.0 * torch.log10(probs_clamped / (1.0 - probs_clamped))

        samples_dict = {}
        stats_dict = {}

        for m_idx, mech in enumerate(MECHANISM_LIST):
            if top_k is not None:
                top_k_mechs = probs.topk(top_k, dim=-1).indices
                if not (top_k_mechs == m_idx).any():
                    samples_dict[mech] = None
                    stats_dict[mech] = None
                    continue

            T = temperature
            if temperature_map is not None and mech in temperature_map:
                T = torch.tensor(temperature_map[mech], dtype=torch.float32)

            s = self.flow_heads[mech].sample(context, n_samples=n_samples,
                                             temperature=T)
            samples_dict[mech] = s
            stats_dict[mech] = {
                'mean': s.mean(dim=1),
                'std': s.std(dim=1),
                'median': s.median(dim=1).values,
                'q05': s.quantile(0.05, dim=1),
                'q95': s.quantile(0.95, dim=1),
            }

        return {
            'mechanism_probs': probs,
            'mechanism_xdB': xdB,
            'mechanism_pred': pred,
            'samples': samples_dict,
            'stats': stats_dict,
        }

    @torch.no_grad()
    def predict_single_mechanism(self, x, mechanism, scan_mask=None,
                                 sigmas=None, flux_scales=None, n_samples=1000,
                                 temperature=1.0, temperature_map=None,
                                 summary=None):
        """Sample parameters assuming a known mechanism."""
        context = self.encode_signal(x, scan_mask=scan_mask, sigmas=sigmas,
                                     flux_scales=flux_scales, summary=summary)
        T = temperature
        if temperature_map is not None and mechanism in temperature_map:
            T = torch.tensor(temperature_map[mechanism], dtype=torch.float32)
        samples = self.flow_heads[mechanism].sample(context, n_samples=n_samples,
                                                    temperature=T)
        return {
            'mean': samples.mean(dim=1),
            'std': samples.std(dim=1),
            'median': samples.median(dim=1).values,
            'q05': samples.quantile(0.05, dim=1),
            'q95': samples.quantile(0.95, dim=1),
            'samples': samples,
        }


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


if __name__ == "__main__":
    n_mechs = len(MECHANISM_LIST)
    B, N_sigma, T = n_mechs, 3, 800

    x = torch.randn(B, N_sigma, 3, T)
    scan_mask = torch.ones(B, N_sigma, T, dtype=torch.bool)
    sigmas = torch.randn(B, N_sigma)
    flux_scales = torch.randn(B, N_sigma)
    mechanism_ids = torch.arange(n_mechs)

    max_dim = max(MECHANISM_PARAMS[m]['dim'] for m in MECHANISM_LIST)
    mech_theta = torch.randn(B, max_dim)
    mech_theta_mask = torch.zeros(B, max_dim, dtype=torch.bool)
    for i, mid in enumerate(mechanism_ids):
        d = MECHANISM_PARAMS[MECHANISM_LIST[mid]]['dim']
        mech_theta_mask[i, :d] = True

    print("=" * 60)
    print("Testing MultiMechanismFlow (multi-scan-rate, Set Transformer)")
    print("=" * 60)

    model = MultiMechanismFlow(
        d_context=128,
        d_model=128,
        n_coupling_layers=8,
        hidden_dim=128,
        coupling_type='affine',
    )

    total_params = count_parameters(model)
    print(f"Total parameters: {total_params:,}")
    print(f"  Encoder: {count_parameters(model.encoder):,}")
    print(f"  Classifier: {count_parameters(model.classifier):,}")
    for mech in MECHANISM_LIST:
        print(f"  Flow ({mech}, dim={MECHANISM_PARAMS[mech]['dim']}): "
              f"{count_parameters(model.flow_heads[mech]):,}")

    out = model(x, mechanism_ids, mech_theta, mech_theta_mask,
                scan_mask=scan_mask, sigmas=sigmas, flux_scales=flux_scales)
    print(f"\nForward pass:")
    print(f"  Logits shape: {out['logits'].shape}")
    print(f"  NLL shape: {out['nll'].shape}")
    print(f"  NLL values: {out['nll']}")

    pred = model.predict(x, scan_mask=scan_mask, sigmas=sigmas,
                         flux_scales=flux_scales, n_samples=100)
    print(f"\nPrediction:")
    print(f"  Mechanism probs shape: {pred['mechanism_probs'].shape}")
    print(f"  Mechanism xdB shape: {pred['mechanism_xdB'].shape}")
    print(f"  Predicted mechanisms: {pred['mechanism_pred']}")
    print(f"  xdB for predicted mechanisms:")
    for i in range(B):
        m_idx = pred['mechanism_pred'][i].item()
        mech = MECHANISM_LIST[m_idx]
        xdB_val = pred['mechanism_xdB'][i, m_idx].item()
        print(f"    sample {i}: {mech} ({xdB_val:+.1f} dB)")
    for mech in MECHANISM_LIST:
        if pred['samples'][mech] is not None:
            print(f"  {mech} samples shape: {pred['samples'][mech].shape}")