"""
Checkpoint loading utilities for BADAS.

Loads PyTorch Lightning checkpoints without Lightning dependency.
"""

import os
import torch
from collections import OrderedDict
from typing import Dict, Tuple

from badas.inference.config import Config
from badas.core import create_model, detect_model_type, VideoClassifierBase


def load_model_from_checkpoint(
    checkpoint_path: str,
    map_location: str = 'cpu',
    attn_implementation: str = None,
    model_name_override: str = None,
) -> Tuple[VideoClassifierBase, Config]:
    """
    Load model from Lightning checkpoint without Lightning dependency.

    Supports multiple model types (V-JEPA2, VideoMAE, etc.) by using the
    model registry. The model type is determined from:
    1. 'model_type' field in checkpoint (if present)
    2. Detection from model_name in config

    Args:
        checkpoint_path: Path to .ckpt file
        map_location: Device to load checkpoint to
        model_name_override: HuggingFace model ID or local path to use instead
            of the backbone path stored in the checkpoint.  Use this when the
            original training path is unavailable (e.g. deploying on Jetson).
            The override must refer to the *same architecture* as the training
            backbone — mismatched architectures will produce a dimension error.

    Returns:
        Tuple of (model, config)
    """
    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

    # Load checkpoint
    checkpoint = torch.load(checkpoint_path, map_location=map_location)

    # Extract hyperparameters and create config
    hparams = checkpoint.get('hyper_parameters', {})
    if not hparams:
        raise ValueError("No hyper_parameters found in checkpoint")

    config = Config.load_from_dict(hparams)

    if attn_implementation is not None:
        config.model.attn_implementation = attn_implementation

    if model_name_override is not None:
        original = config.model.model_name
        config.model.model_name = model_name_override
        print(f"   Backbone override: {original!r} → {model_name_override!r}")

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

    # Create model using registry
    model = create_model(config, model_type=model_type)

    # Extract and clean state dict
    state_dict = checkpoint.get('state_dict', {})
    cleaned_state_dict = _clean_state_dict(state_dict)

    # Load weights
    missing, unexpected = model.load_state_dict(cleaned_state_dict, strict=False)

    if unexpected:
        print(f"   Unexpected keys: {len(unexpected)}")

    # Guard against silent dimension mismatches.  If the classifier or temporal
    # processor weights are missing it almost certainly means the model was
    # built with the wrong feature dimension (e.g. wrong backbone override).
    classifier_missing = [k for k in missing if k.startswith('classifier') or k.startswith('temporal_processor')]
    if classifier_missing:
        raise RuntimeError(
            f"Checkpoint weight mismatch: {len(classifier_missing)} classifier/temporal "
            f"keys could not be loaded (e.g. {classifier_missing[0]!r}).  "
            f"This usually means the backbone feature dimension does not match "
            f"what the model was trained with.  "
            f"backbone='{config.model.model_name}'.  "
            f"If you used model_name_override, verify it refers to the same "
            f"architecture (ViT-S → 384-dim, ViT-B → 768-dim, ViT-L → 1024-dim)."
        )

    if missing:
        print(f"   Missing keys: {len(missing)}")

    return model, config


def _patch_config_from_state_dict(config, cleaned_state_dict: Dict[str, torch.Tensor]) -> None:
    """
    Patch config fields that may be stale/wrong in the checkpoint's hyper_parameters.

    Detects the actual temporal method and classifier dimensions from the saved weights
    so the model is built with the correct architecture regardless of what hyper_parameters say.
    """
    tp_keys = {k for k in cleaned_state_dict if k.startswith('temporal_processor.')}

    # ── Detect actual temporal method from weight key signatures ──────────────
    if any('attention.in_proj_weight' in k for k in tp_keys):
        # nn.MultiheadAttention → AttentionProcessor
        actual_method = 'attention'
    elif any('query_vectors' in k for k in tp_keys):
        # Learnable queries → AttentiveProbe
        actual_method = 'probe'
    elif any('lstm' in k or 'rnn' in k for k in tp_keys):
        actual_method = 'lstm'
    else:
        actual_method = None

    if actual_method is not None and config.model.temporal_method != actual_method:
        config.model.temporal_method = actual_method

    # ── Detect num_classes from classifier output layer ────────────────────────
    # Find the last dense weight inside the classifier sub-tree
    import re
    classifier_weights = {
        k: v for k, v in cleaned_state_dict.items()
        if k.startswith('classifier.') and re.search(r'\d+\.weight$', k)
        and v.ndim == 2  # Linear layers only (exclude BatchNorm etc.)
    }
    if classifier_weights:
        last_key = max(classifier_weights, key=lambda k: int(re.search(r'(\d+)\.weight$', k).group(1)))
        actual_num_classes = cleaned_state_dict[last_key].shape[0]
        config.model.num_classes = actual_num_classes
        # If checkpoint was saved with 2-output head, keep it as 2 (don't let
        # _setup_classification_head silently reduce it to 1 for binary classification)
        if actual_num_classes == 2:
            config.model.force_num_classes = True


def _clean_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """
    Clean state dict by removing Lightning prefixes.

    Lightning saves weights with 'model.' prefix that needs to be removed
    for loading into the raw model.
    """
    cleaned = OrderedDict()

    for key, value in state_dict.items():
        # Remove 'model.' prefix if present (Lightning adds this)
        if key.startswith('model.'):
            new_key = key[6:]  # Remove 'model.'
        # Remove 'student.' prefix for distillation checkpoints
        elif key.startswith('student.'):
            new_key = key[8:]  # Remove 'student.'
        else:
            new_key = key

        # Skip optimizer states, training-only, and distillation-only keys
        if any(skip in key for skip in ['optimizer', 'lr_scheduler', 'epoch', 'global_step',
                                         'feature_projector']):
            continue

        cleaned[new_key] = value

    return cleaned


__all__ = [
    'load_model_from_checkpoint',
    '_clean_state_dict',
]
