"""
models.py — All neural-network model implementations for domain adaptation.

Classes
-------
BaseModel        : shared encoder backbone (all conditions)
DANNModel        : encoder + label classifier + domain discriminator via GRL
CORALModel       : encoder + label classifier + CORAL covariance alignment
RamanSelfSupModel: encoder + label classifier + masked-spectrum reconstruction

Mathematical core
-----------------
DANN     : Gradient Reversal Layer — forward identity, backward = -λ·∇
CORAL    : L_coral = || Cov_s − Cov_t ||_F²  (Frobenius norm)
SelfSup  : L_recon = MSE(recon(x_masked), x_original)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from experiment_config import ExperimentConfig

# ── Gradient Reversal Layer ───────────────────────────────────────────────────

class GradientReversalFunction(torch.autograd.Function):
    """
    Custom gradient operator for DANN.

    Forward : y = x                         (identity)
    Backward: dy/dx = -λ · dy/dx_out        (sign reversal)

    References: Ganin & Lempitsky, JMLR 2016 — "Unsupervised Domain Adaptation
    by Backpropagation"
    """

    @staticmethod
    def forward(ctx, x, lambda_):
        ctx.lambda_ = lambda_
        return x.view_as(x)

    @staticmethod
    def backward(ctx, grad_output):
        return -ctx.lambda_ * grad_output, None

class GradientReversalLayer(nn.Module):
    """Wraps GradientReversalFunction as a differentiable nn.Module."""

    def __init__(self, lambda_=1.0):
        super().__init__()
        self.lambda_ = lambda_

    def forward(self, x):
        return GradientReversalFunction.apply(x, self.lambda_)

# ── Base model (shared encoder backbone) ────────────────────────────────────

class BaseModel(nn.Module):
    """
    Shared encoder + linear classifier used by ALL conditions.
    Architecture: small ResNet-style conv stack → bottleneck → classifier.

    Input : [B, 3, 32, 32] RGB image
    Output: [B, num_classes] logits
    """

    def __init__(self, config):
        super().__init__()
        self.cfg = config

        # ── Conv encoder ──────────────────────────────────────────
        self.encoder = nn.Sequential(
            # Block 1: 32×32 → 16×16
            nn.Conv2d(3, 64, 3, 1, 1), nn.BatchNorm2d(64), nn.ReLU(),
            nn.Conv2d(64, 128, 3, 1, 1), nn.BatchNorm2d(128), nn.ReLU(),
            nn.MaxPool2d(2, 2),
            # Block 2: 16×16 → 8×8
            nn.Conv2d(128, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.ReLU(),
            nn.Conv2d(256, 512, 3, 1, 1), nn.BatchNorm2d(512), nn.ReLU(),
            nn.MaxPool2d(2, 2),
            # Block 3: 8×8 → 4×4
            nn.Conv2d(512, 512, 3, 1, 1), nn.BatchNorm2d(512), nn.ReLU(),
            nn.MaxPool2d(2, 2),
            # Global pooling
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )

        # ── Bottleneck ────────────────────────────────────────────
        self.bottleneck = nn.Sequential(
            nn.Linear(512, config.hidden_dim),
            nn.BatchNorm1d(config.hidden_dim),
            nn.ReLU(),
        )

        # ── Classifier ────────────────────────────────────────────
        self.classifier = nn.Linear(config.hidden_dim, config.num_classes)

    def encode(self, x):
        """Returns bottleneck representation z ∈ ℝ^{hidden_dim}."""
        feat = self.encoder(x)        # [B, 512]
        z    = self.bottleneck(feat)  # [B, hidden_dim]
        return z

    def forward(self, x):
        z      = self.encode(x)
        logits = self.classifier(z)
        return logits

# ── DANN Model ─────────────────────────────────────────────────────────────────

class DANNModel(BaseModel):
    """
    Domain-Adversarial Neural Network (Ganin & Lempitsky, JMLR 2016).

    Adds a domain-discriminator head connected to the bottleneck via GRL.
    The encoder is trained to:
      (a) minimise label-prediction loss  L_CE
      (b) fool the domain discriminator   L_adv
    This produces domain-invariant feature representations without target labels.

    STRATEGY: adversarial gradient reversal on domain features.
    """

    def __init__(self, config):
        super().__init__(config)
        self.grl = GradientReversalLayer(lambda_=config.lambda_domain)
        self.domain_discriminator = nn.Sequential(
            nn.Linear(config.hidden_dim, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 2),    # binary: source=0, target=1
        )

    def forward(self, x, return_domain=False):
        z = self.encode(x)
        class_logits = self.classifier(z)
        if not return_domain:
            return class_logits
        reversed_z = self.grl(z)
        domain_logits = self.domain_discriminator(reversed_z)
        return class_logits, domain_logits

    def add_strategy_loss(self, class_logits, domain_logits,
                          labels, domain_labels, device):
        """
        DANN combined objective:
          L_total = L_CE(class) + λ_d · L_CE(domain)

        The GRL negation during backward pass automatically promotes
        domain invariance in the encoder.
        """
        class_loss  = F.cross_entropy(class_logits, labels)
        domain_loss = F.cross_entropy(domain_logits, domain_labels)
        total_loss  = class_loss + self.cfg.lambda_domain * domain_loss
        return total_loss, class_loss.detach(), domain_loss.detach()

# ── CORAL Model ─────────────────────────────────────────────────────────────────

class CORALModel(BaseModel):
    """
    Correlation Alignment (Sun & Saenko, ECCV 2016).

    Matches second-order statistics (covariance matrices) of source and target
    domain bottleneck features. Loss = Frobenius norm of the difference between
    whitened covariance matrices, computed in mini-batch approximation.

    STRATEGY: covariance-matrix alignment, no adversarial signal.
    """

    def __init__(self, config):
        super().__init__(config)

    @staticmethod
    def coral_loss(z, domain_labels, device):
        """
        Mini-batch CORAL loss.

        For each batch:
          1. Centre:   z_s -= mean(z_s),  z_t -= mean(z_t)
          2. Covariance: C_s = (z_sᵀ z_s)/(n_s-1),  C_t = (z_tᵀ z_t)/(n_t-1)
          3. Frobenius:  L_coral = || C_s − C_t ||_F²

        Reference: Sun, Saenko — "Deep CORAL: Correlation Alignment for Deep
        Domain Adaptation" ECCV 2016.
        """
        if domain_labels is None:
            return torch.tensor(0.0, device=device)

        mask_s = (domain_labels == 0).unsqueeze(1).float()
        mask_t = (domain_labels == 1).unsqueeze(1).float()

        n_s = mask_s.sum().item() + 1e-8
        n_t = mask_t.sum().item() + 1e-8

        # Mean centering
        mean_s = (z * mask_s).sum(0, keepdim=True) / n_s
        mean_t = (z * mask_t).sum(0, keepdim=True) / n_t
        z_s = z - mean_s
        z_t = z - mean_t

        # Covariance (d×d)
        C_s = (z_s.t() @ z_s) / max(n_s - 1, 1)
        C_t = (z_t.t() @ z_t) / max(n_t - 1, 1)

        # Frobenius norm squared
        diff = C_s - C_t
        L_coral = (diff ** 2).sum()
        return L_coral

    def forward(self, x):
        z = self.encode(x)
        return self.classifier(z)

    def add_strategy_loss(self, class_logits, z, labels, domain_labels, device):
        class_loss = F.cross_entropy(class_logits, labels)
        coral_loss_val = self.coral_loss(z, domain_labels, device)
        total_loss = class_loss + self.cfg.lambda_domain * coral_loss_val
        return total_loss, class_loss.detach(), coral_loss_val.detach()

# ── Raman Self-Supervised Model ───────────────────────────────────────────────

class RamanSelfSupModel(BaseModel):
    """
    Proposed method: Raman-only self-supervised pre-training.

    Uses a masked-spectrum reconstruction auxiliary task on the flattened
    bottleneck output to learn domain-general spectral features without label
    supervision. After pre-training the encoder is frozen and only the
    classifier is fine-tuned — matching the clinical scenario where only
    limited labelled Raman spectra exist.

    Pseudo-spectrum: flatten [3,32,32] → 3072-d vector (proxy for Raman
    wavenumber × intensity vector).

    STRATEGY: self-supervised masked reconstruction, frozen encoder.
    """

    def __init__(self, config):
        super().__init__(config)
        self.cfg = config
        self.encoder_frozen = False

        D = config.hidden_dim
        self.spec_projector = nn.Sequential(
            nn.Linear(D, 512), nn.ReLU(),
            nn.Linear(512, 1024), nn.ReLU(),
            nn.Linear(1024, 3072),   # match flattened image dim
        )

    def encode(self, x):
        feat = self.encoder(x)        # [B, 512]
        z    = self.bottleneck(feat)  # [B, hidden_dim]
        return z

    def forward(self, x=None, masked_spectrum=None, labels=None):
        """
        Dual-mode forward:
          Mode 1 (x provided, masked_spectrum=None): standard classification logits
          Mode 2 (masked_spectrum provided): reconstruct full spectrum + class logits

        Always returns a 2-tuple (class_logits, recon) so callers can unpack
        uniformly regardless of mode.  recon is None in Mode 1.
        """
        class_logits = None
        recon = None

        if masked_spectrum is not None:
            # ── Self-supervised mode ────────────────────────────
            # Reshape pseudo-spectrum back to image shape for encoder
            B = masked_spectrum.size(0)
            img_shape = masked_spectrum.view(B, 3, 32, 32)
            z = self.encode(img_shape)
            recon = self.spec_projector(z)              # [B, 3072]
            class_logits = self.classifier(z.detach())  # stop grad → encoder
        elif x is not None:
            # ── Standard classification mode ─────────────────────
            z = self.encode(x)
            class_logits = self.classifier(z)

        return class_logits, recon

    def add_strategy_loss(self, class_logits, recon, targets, labels, device):
        """
        Combined loss:
          L_total = L_CE(class) + λ_self · MSE(recon, targets)

        MSE encourages encoder features whose linear projection can reconstruct
        masked spectral regions (proxy for Raman wavenumber reconstruction).
        """
        class_loss = F.cross_entropy(class_logits, labels)
        # recon is None when called in eval / inference mode
        recon_loss = torch.tensor(0.0, device=device)
        if recon is not None:
            recon_loss = F.mse_loss(recon, targets)
        total_loss = class_loss + self.cfg.lambda_selfsup * recon_loss
        return total_loss, class_loss.detach(), recon_loss.detach()

    def freeze_encoder(self):
        """Freeze encoder + bottleneck; only classifier + projector trainable."""
        for p in self.encoder.parameters():
            p.requires_grad = False
        for p in self.bottleneck.parameters():
            p.requires_grad = False
        self.encoder_frozen = True