"""
Commander AI Lab — Supervised Policy Trainer
═════════════════════════════════════════════
Trains the PolicyNetwork on labeled decision data using
cross-entropy loss with class-weight balancing.

Features:
  - Class weighting to handle imbalanced action distribution
  - Cosine annealing LR schedule with warmup
  - Early stopping on validation loss
  - Per-class accuracy tracking
  - Periodic checkpoint saving (best + every N epochs)
  - CUDA/MPS/CPU auto-detection

Usage:
    python -m ml.training.trainer --data-dir ml/models --epochs 50
    python -m ml.training.trainer --data-dir ml/models --lr 5e-4 --batch-size 128
"""

import argparse
import json
import logging
import os
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np

# Conditional torch import
try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from torch.utils.data import DataLoader, TensorDataset
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

project_root = str(Path(__file__).parent.parent.parent)
if project_root not in sys.path:
    sys.path.insert(0, project_root)

from ml.config.scope import (
    NUM_ACTIONS, STATE_DIMS, IDX_TO_ACTION, TRAINING_CONFIG,
)
from ml.training.policy_network import (
    PolicyNetwork, save_checkpoint, load_checkpoint,
)

logger = logging.getLogger("commander_ai_lab.ml.training")


# ══════════════════════════════════════════════════════════
# Dataset Loading
# ══════════════════════════════════════════════════════════

def load_npz_split(path: str) -> Tuple[np.ndarray, np.ndarray]:
    """Load states and labels from an NPZ file."""
    data = np.load(path, allow_pickle=True)
    states = data["states"].astype(np.float32)
    labels = data["labels"].astype(np.int64)
    return states, labels


def make_dataloader(
    states: np.ndarray,
    labels: np.ndarray,
    batch_size: int,
    shuffle: bool = True,
    device: str = "cpu",
) -> DataLoader:
    """Create a PyTorch DataLoader from numpy arrays."""
    state_tensor = torch.from_numpy(states)
    label_tensor = torch.from_numpy(labels)
    dataset = TensorDataset(state_tensor, label_tensor)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        pin_memory=(device != "cpu"),
        drop_last=False,
    )


def compute_class_weights(labels: np.ndarray, num_classes: int) -> torch.Tensor:
    """
    Compute inverse-frequency class weights for cross-entropy.

    Uses smoothed inverse frequency:
        w_c = N / (num_classes * count_c)
    Capped at 10× to prevent extreme weights on rare classes.
    """
    counts = np.bincount(labels, minlength=num_classes).astype(np.float64)
    # Avoid division by zero for classes with no samples
    counts = np.maximum(counts, 1.0)
    n = len(labels)
    weights = n / (num_classes * counts)
    # Cap extreme weights
    weights = np.minimum(weights, 10.0)
    return torch.from_numpy(weights).float()


# ══════════════════════════════════════════════════════════
# Training Loop
# ══════════════════════════════════════════════════════════

class SupervisedTrainer:
    """
    Trains PolicyNetwork with cross-entropy on expert decision data.

    Features:
      - Class-weighted loss for imbalanced actions
      - Cosine annealing with linear warmup
      - Early stopping on validation loss
      - Best-model and periodic checkpointing
      - Detailed per-epoch metrics
    """

    def __init__(
        self,
        model: PolicyNetwork,
        device: str = "cpu",
        learning_rate: float = None,
        batch_size: int = None,
        epochs: int = None,
        patience: int = None,
        checkpoint_dir: str = None,
        warmup_epochs: int = 3,
        save_every: int = 10,
    ):
        self.model = model.to(device)
        self.device = device
        self.lr = learning_rate or TRAINING_CONFIG.learning_rate
        self.batch_size = batch_size or TRAINING_CONFIG.batch_size
        self.epochs = epochs or TRAINING_CONFIG.epochs
        self.patience = patience or TRAINING_CONFIG.early_stop_patience
        self.checkpoint_dir = checkpoint_dir or TRAINING_CONFIG.checkpoint_dir
        self.warmup_epochs = warmup_epochs
        self.save_every = save_every

        # Will be initialized in train()
        self.optimizer = None
        self.scheduler = None
        self.criterion = None
        self.history: List[Dict] = []

    def train(
        self,
        train_states: np.ndarray,
        train_labels: np.ndarray,
        val_states: np.ndarray,
        val_labels: np.ndarray,
    ) -> Dict:
        """
        Run the full training loop.

        Returns:
            Dict with training summary (best metrics, epoch count, etc.)
        """
        if not TORCH_AVAILABLE:
            raise RuntimeError("PyTorch not installed. Run: pip install torch")

        logger.info("=" * 60)
        logger.info("  Commander AI Lab — Supervised Policy Training")
        logger.info("=" * 60)
        logger.info("  Device:        %s", self.device)
        logger.info("  Train samples: %d", len(train_states))
        logger.info("  Val samples:   %d", len(val_states))
        logger.info("  Batch size:    %d", self.batch_size)
        logger.info("  Learning rate: %s", self.lr)
        logger.info("  Max epochs:    %d", self.epochs)
        logger.info("  Patience:      %d", self.patience)
        logger.info("  Warmup epochs: %d", self.warmup_epochs)
        logger.info("=" * 60)

        # Data loaders
        train_loader = make_dataloader(
            train_states, train_labels, self.batch_size, shuffle=True, device=self.device
        )
        val_loader = make_dataloader(
            val_states, val_labels, self.batch_size, shuffle=False, device=self.device
        )

        # Class-weighted cross-entropy
        class_weights = compute_class_weights(train_labels, NUM_ACTIONS).to(self.device)
        self.criterion = nn.CrossEntropyLoss(weight=class_weights)

        logger.info("  Class weights: %s", 
                     {IDX_TO_ACTION[i].value: f"{w:.2f}" 
                      for i, w in enumerate(class_weights.cpu().numpy())})

        # Optimizer
        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=self.lr,
            weight_decay=1e-4,
        )

        # Cosine annealing with warmup
        # During warmup, LR ramps from lr/10 to lr
        # After warmup, cosine decay to lr/100
        total_steps = self.epochs * len(train_loader)
        warmup_steps = self.warmup_epochs * len(train_loader)

        def lr_lambda(step):
            if step < warmup_steps:
                return 0.1 + 0.9 * (step / max(warmup_steps, 1))
            progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
            return 0.01 + 0.99 * 0.5 * (1 + np.cos(np.pi * progress))

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer, lr_lambda
        )

        # Training state
        best_val_loss = float("inf")
        best_val_acc = 0.0
        best_epoch = 0
        epochs_without_improvement = 0

        os.makedirs(self.checkpoint_dir, exist_ok=True)

        for epoch in range(1, self.epochs + 1):
            t_start = time.time()

            # Train one epoch
            train_loss, train_acc, train_per_class = self._train_epoch(
                train_loader, epoch
            )

            # Validate
            val_loss, val_acc, val_per_class = self._validate(val_loader)

            elapsed = time.time() - t_start
            current_lr = self.optimizer.param_groups[0]["lr"]

            # Record history
            epoch_metrics = {
                "epoch": epoch,
                "train_loss": train_loss,
                "train_acc": train_acc,
                "val_loss": val_loss,
                "val_acc": val_acc,
                "lr": current_lr,
                "time_s": elapsed,
                "train_per_class": train_per_class,
                "val_per_class": val_per_class,
            }
            self.history.append(epoch_metrics)

            # Log
            logger.info(
                "  Epoch %3d/%d | train_loss=%.4f train_acc=%.3f | "
                "val_loss=%.4f val_acc=%.3f | lr=%.2e | %.1fs",
                epoch, self.epochs,
                train_loss, train_acc,
                val_loss, val_acc,
                current_lr, elapsed,
            )

            # Check for improvement
            improved = False
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                best_val_acc = val_acc
                best_epoch = epoch
                epochs_without_improvement = 0
                improved = True

                # Save best model
                save_checkpoint(
                    self.model, self.optimizer, epoch,
                    {"val_loss": val_loss, "val_acc": val_acc,
                     "train_loss": train_loss, "train_acc": train_acc},
                    os.path.join(self.checkpoint_dir, "best_policy.pt"),
                )
            else:
                epochs_without_improvement += 1

            # Periodic checkpoint
            if epoch % self.save_every == 0:
                save_checkpoint(
                    self.model, self.optimizer, epoch,
                    {"val_loss": val_loss, "val_acc": val_acc},
                    os.path.join(self.checkpoint_dir, f"policy_epoch_{epoch:03d}.pt"),
                )

            # Early stopping
            if epochs_without_improvement >= self.patience:
                logger.info(
                    "  Early stopping at epoch %d (no improvement for %d epochs)",
                    epoch, self.patience,
                )
                break

        # Final summary
        logger.info("")
        logger.info("=" * 60)
        logger.info("  Training Complete")
        logger.info("=" * 60)
        logger.info("  Best epoch:     %d", best_epoch)
        logger.info("  Best val_loss:  %.4f", best_val_loss)
        logger.info("  Best val_acc:   %.3f", best_val_acc)
        logger.info("  Total epochs:   %d", len(self.history))
        logger.info("  Checkpoint:     %s", os.path.join(self.checkpoint_dir, "best_policy.pt"))
        logger.info("=" * 60)

        # Save training history
        history_path = os.path.join(self.checkpoint_dir, "training_history.json")
        # Convert per-class dicts for JSON serialization
        serializable_history = []
        for h in self.history:
            entry = {k: v for k, v in h.items()
                     if k not in ("train_per_class", "val_per_class")}
            entry["train_per_class"] = {
                str(k): v for k, v in h["train_per_class"].items()
            }
            entry["val_per_class"] = {
                str(k): v for k, v in h["val_per_class"].items()
            }
            serializable_history.append(entry)

        with open(history_path, "w") as f:
            json.dump(serializable_history, f, indent=2)
        logger.info("  History saved:  %s", history_path)

        return {
            "best_epoch": best_epoch,
            "best_val_loss": best_val_loss,
            "best_val_acc": best_val_acc,
            "total_epochs": len(self.history),
            "checkpoint_path": os.path.join(self.checkpoint_dir, "best_policy.pt"),
        }

    def _train_epoch(
        self,
        loader: DataLoader,
        epoch: int,
    ) -> Tuple[float, float, Dict[int, float]]:
        """Train for one epoch. Returns (loss, accuracy, per_class_acc)."""
        self.model.train()
        total_loss = 0.0
        correct = 0
        total = 0
        per_class_correct = np.zeros(NUM_ACTIONS)
        per_class_total = np.zeros(NUM_ACTIONS)

        for batch_idx, (states, labels) in enumerate(loader):
            states = states.to(self.device)
            labels = labels.to(self.device)

            # Forward
            logits = self.model(states)
            loss = self.criterion(logits, labels)

            # Backward
            self.optimizer.zero_grad()
            loss.backward()

            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)

            self.optimizer.step()
            self.scheduler.step()

            # Metrics
            total_loss += loss.item() * len(labels)
            preds = logits.argmax(dim=-1)
            correct += (preds == labels).sum().item()
            total += len(labels)

            # Per-class tracking
            for c in range(NUM_ACTIONS):
                mask = labels == c
                per_class_total[c] += mask.sum().item()
                per_class_correct[c] += (preds[mask] == c).sum().item()

        avg_loss = total_loss / max(total, 1)
        accuracy = correct / max(total, 1)

        per_class_acc = {}
        for c in range(NUM_ACTIONS):
            if per_class_total[c] > 0:
                per_class_acc[c] = per_class_correct[c] / per_class_total[c]
            else:
                per_class_acc[c] = 0.0

        return avg_loss, accuracy, per_class_acc

    @torch.no_grad()
    def _validate(
        self,
        loader: DataLoader,
    ) -> Tuple[float, float, Dict[int, float]]:
        """Validate on held-out data. Returns (loss, accuracy, per_class_acc)."""
        self.model.eval()
        total_loss = 0.0
        correct = 0
        total = 0
        per_class_correct = np.zeros(NUM_ACTIONS)
        per_class_total = np.zeros(NUM_ACTIONS)

        for states, labels in loader:
            states = states.to(self.device)
            labels = labels.to(self.device)

            logits = self.model(states)
            loss = self.criterion(logits, labels)

            total_loss += loss.item() * len(labels)
            preds = logits.argmax(dim=-1)
            correct += (preds == labels).sum().item()
            total += len(labels)

            for c in range(NUM_ACTIONS):
                mask = labels == c
                per_class_total[c] += mask.sum().item()
                per_class_correct[c] += (preds[mask] == c).sum().item()

        avg_loss = total_loss / max(total, 1)
        accuracy = correct / max(total, 1)

        per_class_acc = {}
        for c in range(NUM_ACTIONS):
            if per_class_total[c] > 0:
                per_class_acc[c] = per_class_correct[c] / per_class_total[c]
            else:
                per_class_acc[c] = 0.0

        return avg_loss, accuracy, per_class_acc


# ══════════════════════════════════════════════════════════
# Evaluation
# ══════════════════════════════════════════════════════════

def evaluate_model(
    model: PolicyNetwork,
    test_states: np.ndarray,
    test_labels: np.ndarray,
    device: str = "cpu",
) -> Dict:
    """
    Full evaluation on test set with detailed metrics.

    Returns:
        Dict with overall accuracy, per-class accuracy, confusion matrix
    """
    model.eval()
    model.to(device)

    states_t = torch.from_numpy(test_states.astype(np.float32)).to(device)
    labels_t = torch.from_numpy(test_labels.astype(np.int64)).to(device)

    with torch.no_grad():
        logits = model(states_t)
        preds = logits.argmax(dim=-1)

    preds_np = preds.cpu().numpy()
    labels_np = test_labels

    # Overall accuracy
    overall_acc = (preds_np == labels_np).mean()

    # Per-class metrics
    per_class = {}
    for c in range(NUM_ACTIONS):
        mask = labels_np == c
        total_c = mask.sum()
        correct_c = (preds_np[mask] == c).sum() if total_c > 0 else 0
        pred_as_c = (preds_np == c).sum()

        precision = correct_c / max(pred_as_c, 1)
        recall = correct_c / max(total_c, 1) 
        f1 = 2 * precision * recall / max(precision + recall, 1e-8)

        action_name = IDX_TO_ACTION[c].value
        per_class[action_name] = {
            "support": int(total_c),
            "correct": int(correct_c),
            "predicted": int(pred_as_c),
            "precision": round(float(precision), 4),
            "recall": round(float(recall), 4),
            "f1": round(float(f1), 4),
        }

    # Confusion matrix
    confusion = np.zeros((NUM_ACTIONS, NUM_ACTIONS), dtype=int)
    for true_c, pred_c in zip(labels_np, preds_np):
        confusion[true_c][pred_c] += 1

    return {
        "overall_accuracy": round(float(overall_acc), 4),
        "total_samples": len(labels_np),
        "per_class": per_class,
        "confusion_matrix": confusion.tolist(),
    }


def print_evaluation(results: Dict):
    """Pretty-print evaluation results."""
    logger.info("")
    logger.info("=" * 65)
    logger.info("  Test Set Evaluation")
    logger.info("=" * 65)
    logger.info(f"  Overall Accuracy: {results['overall_accuracy']:.1%}")
    logger.info(f"  Total Samples:    {results['total_samples']:,}")
    logger.info("")
    logger.info(f"  {'Action':<20} {'Support':>8} {'Prec':>7} {'Recall':>7} {'F1':>7}")
    logger.info("  " + "-" * 55)

    for action_name, m in results["per_class"].items():
        logger.info(
            f"  {action_name:<20} {m['support']:>8} "
            f"{m['precision']:>7.3f} {m['recall']:>7.3f} {m['f1']:>7.3f}"
        )

    logger.info("  " + "-" * 55)

    # Confusion matrix
    logger.info("")
    logger.info("  Confusion Matrix (rows=true, cols=predicted):")
    header = "  " + " " * 18
    for c in range(NUM_ACTIONS):
        short_name = IDX_TO_ACTION[c].value[:6]
        header += f" {short_name:>6}"
    logger.info(header)

    for r in range(NUM_ACTIONS):
        row_name = IDX_TO_ACTION[r].value[:16]
        row_str = f"  {row_name:<18}"
        for c in range(NUM_ACTIONS):
            row_str += f" {results['confusion_matrix'][r][c]:>6}"
        logger.info(row_str)
    logger.info("")


# ══════════════════════════════════════════════════════════
# Device Selection
# ══════════════════════════════════════════════════════════

def get_best_device() -> str:
    """Auto-detect the best available compute device."""
    if not TORCH_AVAILABLE:
        return "cpu"
    if torch.cuda.is_available():
        name = torch.cuda.get_device_name(0)
        mem = torch.cuda.get_device_properties(0).total_mem / (1024**3)
        logger.info("  CUDA device: %s (%.1f GB)", name, mem)
        return "cuda"
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        logger.info("  Using Apple MPS")
        return "mps"
    logger.info("  Using CPU (no GPU detected)")
    return "cpu"


# ══════════════════════════════════════════════════════════
# CLI Entry Point
# ══════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(
        description="Commander AI Lab — Train Policy Network"
    )
    parser.add_argument(
        "--data-dir", default="ml/models",
        help="Directory containing dataset-train.npz / dataset-val.npz / dataset-test.npz",
    )
    parser.add_argument(
        "--epochs", type=int, default=None,
        help=f"Max training epochs (default: {TRAINING_CONFIG.epochs})",
    )
    parser.add_argument(
        "--batch-size", type=int, default=None,
        help=f"Batch size (default: {TRAINING_CONFIG.batch_size})",
    )
    parser.add_argument(
        "--lr", type=float, default=None,
        help=f"Learning rate (default: {TRAINING_CONFIG.learning_rate})",
    )
    parser.add_argument(
        "--patience", type=int, default=None,
        help=f"Early stopping patience (default: {TRAINING_CONFIG.early_stop_patience})",
    )
    parser.add_argument(
        "--checkpoint-dir", default=None,
        help=f"Checkpoint directory (default: {TRAINING_CONFIG.checkpoint_dir})",
    )
    parser.add_argument(
        "--device", default=None,
        help="Compute device: cuda, mps, or cpu (default: auto-detect)",
    )
    parser.add_argument(
        "--eval-only", action="store_true",
        help="Only evaluate an existing checkpoint on test data",
    )
    parser.add_argument(
        "--checkpoint", default=None,
        help="Path to checkpoint for --eval-only (default: best_policy.pt)",
    )
    args = parser.parse_args()

    if not TORCH_AVAILABLE:
        logger.info("PyTorch not installed. Install with:")
        logger.info("  pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121")
        sys.exit(1)

    device = args.device or get_best_device()
    data_dir = args.data_dir
    ckpt_dir = args.checkpoint_dir or TRAINING_CONFIG.checkpoint_dir

    # Load data
    train_path = os.path.join(data_dir, "dataset-train.npz")
    val_path = os.path.join(data_dir, "dataset-val.npz")
    test_path = os.path.join(data_dir, "dataset-test.npz")

    if args.eval_only:
        # Evaluation only
        if not os.path.exists(test_path):
            logger.info(f"Test data not found: {test_path}")
            logger.info("Run: python -m ml.scripts.ml_cli build")
            sys.exit(1)

        ckpt_path = args.checkpoint or os.path.join(ckpt_dir, "best_policy.pt")
        if not os.path.exists(ckpt_path):
            logger.info(f"Checkpoint not found: {ckpt_path}")
            sys.exit(1)

        test_states, test_labels = load_npz_split(test_path)
        model, checkpoint = load_checkpoint(ckpt_path, device)
        logger.info("  Loaded checkpoint from epoch %d", checkpoint.get("epoch", "?"))

        results = evaluate_model(model, test_states, test_labels, device)
        print_evaluation(results)

        # Save evaluation results
        eval_path = os.path.join(ckpt_dir, "eval_results.json")
        with open(eval_path, "w") as f:
            json.dump(results, f, indent=2)
        logger.info(f"  Results saved: {eval_path}")
        return

    # Training mode
    for path, label in [(train_path, "train"), (val_path, "val")]:
        if not os.path.exists(path):
            logger.info(f"{label} data not found: {path}")
            logger.info("Build the dataset first:")
            logger.info("  python -m ml.scripts.ml_cli build --results-dir results/")
            sys.exit(1)

    train_states, train_labels = load_npz_split(train_path)
    val_states, val_labels = load_npz_split(val_path)

    logger.info("  Train: %d samples, %d features", *train_states.shape)
    logger.info("  Val:   %d samples, %d features", *val_states.shape)

    # Verify dimensions match expected state size
    actual_dim = train_states.shape[1]
    expected_dim = STATE_DIMS.total_state_dim
    if actual_dim != expected_dim:
        logger.warning(
            "  State dimension mismatch: data=%d, config=%d. Using data dimension.",
            actual_dim, expected_dim,
        )

    # Create model
    model = PolicyNetwork(input_dim=actual_dim)
    param_count = sum(p.numel() for p in model.parameters())
    logger.info("  Model parameters: %s", f"{param_count:,}")

    # Train
    trainer = SupervisedTrainer(
        model=model,
        device=device,
        learning_rate=args.lr,
        batch_size=args.batch_size,
        epochs=args.epochs,
        patience=args.patience,
        checkpoint_dir=ckpt_dir,
    )

    summary = trainer.train(train_states, train_labels, val_states, val_labels)

    # Evaluate on test set if available
    if os.path.exists(test_path):
        logger.info("")
        logger.info("  Running test set evaluation...")
        test_states, test_labels = load_npz_split(test_path)

        # Reload best checkpoint for evaluation
        best_model, _ = load_checkpoint(summary["checkpoint_path"], device)
        results = evaluate_model(best_model, test_states, test_labels, device)
        print_evaluation(results)

        eval_path = os.path.join(ckpt_dir, "eval_results.json")
        with open(eval_path, "w") as f:
            json.dump(results, f, indent=2)
        logger.info("  Eval results saved: %s", eval_path)
    else:
        logger.info("  No test data found at %s — skipping evaluation.", test_path)


if __name__ == "__main__":
    main()