"""Training loop for the Game Transformer.

Implements:
- Mixed precision training (bfloat16/fp16)
- Cosine LR schedule with warmup
- Gradient clipping
- Multi-objective loss (policy + value + action-value)
- Periodic evaluation and checkpointing
"""

from __future__ import annotations

import json
import math
import time
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

if TYPE_CHECKING:
    from essence_wars.transformer.model import GameTransformer


@dataclass
class TrainConfig:
    """Training hyperparameters."""

    # Data
    data_paths: list[str] | None = None
    max_seq_len: int = 128
    val_fraction: float = 0.1

    # Optimization
    batch_size: int = 256
    epochs: int = 10
    learning_rate: float = 3e-4
    weight_decay: float = 0.1
    warmup_steps: int = 200
    grad_clip: float = 1.0
    betas: tuple[float, float] = (0.9, 0.95)

    # Loss weights
    policy_weight: float = 1.0
    value_weight: float = 1.0
    action_value_weight: float = 0.0  # Enable when AV data available

    # Logging & checkpointing
    log_interval: int = 50
    eval_interval: int = 500
    save_interval: int = 2000
    output_dir: str = "experiments/transformer"

    # Hardware
    device: str = "auto"
    mixed_precision: bool = True
    num_workers: int = 4

    # W&B
    wandb_project: str | None = None  # Set to enable W&B logging
    wandb_run_name: str | None = None


class Trainer:
    """Training loop for the Game Transformer."""

    def __init__(
        self,
        model: GameTransformer,
        config: TrainConfig,
        train_dataset: Dataset[Any],
        val_dataset: Dataset[Any] | None = None,
    ) -> None:
        self.config = config
        self.output_dir = Path(config.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Device
        if config.device == "auto":
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = torch.device(config.device)

        self.model = model.to(self.device)

        # Optimizer
        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=config.learning_rate,
            weight_decay=config.weight_decay,
            betas=config.betas,
        )

        # Data loaders
        self.train_loader = DataLoader(
            train_dataset,
            batch_size=config.batch_size,
            shuffle=True,
            num_workers=config.num_workers,
            pin_memory=True,
            drop_last=True,
        )
        self.val_loader = None
        if val_dataset and len(val_dataset) > 0:  # type: ignore[arg-type]
            self.val_loader = DataLoader(
                val_dataset,
                batch_size=config.batch_size,
                shuffle=False,
                num_workers=config.num_workers,
                pin_memory=True,
            )

        # Mixed precision
        self.use_amp = config.mixed_precision and self.device.type == "cuda"
        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        self.amp_dtype = dtype if self.use_amp else torch.float32
        self.scaler = torch.amp.GradScaler("cuda", enabled=self.use_amp and dtype == torch.float16)  # type: ignore[attr-defined]

        # Training state
        self.global_step = 0
        self.best_val_loss = float("inf")
        self.train_log: list[dict[str, Any]] = []

        # W&B
        self.wandb_run = None
        if config.wandb_project:
            try:
                import wandb

                self.wandb_run = wandb.init(
                    project=config.wandb_project,
                    name=config.wandb_run_name,
                    config={
                        "model": {
                            "n_layers": model.config.n_layers,
                            "n_heads": model.config.n_heads,
                            "embed_dim": model.config.embed_dim,
                            "ffn_dim": model.config.ffn_dim,
                            "vocab_size": model.config.vocab_size,
                            "max_seq_len": model.config.max_seq_len,
                            "params": model.count_parameters(),
                        },
                        "training": {
                            "batch_size": config.batch_size,
                            "epochs": config.epochs,
                            "learning_rate": config.learning_rate,
                            "warmup_steps": config.warmup_steps,
                            "weight_decay": config.weight_decay,
                            "grad_clip": config.grad_clip,
                        },
                        "data": {
                            "train_size": len(train_dataset),  # type: ignore[arg-type]
                            "val_size": len(val_dataset) if val_dataset else 0,  # type: ignore[arg-type]
                            "max_seq_len": config.max_seq_len,
                        },
                        "device": str(self.device),
                        "mixed_precision": self.use_amp,
                    },
                    resume="allow",
                )
            except ImportError:
                print("  wandb not installed, skipping W&B logging")

    def _get_lr(self) -> float:
        """Cosine LR schedule with linear warmup."""
        cfg = self.config
        total_steps = len(self.train_loader) * cfg.epochs

        if self.global_step < cfg.warmup_steps:
            # Linear warmup
            return cfg.learning_rate * self.global_step / max(1, cfg.warmup_steps)
        else:
            # Cosine decay to 10% of peak
            progress = (self.global_step - cfg.warmup_steps) / max(
                1, total_steps - cfg.warmup_steps
            )
            return cfg.learning_rate * (0.1 + 0.9 * (1 + math.cos(math.pi * progress)) / 2)

    def _compute_loss(
        self, outputs: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]
    ) -> tuple[torch.Tensor, dict[str, float]]:
        """Compute multi-objective training loss."""
        cfg = self.config
        losses: dict[str, float] = {}

        # Policy loss: cross-entropy with legal mask
        logits = outputs["policy_logits"]
        legal = batch["legal_mask"].to(self.device)
        masked_logits = logits.masked_fill(legal == 0, float("-inf"))
        policy_loss = F.cross_entropy(masked_logits, batch["action_target"].to(self.device))
        losses["policy"] = policy_loss.item()

        total = cfg.policy_weight * policy_loss

        # Value loss: MSE
        value_loss = F.mse_loss(outputs["value"], batch["value_target"].to(self.device))
        losses["value"] = value_loss.item()
        total = total + cfg.value_weight * value_loss

        # Action-value loss (optional)
        if cfg.action_value_weight > 0:
            # For now, action-value targets would need to be in the batch
            # This is a placeholder for when AV data is available
            pass

        losses["total"] = total.item()
        return total, losses

    @torch.no_grad()
    def evaluate(self) -> dict[str, float]:
        """Run evaluation on validation set."""
        if self.val_loader is None:
            return {}

        self.model.eval()
        total_loss = 0.0
        total_policy_loss = 0.0
        total_value_loss = 0.0
        correct = 0
        total = 0
        n_batches = 0

        for batch in self.val_loader:
            input_ids = batch["input_ids"].to(self.device)
            attn_mask = batch["attention_mask"].to(self.device)

            with torch.amp.autocast("cuda", dtype=self.amp_dtype, enabled=self.use_amp):  # type: ignore[attr-defined]
                outputs = self.model(input_ids, attn_mask)

            _, losses = self._compute_loss(outputs, batch)
            total_loss += losses["total"]
            total_policy_loss += losses["policy"]
            total_value_loss += losses["value"]

            # Policy accuracy (top-1 among legal actions)
            logits = outputs["policy_logits"]
            legal = batch["legal_mask"].to(self.device)
            masked_logits = logits.masked_fill(legal == 0, float("-inf"))
            preds = masked_logits.argmax(dim=-1)
            correct += (preds == batch["action_target"].to(self.device)).sum().item()
            total += len(preds)
            n_batches += 1

        self.model.train()

        return {
            "val_loss": total_loss / max(1, n_batches),
            "val_policy_loss": total_policy_loss / max(1, n_batches),
            "val_value_loss": total_value_loss / max(1, n_batches),
            "val_accuracy": correct / max(1, total),
        }

    def save_checkpoint(self, name: str) -> Path:
        """Save model checkpoint."""
        path = self.output_dir / f"{name}.pt"
        torch.save(
            {
                "model_state_dict": self.model.state_dict(),
                "optimizer_state_dict": self.optimizer.state_dict(),
                "config": self.model.config,
                "global_step": self.global_step,
                "best_val_loss": self.best_val_loss,
            },
            path,
        )
        return path

    def train(self) -> dict[str, Any]:
        """Run full training loop.

        Returns:
            Summary dict with final metrics.
        """
        cfg = self.config
        total_steps = len(self.train_loader) * cfg.epochs
        n_params = self.model.count_parameters()

        print("Training Game Transformer")
        print(f"  Model: {n_params:,} params ({n_params / 1e6:.1f}M)")
        print(f"  Device: {self.device}")
        print(f"  Mixed precision: {self.use_amp} ({self.amp_dtype})")
        print(f"  Train: {len(self.train_loader.dataset):,} positions")  # type: ignore[arg-type]
        if self.val_loader:
            print(f"  Val: {len(self.val_loader.dataset):,} positions")  # type: ignore[arg-type]
        print(f"  Batch size: {cfg.batch_size}")
        print(f"  Epochs: {cfg.epochs}")
        print(f"  Total steps: {total_steps:,}")
        print(f"  LR: {cfg.learning_rate} (cosine, {cfg.warmup_steps} warmup)")
        print(f"  Output: {self.output_dir}")
        print()

        self.model.train()
        start_time = time.time()
        running_loss = 0.0
        running_policy = 0.0
        running_value = 0.0
        log_count = 0

        for epoch in range(cfg.epochs):
            for batch in self.train_loader:
                # Update learning rate
                lr = self._get_lr()
                for param_group in self.optimizer.param_groups:
                    param_group["lr"] = lr

                # Forward
                input_ids = batch["input_ids"].to(self.device)
                attn_mask = batch["attention_mask"].to(self.device)

                with torch.amp.autocast("cuda", dtype=self.amp_dtype, enabled=self.use_amp):  # type: ignore[attr-defined]
                    outputs = self.model(input_ids, attn_mask)
                    loss, losses = self._compute_loss(outputs, batch)

                # Backward
                self.optimizer.zero_grad()
                self.scaler.scale(loss).backward()  # type: ignore[no-untyped-call]
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), cfg.grad_clip)
                self.scaler.step(self.optimizer)
                self.scaler.update()

                # Track running averages
                running_loss += losses["total"]
                running_policy += losses["policy"]
                running_value += losses["value"]
                log_count += 1
                self.global_step += 1

                # Logging
                if self.global_step % cfg.log_interval == 0:
                    avg_loss = running_loss / log_count
                    avg_policy = running_policy / log_count
                    avg_value = running_value / log_count
                    elapsed = time.time() - start_time
                    sps = self.global_step * cfg.batch_size / elapsed

                    entry = {
                        "step": self.global_step,
                        "epoch": epoch + 1,
                        "loss": round(avg_loss, 4),
                        "policy_loss": round(avg_policy, 4),
                        "value_loss": round(avg_value, 4),
                        "lr": lr,
                        "samples_per_sec": round(sps),
                    }
                    self.train_log.append(entry)

                    print(
                        f"  Step {self.global_step:>6}/{total_steps} "
                        f"| E{epoch + 1} "
                        f"| Loss: {avg_loss:.4f} "
                        f"(policy={avg_policy:.3f}, value={avg_value:.3f}) "
                        f"| LR: {lr:.1e} "
                        f"| {sps:,.0f} samp/s"
                    )

                    if self.wandb_run:
                        self.wandb_run.log(entry, step=self.global_step)

                    running_loss = 0.0
                    running_policy = 0.0
                    running_value = 0.0
                    log_count = 0

                # Evaluation
                if self.global_step % cfg.eval_interval == 0 and self.val_loader:
                    val_metrics = self.evaluate()
                    val_loss = val_metrics["val_loss"]
                    val_acc = val_metrics["val_accuracy"]

                    print(
                        f"    -> Val: loss={val_loss:.4f}, "
                        f"accuracy={val_acc:.1%}"
                    )

                    if self.wandb_run:
                        self.wandb_run.log(val_metrics, step=self.global_step)

                    if val_loss < self.best_val_loss:
                        self.best_val_loss = val_loss
                        self.save_checkpoint("best")
                        print(f"    -> New best! Saved to {self.output_dir}/best.pt")

                    self.train_log[-1].update(val_metrics)

                # Periodic checkpoint
                if self.global_step % cfg.save_interval == 0:
                    self.save_checkpoint(f"step_{self.global_step}")

        # Final evaluation
        total_time = time.time() - start_time
        final_metrics = self.evaluate() if self.val_loader else {}

        # Save final model
        self.save_checkpoint("final")

        # Save training log
        log_path = self.output_dir / "training_log.json"
        with log_path.open("w") as f:
            json.dump(self.train_log, f, indent=2)

        summary = {
            "total_steps": self.global_step,
            "total_time_seconds": round(total_time),
            "best_val_loss": round(self.best_val_loss, 4),
            "final_metrics": final_metrics,
            "model_params": n_params,
            "output_dir": str(self.output_dir),
        }

        print("\nTraining complete!")
        print(f"  Total steps: {self.global_step:,}")
        print(f"  Total time: {total_time:.0f}s ({total_time / 3600:.1f}h)")
        print(f"  Best val loss: {self.best_val_loss:.4f}")
        if final_metrics:
            print(f"  Final val accuracy: {final_metrics.get('val_accuracy', 0):.1%}")
        print(f"  Output: {self.output_dir}")

        summary_path = self.output_dir / "training_summary.json"
        with summary_path.open("w") as f:
            json.dump(summary, f, indent=2)

        if self.wandb_run:
            self.wandb_run.log({"final/" + k: v for k, v in final_metrics.items()})
            self.wandb_run.summary.update(summary)
            self.wandb_run.finish()

        return summary
