src/mavt/training/lightning_module.py

"""PyTorch Lightning Module for MAVT training.

Supports 3-stage curriculum via `training_stage` parameter:
  1 — image only,  SigLIP2 fully frozen,  LR = 1e-4
  2 — +video,      SigLIP2 last 4 unfrozen, LR = 5e-5
  3 — +3D,         SigLIP2 fully unfrozen, LR = 2e-5

To move between stages: start training with the next stage config and pass
`--ckpt_path <prev_stage_checkpoint>` to LightningCLI.
"""

from __future__ import annotations
from typing import Any, Dict, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
import lightning as L
from lightning.pytorch.utilities import grad_norm

from mavt.model.mavt import MAVT
from mavt.losses.losses import MAVTLoss


_STAGE_LR = {1: 1e-4, 2: 5e-5, 3: 2e-5}
_STAGE_SIGLIP2_FROZEN_BLOCKS = {1: 10, 2: 6, 3: 0}  # number of transformer blocks frozen
_STAGE_W_SEM = {1: 0.5, 2: 0.3, 3: 0.2}             # default cosine-distill weight per stage


class MAVTLightningModule(L.LightningModule):
    """Lightning module for end-to-end MAVT training."""

    def __init__(
        self,
        # Model
        embed_dim: int = 1152,
        num_heads: int = 16,
        num_blocks: int = 12,
        patch_size: int = 16,
        t_patch: int = 2,
        latent_dim: int = 32,
        kl_weight: float = 1e-4,
        semantic_dim: int = 768,
        dec_dim: int = 768,
        num_dec_attn_blocks: int = 4,
        r_s: int = 2,
        r_t: int = 1,
        use_gradient_checkpointing: bool = True,
        mlp_ratio: float = 4.0,
        dropout: float = 0.0,
        # C-D split (cd-split commit 6368dfb)
        local_detail_window_size: int = 1,
        local_detail_temporal_window_size: int = 1,
        # Loss
        w_l1: float = 1.0,
        w_lpips: float = 0.1,
        w_kl: float = 1.0,      # passthrough — KL pre-scaled by VAEHead.kl_weight
        w_clip: float = 0.0,
        w_sem: float = 0.0,
        w_aux: float = 0.01,
        w_temp: float = 0.0,    # temporal consistency weight (video only)
        use_lpips: bool = True,
        use_clip: bool = False,
        active_modalities: list = None,  # e.g. ['image', 'video']; None → all three
        # Curriculum
        training_stage: int = 1,
        siglip2_model_name: str = "google/siglip2-base-patch16-224",
        init_siglip2: bool = True,
        use_semantic_distill: bool = False,
        # Cross-stage weight transfer (loads model weights only, NOT optimizer
        # / scheduler / step state — use --ckpt_path for true resume instead).
        init_from_ckpt: Optional[str] = None,
        # Optimiser
        weight_decay: float = 0.01,
        grad_clip: float = 1.0,
        warmup_steps: int = 1000,
        total_steps: int = 200_000,
    ):
        super().__init__()
        self.save_hyperparameters()

        self.model = MAVT(
            embed_dim=embed_dim, num_heads=num_heads, num_blocks=num_blocks,
            patch_size=patch_size, t_patch=t_patch,
            latent_dim=latent_dim, kl_weight=kl_weight,
            semantic_dim=semantic_dim, dec_dim=dec_dim,
            num_dec_attn_blocks=num_dec_attn_blocks, r_s=r_s, r_t=r_t,
            use_gradient_checkpointing=use_gradient_checkpointing,
            mlp_ratio=mlp_ratio, dropout=dropout,
            local_detail_window_size=local_detail_window_size,
            local_detail_temporal_window_size=local_detail_temporal_window_size,
        )

        _active_mods = tuple(active_modalities) if active_modalities else ('image', 'video', 'threed')
        self.loss_fn = MAVTLoss(
            w_l1=w_l1, w_lpips=w_lpips, w_kl=w_kl,
            w_clip=w_clip, w_sem=w_sem, w_aux=w_aux,
            w_temp=w_temp,
            use_lpips=use_lpips, use_clip=use_clip,
            active_modalities=_active_mods,
        )

        # Frozen vision teacher (loaded lazily in setup() to keep __init__ light)
        self.semantic_teacher: Optional[nn.Module] = None
        self._teacher_image_size: int = 224

    # ------------------------------------------------------------------ #
    #  Setup                                                               #
    # ------------------------------------------------------------------ #

    def setup(self, stage: str) -> None:
        hp = self.hparams
        if stage != 'fit':
            return
        # Eager slot pooler creation — must run BEFORE configure_optimizers
        # so the new params land in the optimizer's param_groups.
        self._prepare_cd_split_poolers()
        # Sync EMA modalities from DataModule — single source of truth
        self._sync_ema_modalities()
        if hp.init_siglip2:
            frozen = _STAGE_SIGLIP2_FROZEN_BLOCKS[hp.training_stage]
            self.model.load_siglip2_weights(hp.siglip2_model_name, frozen)
        if hp.use_semantic_distill and self.semantic_teacher is None:
            self._load_semantic_teacher(hp.siglip2_model_name)
        # Cross-stage weight transfer (after siglip2 / teacher are in place so
        # they get overwritten by ckpt values when present).
        if hp.init_from_ckpt:
            self._load_weights_from_ckpt(hp.init_from_ckpt)

    def _load_weights_from_ckpt(self, path: str) -> None:
        # weights_only=False is required because Lightning ckpts contain a
        # full pickle (state_dict + hparams + callbacks). Source is our own
        # filesystem so untrusted-pickle risk is N/A.
        ckpt = torch.load(path, map_location='cpu', weights_only=False)
        sd = ckpt.get('state_dict', ckpt)
        missing, unexpected = self.load_state_dict(sd, strict=False)
        kept_missing = [
            k for k in missing
            if not k.startswith('semantic_teacher.')
            and not k.startswith('model.cd_split._content_poolers.')
            and not k.startswith('model.cd_split._detail_poolers.')
        ]
        print(f"[init_from_ckpt] loaded {path}")
        print(f"[init_from_ckpt] missing (kept random init): {len(kept_missing)} keys "
              f"+ {len(missing) - len(kept_missing)} expected (teacher/new poolers)")
        if unexpected:
            print(f"[init_from_ckpt] unexpected (dropped): {len(unexpected)} keys")
        print("[init_from_ckpt] NOTE: optimizer state / LR scheduler / global_step are NOT "
              "restored — this is a soft restart. For exact resume of the SAME stage, "
              "use Lightning --ckpt_path instead.")

    def _prepare_cd_split_poolers(self) -> None:
        """Read active modality + resolution from the attached DataModule and
        eagerly create every SlotPooler the trainer will need."""
        dm = getattr(self.trainer, 'datamodule', None)
        if dm is None or not hasattr(dm, 'hparams'):
            return
        dm_hp = dm.hparams
        active = getattr(dm_hp, 'active_modalities', None) or []
        specs = []
        for modality in active:
            if modality == 'image':
                specs.append({
                    'modality': 'image',
                    'resolution': dm_hp.image_resolution,
                })
            elif modality == 'video':
                specs.append({
                    'modality': 'video',
                    'resolution': dm_hp.video_resolution,
                    'frames':     dm_hp.video_frames,
                    't_patch':    self.hparams.t_patch,
                })
            elif modality == 'threed':
                specs.append({
                    'modality': 'threed',
                    'resolution': dm_hp.triplane_res,
                })
        if specs:
            self.model.prepare_for_modalities(specs)

    def _sync_ema_modalities(self) -> None:
        """Use DataModule as single source of truth for active_modalities.

        Overrides whatever was set in model config so data.active_modalities
        and the EMA weighter never drift apart.
        """
        dm = getattr(self.trainer, 'datamodule', None)
        if dm is not None and hasattr(dm, 'hparams'):
            active = getattr(dm.hparams, 'active_modalities', None)
            if active:
                self.loss_fn.ema_weighter.active_modalities = tuple(active)
                return
        # Fallback: use model config param if DataModule not available
        fallback = getattr(self.hparams, 'active_modalities', None)
        if fallback:
            self.loss_fn.ema_weighter.active_modalities = tuple(fallback)

    def _load_semantic_teacher(self, model_name: str) -> None:
        """Load frozen SigLIP2 vision tower as teacher for cosine distillation."""
        try:
            from transformers import AutoModel
            siglip = AutoModel.from_pretrained(model_name)
            teacher = siglip.vision_model
            for p in teacher.parameters():
                p.requires_grad_(False)
            teacher.eval()
            self.semantic_teacher = teacher
            try:
                self._teacher_image_size = int(siglip.config.vision_config.image_size)
            except AttributeError:
                self._teacher_image_size = 224
        except Exception as exc:  # noqa: BLE001
            print(f"[lightning_module] semantic teacher load failed ({exc}); "
                  f"distillation disabled this run")
            self.semantic_teacher = None

    def train(self, mode: bool = True):  # type: ignore[override]
        """Keep frozen teacher in eval mode regardless of train()/eval() calls."""
        super().train(mode)
        if self.semantic_teacher is not None:
            self.semantic_teacher.eval()
        return self

    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        """Strip frozen teacher weights from checkpoints to keep them small."""
        state = checkpoint.get('state_dict', {})
        for k in list(state.keys()):
            if k.startswith('semantic_teacher.'):
                del state[k]

    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        """Inject the freshly-loaded teacher's params back into the checkpoint
        state_dict so Lightning's strict load does not fail with missing keys.

        Order at training-resume time:
          1. setup('fit')               → teacher loaded from HF (deterministic)
          2. on_load_checkpoint(ckpt)   → we inject teacher keys here
          3. self.load_state_dict(ckpt) → strict=True now sees a complete dict

        Teacher weights are deterministic from the HF model name, so injecting
        the current values is equivalent to whatever the checkpoint would have
        contained — no behavior change, only avoids the strict-load error.
        """
        if self.semantic_teacher is None:
            return
        sd = checkpoint.setdefault('state_dict', {})
        for k, v in self.semantic_teacher.state_dict().items():
            sd.setdefault(f'semantic_teacher.{k}', v)

    # ------------------------------------------------------------------ #
    #  Training step                                                       #
    # ------------------------------------------------------------------ #

    def _step(self, batch: Dict, log_prefix: str) -> torch.Tensor:
        x       = batch['data']
        modality = batch['modality']

        out = self.model(x, modality, decode=True)

        # For video: decoder reconstructs in patch-grid temporal space (Tp = T//t_patch).
        # Target must match. Use first frame of each temporal patch group.
        if modality == 'video':
            t_patch = self.hparams.t_patch
            target = x[:, :, ::t_patch]  # (B, 3, Tp, H, W)
        else:
            target = x  # image and threed pass through as-is

        # Vision-vision distillation: forward x_proxy through frozen teacher.
        teacher_embed: Optional[torch.Tensor] = None
        if self.semantic_teacher is not None:
            with torch.no_grad():
                proxy = self._make_teacher_input(x, modality, self._teacher_image_size)
                teacher_embed = self.semantic_teacher(pixel_values=proxy).pooler_output

        losses = self.loss_fn(
            pred=out.reconstruction,
            target=target,
            loss_kl=out.loss_kl,
            slot_diversity=out.cd_metrics['slot_diversity'],
            modality=modality,
            semantic_embed=out.semantic,
            teacher_embed=teacher_embed,
        )

        # Logging — aggregate (all modalities combined)
        for k, v in losses.items():
            self.log(f'{log_prefix}/{k}', v, on_step=True, on_epoch=True,
                     prog_bar=(k == 'loss'), sync_dist=True)
        # Per-modality breakdown (diagnose image vs video separately)
        for k in ('loss', 'loss_recon', 'loss_l1', 'loss_kl'):
            if k in losses:
                self.log(f'{log_prefix}/{k}_{modality}', losses[k],
                         on_step=True, on_epoch=True, sync_dist=True)
        for k, v in out.cd_metrics.items():
            self.log(f'{log_prefix}/cd_{k}', v, on_step=False, on_epoch=True,
                     sync_dist=True)
        self.log(f'{log_prefix}/modality_{modality}', 1.0,
                 on_step=False, on_epoch=True, sync_dist=False)

        return losses['loss']

    def training_step(self, batch: Dict, batch_idx: int) -> torch.Tensor:
        return self._step(batch, 'train')

    def validation_step(self, batch: Dict, batch_idx: int) -> None:
        with torch.no_grad():
            self._step(batch, 'val')
            # Log sample reconstructions to wandb/tensorboard every N steps
            if batch_idx == 0:
                self._log_images(batch)

    # ------------------------------------------------------------------ #
    #  Teacher input proxy                                                 #
    # ------------------------------------------------------------------ #

    @staticmethod
    def _make_teacher_input(x: torch.Tensor, modality: str,
                            target_size: int) -> torch.Tensor:
        """Project the multi-modal input down to a single (B, 3, S, S) image
        the SigLIP2 vision teacher can consume.

        image  : x as-is.
        video  : middle frame.
        threed : XY plane (front view, closest to natural-image distribution).
        """
        if modality == 'image':
            proxy = x
        elif modality == 'video':
            T = x.shape[2]
            proxy = x[:, :, T // 2]                # (B, 3, H, W)
        elif modality == 'threed':
            proxy = x[:, 0]                         # (B, 3, H, W) plane XY
        else:
            raise ValueError(f"Unknown modality: {modality}")

        if proxy.shape[-1] != target_size or proxy.shape[-2] != target_size:
            proxy = F.interpolate(
                proxy, size=(target_size, target_size),
                mode='bilinear', align_corners=False,
            )
        return proxy

    # ------------------------------------------------------------------ #
    #  Visualisation                                                       #
    # ------------------------------------------------------------------ #

    def _log_images(self, batch: Dict, n: int = 4) -> None:
        try:
            x       = batch['data'][:n]
            modality = batch['modality']
            out = self.model(x, modality, decode=True)

            if modality == 'image':
                grid_in  = _to_grid(x)
                grid_out = _to_grid(out.reconstruction)
            elif modality == 'video':
                # Log a temporal strip: 4 evenly-spaced frames per clip stacked
                # horizontally so reviewers can spot temporal coherence (vs.
                # only a single first frame).
                B, C, T, H, W = x.shape
                Tp = out.reconstruction.shape[2]
                t_in = torch.linspace(0, T - 1, 4).long()
                t_out = torch.linspace(0, Tp - 1, 4).long()
                in_strip  = x[:, :, t_in].permute(0, 2, 1, 3, 4).reshape(B * 4, C, H, W)
                out_strip = out.reconstruction[:, :, t_out].permute(0, 2, 1, 3, 4).reshape(B * 4, C, H, W)
                grid_in  = _to_grid(in_strip,  nrow=4)
                grid_out = _to_grid(out_strip, nrow=4)
            elif modality == 'threed':
                # Log XY plane (plane index 0)
                grid_in  = _to_grid(x[:, 0])
                grid_out = _to_grid(out.reconstruction[:, 0])
            else:
                return

            loggers = self.loggers if isinstance(self.loggers, (list, tuple)) else [self.loggers]
            for logger in loggers:
                if hasattr(logger, 'log_image'):
                    logger.log_image(key=f'val/{modality}_input',  images=[grid_in])
                    logger.log_image(key=f'val/{modality}_recon',  images=[grid_out])
        except Exception:  # noqa: BLE001
            pass

    # ------------------------------------------------------------------ #
    #  Optimiser                                                           #
    # ------------------------------------------------------------------ #

    def configure_optimizers(self):
        hp = self.hparams
        lr = _STAGE_LR[hp.training_stage]

        # Separate RGAT params for potential different LR (currently same LR)
        rgat_params, other_params = [], []
        for name, p in self.model.named_parameters():
            if not p.requires_grad:
                continue
            if 'rgat' in name.lower() or 'rgat4d' in name.lower():
                rgat_params.append(p)
            else:
                other_params.append(p)

        param_groups = [{'params': other_params, 'lr': lr}]
        if rgat_params:
            param_groups.append({'params': rgat_params, 'lr': lr})

        optimizer = torch.optim.AdamW(param_groups, weight_decay=hp.weight_decay)

        # Linear warmup + cosine decay
        def lr_lambda(step: int) -> float:
            if step < hp.warmup_steps:
                return step / max(1, hp.warmup_steps)
            progress = (step - hp.warmup_steps) / max(1, hp.total_steps - hp.warmup_steps)
            return max(0.0, 0.5 * (1.0 + torch.cos(torch.tensor(torch.pi * progress)).item()))

        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'interval': 'step',
                'frequency': 1,
            },
        }

    def configure_gradient_clipping(self, optimizer, gradient_clip_val=None, gradient_clip_algorithm=None):
        self.clip_gradients(optimizer, gradient_clip_val=self.hparams.grad_clip,
                            gradient_clip_algorithm='norm')


# --------------------------------------------------------------------------- #
#  Visualisation helper                                                         #
# --------------------------------------------------------------------------- #

def _to_grid(x: torch.Tensor, nrow: int = 4) -> Any:
    """Convert (B, 3, H, W) tensor to a PIL Image grid for logging."""
    try:
        from torchvision.utils import make_grid
        from PIL import Image
        import numpy as np
        x = x.detach().cpu().float().clamp(-1, 1)
        x = (x + 1) / 2                         # [0, 1]
        grid = make_grid(x, nrow=nrow, normalize=False)
        arr = (grid.permute(1, 2, 0).numpy() * 255).astype('uint8')
        return Image.fromarray(arr)
    except Exception:  # noqa: BLE001
        return None