"""
Knowledge Distillation for Video Classification
================================================

Distills knowledge from a V-JEPA2 teacher to a VideoMAEv2 student model.

The student (VideoMAEv2-Base, 86M params) is already pre-trained on video data,
making it an efficient starting point for collision anticipation.

Usage:
    python -m badas.training.distillation --config configs/distillation.yml
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torchmetrics import Accuracy, AUROC, AveragePrecision, Precision, Recall
from dataclasses import dataclass, field, asdict
from typing import Optional, List, Tuple, Dict, Any, Union
import time
import yaml
import os

from transformers import AutoConfig

from badas.training.config import Config, ModelConfig, DataConfig, TrainingConfig, SystemConfig, LoggingConfig
from badas.core import (
    create_model,
    detect_model_type,
    VJEPA2Classifier,
    VideoMAEv2Classifier,
    VideoClassifierBase,
)


# =============================================================================
# CONFIGURATION
# =============================================================================

@dataclass
class DistillationConfig:
    """Configuration for knowledge distillation."""

    # Teacher model (V-JEPA2)
    teacher_checkpoint: str = ""
    teacher_model_name: str = "facebook/vjepa2-vitl-fpc16-256-ssv2"
    freeze_teacher: bool = True

    # Student model (VideoMAEv2)
    student_model_name: str = "OpenGVLab/VideoMAEv2-Base"

    # Temperature for soft labels
    temperature: float = 4.0

    # Loss weights
    alpha_hard: float = 0.5    # Ground truth loss
    alpha_logit: float = 0.3   # Logit-level KD loss
    alpha_feature: float = 0.15  # Feature-level distillation
    alpha_attention: float = 0.05  # Attention transfer

    # Feature distillation settings
    feature_layers: Union[List[int], str] = "auto"
    num_feature_layers: int = 4
    feature_projection: str = "linear"

    # Attention distillation settings
    attention_layers: Union[List[int], str] = "auto"
    num_attention_layers: int = 2

    # Progressive training
    warmup_epochs_hard_only: int = 0

    # SSL backbone initialization
    ssl_backbone_checkpoint: str = ""     # Path to SSL .ckpt; encoder.* → student backbone

    # Phase 2 hard-only cooldown
    cooldown_steps_hard_only: int = 0     # Final N steps: hard BCE only (phase 2)

    # Offline distillation: path to precomputed teacher output cache directory.
    # If the directory exists, teacher is NOT loaded during training.
    # Build the cache first: python scripts/precompute_teacher_outputs.py --config ...
    # Default: alongside the clips in badas_clips_output/
    teacher_cache_path: str = ""


@dataclass
class FullDistillationConfig:
    """Complete configuration for distillation training."""

    distillation: DistillationConfig = field(default_factory=DistillationConfig)
    model: ModelConfig = field(default_factory=ModelConfig)
    data: DataConfig = field(default_factory=DataConfig)
    training: TrainingConfig = field(default_factory=TrainingConfig)
    system: SystemConfig = field(default_factory=SystemConfig)
    logging: LoggingConfig = field(default_factory=LoggingConfig)

    def save(self, path: str):
        """Save configuration to YAML file"""
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w") as f:
            yaml.dump(asdict(self), f, default_flow_style=False, indent=2)

    @classmethod
    def load(cls, path: str) -> "FullDistillationConfig":
        """Load configuration from YAML file"""
        with open(path, "r") as f:
            data = yaml.safe_load(f)

        return cls(
            distillation=DistillationConfig(**data.get("distillation", {})),
            model=ModelConfig(**data.get("model", {})),
            data=DataConfig(**data.get("data", {})),
            training=TrainingConfig(**data.get("training", {})),
            system=SystemConfig(**data.get("system", {})),
            logging=LoggingConfig(**data.get("logging", {})),
        )

    def to_student_config(self) -> Config:
        """Convert to standard Config for compatibility."""
        return Config(
            model=self.model,
            data=self.data,
            training=self.training,
            system=self.system,
            logging=self.logging,
        )


# =============================================================================
# FEATURE PROJECTION
# =============================================================================

class FeatureProjector(nn.Module):
    """Projects student features to match teacher dimensions."""

    def __init__(self, student_dim: int, teacher_dim: int, projection_type: str = "linear"):
        super().__init__()

        if student_dim == teacher_dim:
            self.projector = nn.Identity()
        elif projection_type == "linear":
            self.projector = nn.Linear(student_dim, teacher_dim)
        elif projection_type == "mlp":
            hidden_dim = (student_dim + teacher_dim) // 2
            self.projector = nn.Sequential(
                nn.Linear(student_dim, hidden_dim),
                nn.GELU(),
                nn.LayerNorm(hidden_dim),
                nn.Linear(hidden_dim, teacher_dim),
            )
        else:
            self.projector = nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.projector(x)


class MultiLayerProjector(nn.Module):
    """Handles projection for multiple layers."""

    def __init__(self, student_dim: int, teacher_dim: int, num_layers: int, projection_type: str = "linear"):
        super().__init__()
        self.projectors = nn.ModuleList([
            FeatureProjector(student_dim, teacher_dim, projection_type)
            for _ in range(num_layers)
        ])

    def forward(self, features: List[torch.Tensor]) -> List[torch.Tensor]:
        return [proj(feat) for proj, feat in zip(self.projectors, features)]


# =============================================================================
# DISTILLATION LOSS
# =============================================================================

class DistillationLoss(nn.Module):
    """Combined loss for knowledge distillation."""

    def __init__(self, config: DistillationConfig):
        super().__init__()
        self.config = config
        self.temperature = config.temperature
        self.alpha_hard = config.alpha_hard
        self.alpha_logit = config.alpha_logit
        self.alpha_feature = config.alpha_feature
        self.alpha_attention = config.alpha_attention

        self.ce_loss = nn.BCEWithLogitsLoss()
        self.kl_loss = nn.KLDivLoss(reduction="batchmean")
        self.mse_loss = nn.MSELoss()

    def forward(
        self,
        student_logits: torch.Tensor,
        teacher_logits: torch.Tensor,
        labels: torch.Tensor,
        student_features: Optional[List[torch.Tensor]] = None,
        teacher_features: Optional[List[torch.Tensor]] = None,
        student_attentions: Optional[List[torch.Tensor]] = None,
        teacher_attentions: Optional[List[torch.Tensor]] = None,
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """Compute combined distillation loss."""
        loss_dict = {}
        total_loss = 0.0

        # 1. HARD LABEL LOSS
        hard_loss = self.ce_loss(student_logits, labels.float())
        loss_dict["hard_loss"] = hard_loss.item()
        total_loss = total_loss + self.alpha_hard * hard_loss

        # 2. LOGIT DISTILLATION
        if self.alpha_logit > 0:
            T = self.temperature
            student_soft = torch.stack([
                F.log_softmax(torch.stack([-student_logits, student_logits], dim=-1) / T, dim=-1)
            ], dim=0).squeeze(0)
            teacher_soft = torch.stack([
                F.softmax(torch.stack([-teacher_logits, teacher_logits], dim=-1) / T, dim=-1)
            ], dim=0).squeeze(0)

            logit_loss = self.kl_loss(student_soft, teacher_soft.detach()) * (T ** 2)
            loss_dict["logit_loss"] = logit_loss.item()
            total_loss = total_loss + self.alpha_logit * logit_loss

        # 3. FEATURE DISTILLATION
        if self.alpha_feature > 0 and student_features and teacher_features:
            feature_loss = 0.0
            num_layers = min(len(student_features), len(teacher_features))

            for s_feat, t_feat in zip(student_features, teacher_features):
                s_norm = F.normalize(s_feat, dim=-1)
                t_norm = F.normalize(t_feat.detach(), dim=-1)
                feature_loss = feature_loss + self.mse_loss(s_norm, t_norm)

            feature_loss = feature_loss / max(num_layers, 1)
            loss_dict["feature_loss"] = feature_loss.item()
            total_loss = total_loss + self.alpha_feature * feature_loss

        # 4. ATTENTION TRANSFER
        if self.alpha_attention > 0 and student_attentions and teacher_attentions:
            attention_loss = 0.0
            num_layers = min(len(student_attentions), len(teacher_attentions))

            for s_attn, t_attn in zip(student_attentions, teacher_attentions):
                s_mean = s_attn.mean(dim=1).flatten(1)
                t_mean = t_attn.mean(dim=1).flatten(1).detach()
                s_log = F.log_softmax(s_mean, dim=-1)
                t_soft = F.softmax(t_mean, dim=-1)
                attention_loss = attention_loss + self.kl_loss(s_log, t_soft)

            attention_loss = attention_loss / max(num_layers, 1)
            loss_dict["attention_loss"] = attention_loss.item()
            total_loss = total_loss + self.alpha_attention * attention_loss

        loss_dict["total_loss"] = total_loss.item()
        return total_loss, loss_dict


# =============================================================================
# DISTILLATION MODULE
# =============================================================================

class DistillationModule(pl.LightningModule):
    """
    PyTorch Lightning module for distillation.

    Teacher: V-JEPA2 (bucharest-mongoose)
    Student: VideoMAEv2-Base (86M params, pre-trained)
    """

    def __init__(self, config: FullDistillationConfig):
        super().__init__()
        self.config = config
        self.distill_config = config.distillation

        self.save_hyperparameters(asdict(config))

        # Offline cache: if cache exists, teacher is not loaded during training.
        self.teacher_cache = self._try_load_teacher_cache()
        self._offline = self.teacher_cache is not None

        # Initialize teacher (skipped when offline cache is ready)
        if not self._offline:
            self._init_teacher()
        else:
            # Still need teacher_dim and teacher_layers for projector / layer mapping.
            # Read them from the HF config without downloading the full model.
            self._init_teacher_metadata_only()

        # Initialize student
        self._init_student()

        # Initialize projectors
        self._init_projectors()

        # Setup layer indices for distillation
        self._setup_layer_indices()

        # Initialize loss
        self.criterion = DistillationLoss(self.distill_config)

        # Metrics
        self._init_metrics()

    def _try_load_teacher_cache(self):
        """Load TeacherOutputCache if the cache directory exists and is complete."""
        path = self.distill_config.teacher_cache_path
        if not path:
            return None
        from badas.training.teacher_cache import TeacherOutputCache
        index_file = os.path.join(path, TeacherOutputCache.INDEX_FILE)
        if os.path.exists(index_file):
            print(f"\n[Offline distillation] Loading teacher cache: {path}")
            return TeacherOutputCache(path)
        print(f"\n[Offline distillation] Cache not found at {path} — will run teacher online.")
        return None

    def _init_teacher_metadata_only(self):
        """Read teacher architecture metadata (dim, layers) without loading the model."""
        print(f"\nTeacher metadata: {self.distill_config.teacher_model_name} (offline mode, not loaded)")
        self._teacher_hf_config = AutoConfig.from_pretrained(
            self.distill_config.teacher_model_name,
            trust_remote_code=True,
        )
        self.teacher_dim    = self._teacher_hf_config.hidden_size
        self.teacher_layers = self._teacher_hf_config.num_hidden_layers
        print(f"  Dimension: {self.teacher_dim}, Layers: {self.teacher_layers}")
        # Assign dummy teacher so save_inference_checkpoint / _report_architecture don't break
        self.teacher = None

    def _init_teacher(self):
        """Initialize V-JEPA2 teacher model."""
        print(f"\nLoading teacher: {self.distill_config.teacher_model_name}")

        self._teacher_hf_config = AutoConfig.from_pretrained(
            self.distill_config.teacher_model_name,
            trust_remote_code=True
        )

        self._teacher_config = Config(
            model=ModelConfig(
                model_name=self.distill_config.teacher_model_name,
                num_classes=self.config.model.num_classes,
                use_future_prediction=self.config.model.use_future_prediction,
                temporal_method=self.config.model.temporal_method,
                head_type=self.config.model.head_type,
                head_hidden_dim=self.config.model.head_hidden_dim,
                head_num_layers=self.config.model.head_num_layers,
                head_dropout=self.config.model.head_dropout,
            ),
            data=self.config.data,
        )

        if self.distill_config.teacher_checkpoint:
            self.teacher = self._load_checkpoint(
                self.distill_config.teacher_checkpoint,
                self._teacher_config
            )
        else:
            self.teacher = create_model(self._teacher_config)

        self.teacher_dim = self._teacher_hf_config.hidden_size  # 1024 for ViT-L
        self.teacher_layers = self._teacher_hf_config.num_hidden_layers  # 24 for ViT-L

        print(f"  Dimension: {self.teacher_dim}")
        print(f"  Layers: {self.teacher_layers}")

        # Freeze teacher
        if self.distill_config.freeze_teacher:
            for param in self.teacher.parameters():
                param.requires_grad = False
            self.teacher.eval()

    def _init_student(self):
        """Initialize VideoMAEv2 student model."""
        print(f"\nLoading student: {self.distill_config.student_model_name}")

        self._student_config = Config(
            model=ModelConfig(
                model_name=self.distill_config.student_model_name,
                num_classes=self.config.model.num_classes,
                use_future_prediction=False,  # VideoMAEv2 doesn't have predictor
                temporal_method=self.config.model.temporal_method,
                head_type=self.config.model.head_type,
                head_hidden_dim=self.config.model.head_hidden_dim,
                head_num_layers=self.config.model.head_num_layers,
                head_dropout=self.config.model.head_dropout,
            ),
            data=self.config.data,
        )

        self.student = create_model(self._student_config)
        self.student_dim = self.student.feature_dim

        # Load SSL backbone weights if provided
        if self.distill_config.ssl_backbone_checkpoint:
            self._load_ssl_backbone()

        # Get layer count from HuggingFace config
        student_hf_config = AutoConfig.from_pretrained(
            self.distill_config.student_model_name,
            trust_remote_code=True
        )
        if hasattr(student_hf_config, 'model_config'):
            self.student_layers = student_hf_config.model_config.get('depth', 12)
        else:
            self.student_layers = getattr(student_hf_config, 'num_hidden_layers', 12)

        # Calculate compression
        student_params = sum(p.numel() for p in self.student.parameters())
        if self.teacher is not None:
            teacher_params = sum(p.numel() for p in self.teacher.parameters())
            print(f"  Parameters: {student_params:,} (teacher: {teacher_params:,})")
            print(f"  Compression: {teacher_params / student_params:.2f}x")
        else:
            print(f"  Parameters: {student_params:,} (teacher: offline cache)")

    def _load_ssl_backbone(self):
        """Load SSL pre-trained weights into student backbone.

        Expects a checkpoint whose state_dict has keys prefixed with 'encoder.'
        matching the HuggingFace VJEPA2 backbone key names 1-to-1.
        Position embeddings are intentionally skipped (strict=False) when the
        temporal grid differs between SSL pre-training and fine-tuning.
        """
        path = self.distill_config.ssl_backbone_checkpoint
        print(f"\n[SSL] Loading backbone weights from: {path}")
        ckpt = torch.load(path, map_location='cpu')
        state_dict = ckpt.get('state_dict', ckpt)
        encoder_sd = {k: v for k, v in state_dict.items() if k.startswith('encoder.')}
        missing, unexpected = self.student.backbone.load_state_dict(encoder_sd, strict=False)
        print(f"[SSL] Loaded {len(encoder_sd)} keys | missing={len(missing)} unexpected={len(unexpected)}")

    def _init_projectors(self):
        """Initialize feature projectors."""
        if self.student_dim != self.teacher_dim and self.distill_config.alpha_feature > 0:
            num_layers = (
                self.distill_config.num_feature_layers
                if self.distill_config.feature_layers == "auto"
                else len(self.distill_config.feature_layers)
            )
            self.feature_projector = MultiLayerProjector(
                self.student_dim,
                self.teacher_dim,
                num_layers,
                self.distill_config.feature_projection
            )
        else:
            self.feature_projector = None

    def _setup_layer_indices(self):
        """Setup which layers to use for distillation."""
        # Feature layers
        if self.distill_config.feature_layers == "auto":
            n = self.distill_config.num_feature_layers
            self.teacher_feature_layers = [
                int(i * (self.teacher_layers - 1) / (n - 1)) for i in range(n)
            ]
            self.student_feature_layers = [
                int(i * (self.student_layers - 1) / (n - 1)) for i in range(n)
            ]
        else:
            self.teacher_feature_layers = self.distill_config.feature_layers
            self.student_feature_layers = [
                int(l * self.student_layers / self.teacher_layers)
                for l in self.teacher_feature_layers
            ]

        # Attention layers
        if self.distill_config.attention_layers == "auto":
            n = self.distill_config.num_attention_layers
            self.teacher_attention_layers = [
                int(i * (self.teacher_layers - 1) / (n - 1)) for i in range(n)
            ]
            self.student_attention_layers = [
                int(i * (self.student_layers - 1) / (n - 1)) for i in range(n)
            ]
        else:
            self.teacher_attention_layers = self.distill_config.attention_layers
            self.student_attention_layers = [
                int(l * self.student_layers / self.teacher_layers)
                for l in self.teacher_attention_layers
            ]

        print(f"\nDistillation layer mapping:")
        print(f"  Features: teacher {self.teacher_feature_layers} -> student {self.student_feature_layers}")
        print(f"  Attention: teacher {self.teacher_attention_layers} -> student {self.student_attention_layers}")

    def _init_metrics(self):
        """Initialize metrics."""
        task = "binary" if self.config.model.num_classes == 2 else "multiclass"
        num_classes = self.config.model.num_classes if task == "multiclass" else None

        self.train_accuracy = Accuracy(task=task, num_classes=num_classes)
        self.val_accuracy = Accuracy(task=task, num_classes=num_classes)
        self.val_precision = Precision(task=task, num_classes=num_classes)
        self.val_recall = Recall(task=task, num_classes=num_classes)
        self.val_auc = AUROC(task=task, num_classes=num_classes)
        self.val_ap = AveragePrecision(task=task, num_classes=num_classes)

    def _load_checkpoint(self, path: str, config: Config) -> VideoClassifierBase:
        """Load model from checkpoint using registry."""
        checkpoint = torch.load(path, map_location="cpu")

        # Determine model type
        model_type = checkpoint.get("model_type")
        if model_type is None:
            model_type = detect_model_type(config.model.model_name)

        model = create_model(config, model_type=model_type)

        state_dict = checkpoint.get("state_dict", checkpoint)
        cleaned = {k.replace("model.", ""): v for k, v in state_dict.items()}
        model.load_state_dict(cleaned, strict=False)
        return model

    def save_inference_checkpoint(self, path: str, which: str = "student"):
        """
        Save an inference-ready checkpoint for deployment.

        Args:
            path: Output path for the checkpoint
            which: "student" or "teacher"
        """
        if which == "student":
            model = self.student
            config = self._student_config
        elif which == "teacher":
            model = self.teacher
            config = self._teacher_config
        else:
            raise ValueError(f"which must be 'student' or 'teacher', got '{which}'")

        # Create inference-compatible checkpoint
        checkpoint = {
            "model_type": model.MODEL_TYPE,
            "hyper_parameters": asdict(config),
            "state_dict": {f"model.{k}": v for k, v in model.state_dict().items()},
        }

        os.makedirs(os.path.dirname(path) if os.path.dirname(path) else ".", exist_ok=True)
        torch.save(checkpoint, path)
        print(f"Saved {which} inference checkpoint: {path}")

    def _extract_teacher_intermediates(
        self,
        x: torch.Tensor,
    ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]:
        """Extract logits, features, and attentions from teacher."""
        need_hidden = self.distill_config.alpha_feature > 0
        need_attention = self.distill_config.alpha_attention > 0

        # Teacher forward (V-JEPA2)
        backbone_output = self.teacher.backbone(
            pixel_values_videos=x,
            skip_predictor=True,
            output_hidden_states=need_hidden,
            output_attentions=need_attention,
        )
        hidden_states = backbone_output.hidden_states if need_hidden else None
        attentions = backbone_output.attentions if need_attention else None
        logits = self.teacher(x)

        # Extract specific layers
        features = []
        if hidden_states is not None:
            for idx in self.teacher_feature_layers:
                if idx < len(hidden_states):
                    features.append(hidden_states[idx])

        attn_maps = []
        if attentions is not None:
            for idx in self.teacher_attention_layers:
                if idx < len(attentions):
                    attn_maps.append(attentions[idx])

        return logits, features, attn_maps

    def _extract_student_intermediates(
        self,
        x: torch.Tensor,
    ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]:
        """Extract logits, features, and attentions from student.

        For VJEPA2Classifier students a second backbone-only pass is used to
        retrieve hidden states / attentions when distillation needs them.
        For other students (e.g. VideoMAEv2) intermediate outputs are not
        available and empty lists are returned.
        """
        need_hidden = (
            self.distill_config.alpha_feature > 0
            and isinstance(self.student, VJEPA2Classifier)
        )
        need_attention = (
            self.distill_config.alpha_attention > 0
            and isinstance(self.student, VJEPA2Classifier)
        )

        logits = self.student(x)   # full forward (temporal pool + head)

        if not need_hidden and not need_attention:
            return logits, [], []

        # Second backbone-only pass for intermediates (student is small; consistent
        # with how teacher already does two passes in _extract_teacher_intermediates)
        backbone_output = self.student.backbone(
            pixel_values_videos=x,
            skip_predictor=True,
            output_hidden_states=need_hidden,
            output_attentions=need_attention,
        )
        hidden = (
            [backbone_output.hidden_states[i]
             for i in self.student_feature_layers
             if i < len(backbone_output.hidden_states)]
            if need_hidden else []
        )
        attns = (
            [backbone_output.attentions[i]
             for i in self.student_attention_layers
             if i < len(backbone_output.attentions)]
            if need_attention else []
        )
        return logits, hidden, attns

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through student."""
        return self.student(x)

    def training_step(self, batch, batch_idx):
        """Training step with distillation (online or offline via cache)."""
        videos, labels, paths = batch

        # Phase 2 cooldown: hard BCE only (no teacher forward / cache lookup)
        in_cooldown = (
            self.distill_config.cooldown_steps_hard_only > 0
            and self.trainer.max_steps > 0
            and self.trainer.global_step >= (
                self.trainer.max_steps - self.distill_config.cooldown_steps_hard_only
            )
        )
        self._in_cooldown = in_cooldown  # exposed for logging callback

        in_warmup = (
            self.distill_config.warmup_epochs_hard_only > 0 and
            self.current_epoch < self.distill_config.warmup_epochs_hard_only
        )

        # Student forward
        student_logits, student_features, student_attns = self._extract_student_intermediates(videos)

        if in_cooldown:
            loss = F.binary_cross_entropy_with_logits(student_logits, labels.float())
            loss_dict = {"hard_loss": loss.item(), "total_loss": loss.item()}
            self.log("train_phase", 2.0, prog_bar=False)

        elif self._offline:
            # ── Offline distillation: precomputed teacher outputs ────────────
            teacher_logits, teacher_features_cached = self.teacher_cache.lookup_batch(
                paths, device=videos.device
            )
            # Mean-pool student features to match global cached teacher features
            if student_features:
                student_features_pooled = [f.mean(dim=1) for f in student_features]
                if self.feature_projector is not None:
                    student_features_pooled = self.feature_projector(student_features_pooled)
            else:
                student_features_pooled = []

            loss, loss_dict = self.criterion(
                student_logits, teacher_logits, labels,
                student_features_pooled, teacher_features_cached,
                [], [],   # attention: not cached (disabled in config anyway)
            )

        elif in_warmup:
            loss = F.binary_cross_entropy_with_logits(student_logits, labels.float())
            loss_dict = {"hard_loss": loss.item(), "total_loss": loss.item()}

        else:
            # ── Online distillation: run teacher every step ──────────────────
            # Project full-sequence student features before teacher forward
            if self.feature_projector is not None and student_features:
                student_features = self.feature_projector(student_features)

            with torch.no_grad():
                teacher_logits, teacher_features, teacher_attns = self._extract_teacher_intermediates(videos)

            loss, loss_dict = self.criterion(
                student_logits, teacher_logits, labels,
                student_features, teacher_features,
                student_attns, teacher_attns,
            )

        preds = (student_logits > 0.0).long()
        self.train_accuracy(preds, labels)

        self.log("train_loss", loss_dict["total_loss"], prog_bar=True)
        self.log("train_hard_loss", loss_dict["hard_loss"])
        if "logit_loss" in loss_dict:
            self.log("train_logit_loss", loss_dict["logit_loss"])
        if "feature_loss" in loss_dict:
            self.log("train_feature_loss", loss_dict["feature_loss"])
        self.log("train_acc", self.train_accuracy, on_epoch=True, prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        """Validation step."""
        videos, labels, _ = batch   # path unused in validation

        student_logits = self(videos)
        loss = F.binary_cross_entropy_with_logits(student_logits, labels.float())

        preds = (student_logits > 0.0).long()
        probs = torch.sigmoid(student_logits)

        self.val_accuracy(preds, labels)
        self.val_precision(preds, labels)
        self.val_recall(preds, labels)
        self.val_auc(probs, labels)
        self.val_ap(probs, labels)

        self.log("val_loss", loss, on_epoch=True, prog_bar=True, sync_dist=True)
        self.log("val_acc", self.val_accuracy, on_epoch=True, prog_bar=True, sync_dist=True)
        self.log("val_precision", self.val_precision, on_epoch=True, sync_dist=True)
        self.log("val_recall", self.val_recall, on_epoch=True, sync_dist=True)
        self.log("val_auc", self.val_auc, on_epoch=True, sync_dist=True)
        self.log("val_ap", self.val_ap, on_epoch=True, sync_dist=True)

        return loss

    def configure_optimizers(self):
        """Configure optimizer."""
        params = list(self.student.parameters())
        if self.feature_projector is not None:
            params += list(self.feature_projector.parameters())

        optimizer = AdamW(
            params,
            lr=self.config.training.learning_rate,
            weight_decay=self.config.training.weight_decay,
        )

        scheduler = CosineAnnealingLR(
            optimizer,
            T_max=self.config.training.scheduler_T_max,
            eta_min=self.config.training.scheduler_eta_min,
        )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": self.config.training.scheduler_step_size,
            },
        }

    def on_train_epoch_start(self):
        """Called at epoch start."""
        if self.current_epoch == 0:
            self._report_architecture()

    def _report_architecture(self):
        """Print architecture summary."""
        student_params    = sum(p.numel() for p in self.student.parameters())
        student_trainable = sum(p.numel() for p in self.student.parameters() if p.requires_grad)

        print(f"\n{'='*60}")
        print("DISTILLATION SUMMARY")
        print(f"{'='*60}")
        if self._offline:
            print(f"Teacher: {self.distill_config.teacher_model_name} [OFFLINE CACHE]")
            print(f"  Cache: {self.distill_config.teacher_cache_path}")
            print(f"  Cached clips: {len(self.teacher_cache):,}")
        else:
            teacher_params = sum(p.numel() for p in self.teacher.parameters())
            print(f"Teacher: {self.distill_config.teacher_model_name}")
            print(f"  Parameters: {teacher_params:,}")
        print(f"  Dimension: {self.teacher_dim}, Layers: {self.teacher_layers}")
        print(f"\nStudent: {self.distill_config.student_model_name}")
        print(f"  Dimension: {self.student_dim}, Layers: {self.student_layers}")
        print(f"  Parameters: {student_params:,} ({student_trainable:,} trainable)")
        print(f"\nDistillation mode: {'OFFLINE (cache)' if self._offline else 'ONLINE (live teacher)'}")
        print(f"\nLoss Weights:")
        print(f"  Hard (GT):    {self.distill_config.alpha_hard}")
        print(f"  Logit KD:     {self.distill_config.alpha_logit}")
        print(f"  Feature:      {self.distill_config.alpha_feature}")
        print(f"  Attention:    {self.distill_config.alpha_attention}")
        print(f"  Temperature:  {self.distill_config.temperature}")
        print(f"{'='*60}\n")

    def build_teacher_cache(self, cache_dir: Optional[str] = None, batch_size: int = 256,
                             num_workers: int = 8) -> None:
        """
        Precompute teacher outputs for all training clips and save to disk.

        After this completes, set teacher_cache_path in the config to cache_dir
        and re-run training — the teacher will not be loaded and each step will
        be ~2x faster.

        Args:
            cache_dir:   Directory to write cache files. Defaults to
                         teacher_cache_path from config, or
                         <data_root>/teacher_cache/ if not set.
            batch_size:  Clips per GPU batch (no grads → use large value, e.g. 256).
            num_workers: DataLoader worker count.
        """
        if self._offline:
            print("[build_teacher_cache] Already in offline mode — cache already exists.")
            return

        if cache_dir is None:
            cache_dir = self.distill_config.teacher_cache_path or os.path.join(
                self.config.data.data_root, "teacher_cache"
            )

        from badas.training.teacher_cache import TeacherOutputCache
        from badas.training.dataset import VideoDataset

        # Build dataset without augmentation, no resampling, with paths
        dataset = VideoDataset(
            config=self.config.to_student_config(),
            csv_path=self.config.data.csv_path,
            data_root=self.config.data.data_root,
            split_filter="train",
            resample=False,
            transform=None,
            return_path=True,
        )

        device = "cuda" if torch.cuda.is_available() else "cpu"
        TeacherOutputCache.build(
            teacher=self.teacher,
            teacher_feature_layers=self.teacher_feature_layers,
            dataset=dataset,
            cache_dir=cache_dir,
            batch_size=batch_size,
            num_workers=num_workers,
            device=device,
        )
        print(f"\n[build_teacher_cache] Done. Set teacher_cache_path: \"{cache_dir}\" in your config.")


# =============================================================================
# LOGGING CALLBACK
# =============================================================================

class DistillationLogCallback(pl.Callback):
    """
    Prints timestamped, human-readable progress to stdout every
    log_every_n_steps steps and after every validation.

    Works correctly in log files (no \\r trickery).
    """

    def __init__(self, max_steps: int, log_every_n_steps: int = 10):
        self._max_steps       = max_steps
        self._log_every       = log_every_n_steps
        self._t_start         = None
        self._last_step_time  = None
        self._last_logged_step = -1

    @staticmethod
    def _ts() -> str:
        return time.strftime("%H:%M:%S")

    def on_train_start(self, trainer, pl_module):
        self._t_start = time.time()
        self._last_step_time = self._t_start

    @pl.utilities.rank_zero_only
    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        step = trainer.global_step
        # Only log once per optimizer step (avoids duplicates with accumulate_grad_batches>1)
        if step == self._last_logged_step or step % self._log_every != 0 or step == 0:
            return
        self._last_logged_step = step

        now     = time.time()
        elapsed = now - self._t_start
        sps     = step / max(elapsed, 1e-3)           # steps per second
        eta_s   = (self._max_steps - step) / max(sps, 1e-6)

        m = trainer.callback_metrics
        loss      = float(m.get("train_loss",        m.get("loss",       0)))
        loss_hard = float(m.get("train_hard_loss",   m.get("loss_hard",  0)))
        loss_kd   = float(m.get("train_logit_loss",  m.get("loss_kd",    0)))
        loss_feat = float(m.get("train_feature_loss", m.get("loss_feat", 0)))
        lr        = float(m.get("lr-AdamW",           m.get("lr",        0)))

        phase = "cooldown" if getattr(pl_module, '_in_cooldown', False) else "distill"

        print(
            f"[{self._ts()}] step {step:4d}/{self._max_steps}"
            f" | loss {loss:.4f}"
            f" (hard {loss_hard:.4f}  kd {loss_kd:.4f}  feat {loss_feat:.4f})"
            f" | lr {lr:.2e}"
            f" | {phase}"
            f" | ETA {eta_s/60:.0f}m",
            flush=True,
        )

    @pl.utilities.rank_zero_only
    def on_validation_epoch_end(self, trainer, pl_module):
        if trainer.sanity_checking:
            return
        step    = trainer.global_step
        elapsed = time.time() - (self._t_start or time.time())
        m       = trainer.callback_metrics
        val_ap  = float(m.get("val_ap",        0))
        recall  = float(m.get("val_recall",    0))
        prec    = float(m.get("val_precision", 0))

        print(
            f"\n[{self._ts()}] ── VAL  step {step:4d}/{self._max_steps}"
            f" | val_ap {val_ap:.4f}"
            f"  recall {recall:.4f}"
            f"  prec {prec:.4f}"
            f"  elapsed {elapsed/60:.0f}m ──\n",
            flush=True,
        )


# =============================================================================
# MAIN
# =============================================================================

def setup_callbacks(config: FullDistillationConfig, run_name: str):
    """Setup training callbacks."""
    from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor, EarlyStopping

    callbacks = []

    checkpoint_dir = config.logging.output_dir or os.path.join("training/logs/experiments", run_name)
    callbacks.append(ModelCheckpoint(
        dirpath=checkpoint_dir,
        filename=f"{run_name}-{{step}}-{{val_ap:.4f}}-{{val_recall:.4f}}-{{val_precision:.4f}}",
        monitor="val_ap",
        mode="max",
        save_top_k=3,
        save_last=True,
    ))

    callbacks.append(LearningRateMonitor(logging_interval="step"))
    callbacks.append(DistillationLogCallback(
        max_steps=config.training.max_steps,
        log_every_n_steps=config.training.log_every_n_steps,
    ))

    if not config.training.disable_early_stopping:
        callbacks.append(EarlyStopping(
            monitor=config.training.early_stopping_metric,
            patience=config.training.early_stopping_patience,
            mode=config.training.early_stopping_mode,
            min_delta=config.training.early_stopping_min_delta,
        ))

    return callbacks


def main():
    """Main training function."""
    import argparse

    parser = argparse.ArgumentParser(description="Knowledge Distillation (V-JEPA2 -> VideoMAEv2)")
    parser.add_argument("--config", type=str, required=True)
    parser.add_argument("--resume", type=str, default=None, help="Path to checkpoint to resume from")
    args = parser.parse_args()

    config = FullDistillationConfig.load(args.config)
    pl.seed_everything(config.system.seed)

    from badas.training.utils import generate_unique_run_name
    run_name = config.logging.run_name or generate_unique_run_name()
    print(f"\nRun: {run_name}")

    # Build output dir (same pattern as regular training)
    base_dir = config.logging.output_dir or "training/logs/experiments"
    output_dir = os.path.join(base_dir, run_name)
    config.logging.output_dir = output_dir
    os.makedirs(output_dir, exist_ok=True)
    print(f"Output directory: {output_dir}")

    from badas.training.lightning_data_module import VideoDataModule

    data_module = VideoDataModule(config=config.to_student_config(), return_path=True)

    model = DistillationModule(config)
    callbacks = setup_callbacks(config, run_name)  # uses config.logging.output_dir set above

    if config.logging.use_wandb:
        from pytorch_lightning.loggers import WandbLogger
        logger = WandbLogger(
            project=config.logging.wandb_project,
            entity=config.logging.wandb_entity,
            name=run_name,
            tags=config.logging.wandb_tags + ["distillation", "videomaev2"],
            save_dir=output_dir,
        )
    else:
        from pytorch_lightning.loggers import TensorBoardLogger
        logger = TensorBoardLogger(save_dir=output_dir, name="")

    trainer = pl.Trainer(
        max_steps=config.training.max_steps,
        val_check_interval=config.training.val_check_interval,
        log_every_n_steps=config.training.log_every_n_steps,
        precision=config.system.precision,
        accelerator=config.system.accelerator,
        devices=config.system.devices,
        strategy="ddp_find_unused_parameters_true",
        callbacks=callbacks,
        logger=logger,
        **config.training.trainer_kwargs,
    )

    config_path = os.path.join(config.logging.output_dir or "logs", run_name, "config.yml")
    os.makedirs(os.path.dirname(config_path), exist_ok=True)
    config.save(config_path)

    if args.resume:
        print(f"\nResuming from checkpoint: {args.resume}")
    print("\nStarting distillation (V-JEPA2 -> VideoMAEv2)...")
    trainer.fit(model, data_module, ckpt_path=args.resume)
    print(f"\nComplete! Best: {callbacks[0].best_model_path}")


if __name__ == "__main__":
    main()
