# src/utils/checkpoint.py
"""Shared checkpoint loading helpers for inference and backtest."""
from __future__ import annotations

import logging
from pathlib import Path
from typing import Dict, List, Tuple, TYPE_CHECKING

import torch
import numpy as np

if TYPE_CHECKING:
    import torch.nn as nn

import config as cfg

logger = logging.getLogger(__name__)

_VALID_KEYS = {
    "model", "scaler", "feat_cols", "val_auc",
    "temperature_long", "temperature_short", "rr_ratio",
    "use_conv", "use_attention", "use_context",
    "hidden_dim", "context_capacity", "context_hidden_dim",
    "input_dim",
}


def load_checkpoints(base_dir: str | Path) -> List[Dict[str, object]]:
    """Load & return a sorted list of checkpoint dicts from *base_dir*."""
    base_path = Path(base_dir)
    if not base_path.is_dir():
        raise FileNotFoundError(f"Models directory not found: {base_path}")

    checkpoints = sorted(base_path.glob("seed_*/best.pt"))
    if not checkpoints:
        raise FileNotFoundError(f"No checkpoints under {base_path} — run --mode train first")

    loaded: List[Dict[str, object]] = []
    for ckpt_path in checkpoints:
        ckpt = torch.load(ckpt_path, map_location=cfg.DEVICE, weights_only=False)
        loaded.append({k: v for k, v in ckpt.items() if k in _VALID_KEYS})
    return loaded


def get_val_aucs(checkpoints: List[Dict[str, object]]) -> np.ndarray:
    """Extract ``val_auc`` from each checkpoint as ndarray."""
    return np.array([c.get('val_auc', 0.5) for c in checkpoints])


def aggregate_probs(
    probs_long: np.ndarray,
    probs_short: np.ndarray,
    val_aucs: np.ndarray,
) -> Tuple[float, float]:
    """Weighted (by AUC) or mean aggregate of ensemble probabilities."""
    if cfg.ENSEMBLE_WEIGHT_BY_AUC and val_aucs.size > 0 and val_aucs.sum() > 0:
        weights = val_aucs / val_aucs.sum()
        return float(np.sum(probs_long * weights)), float(np.sum(probs_short * weights))
    return float(probs_long.mean()), float(probs_short.mean())


def build_model_from_ckpt(ckpt: Dict[str, object]) -> 'nn.Module':
    """Instantiate a model (LSTM or ContextEnhanced) from a checkpoint dict
    and load weights into it.

    Returns the model on ``cfg.DEVICE`` in eval mode.
    """
    use_context = ckpt.get('use_context', cfg.USE_CONTEXT)
    use_conv = ckpt.get('use_conv', True)
    use_attention = ckpt.get('use_attention', True)

    if use_context:
        from src.models.context import ContextEnhancedLSTMModel
        model = ContextEnhancedLSTMModel(
            input_dim=int(ckpt.get('input_dim', 1)),
            hidden_dim=int(ckpt.get('hidden_dim', 48)),
            num_layers=2,
            dropout=0.35,
            use_conv=use_conv,
            use_attention=use_attention,
            use_context=True,
            context_capacity=int(ckpt.get('context_capacity', 64)),
            context_hidden_dim=int(ckpt.get('context_hidden_dim', 64)),
        ).to(cfg.DEVICE)
    else:
        from src.models.lstm import DualHeadLSTMModel
        model = DualHeadLSTMModel(
            input_dim=int(ckpt.get('input_dim', 1)),
            hidden_dim=int(ckpt.get('hidden_dim', 48)),
            use_conv=use_conv,
            use_attention=use_attention,
        ).to(cfg.DEVICE)

    model.load_state_dict(ckpt['model'], strict=False)
    model.eval()
    return model