"""AlphaZero training loop.

Provides the main training loop for AlphaZero:
1. Self-play: Generate games using current network + MCTS
2. Store (state, mcts_policy, outcome) in replay buffer
3. Train network on batches from replay buffer
4. Repeat
"""

from __future__ import annotations

import time
from typing import TYPE_CHECKING, Any

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

from essence_wars._core import STATE_TENSOR_SIZE, PyGame
from essence_wars.agents.networks import AlphaZeroNetwork
from essence_wars.env import EssenceWarsEnv

from ..utils import RunningMeanStd
from .buffer import DualReplayBuffer, ReplayBuffer, ReplaySample
from .config import AlphaZeroConfig
from .mcts import NeuralMCTS

if TYPE_CHECKING:
    from numpy.typing import NDArray


class AlphaZeroTrainer:
    """AlphaZero training loop.

    Training process:
    1. Self-play: Generate games using current network + MCTS
    2. Store (state, mcts_policy, outcome) in replay buffer
    3. Train network on batches from replay buffer
    4. Repeat

    Args:
        config: AlphaZeroConfig with training parameters
        network: Optional pre-built network
        writer: Optional TensorBoard SummaryWriter

    Example:
        trainer = AlphaZeroTrainer(AlphaZeroConfig())
        trainer.train()
        trainer.save("checkpoints/alphazero_model.pt")
    """

    def __init__(
        self,
        config: AlphaZeroConfig | None = None,
        network: AlphaZeroNetwork | None = None,
        writer: Any = None,
    ) -> None:
        self.config = config or AlphaZeroConfig()
        self.device = torch.device(self.config.device)

        # Create network
        if network is not None:
            self.network = network.to(self.device)
        else:
            self.network = AlphaZeroNetwork(
                hidden_dim=self.config.hidden_dim,
                num_blocks=self.config.num_blocks,
            ).to(self.device)

        # Optimizer
        self.optimizer = optim.Adam(
            self.network.parameters(),
            lr=self.config.learning_rate,
            weight_decay=self.config.weight_decay,
        )

        # Replay buffer (single buffer for standard training)
        self.replay_buffer = ReplayBuffer(capacity=self.config.replay_buffer_size)

        # Dual replay buffer (for BC warm-start, initialized separately)
        self.dual_buffer: DualReplayBuffer | None = None

        # Observation normalizer
        self.obs_normalizer = RunningMeanStd((STATE_TENSOR_SIZE,)) if self.config.normalize_obs else None

        # MCTS
        self.mcts = NeuralMCTS(
            network=self.network,
            config=self.config,
            obs_normalizer=self.obs_normalizer,
        )

        # TensorBoard writer
        self.writer = writer

        # Training state
        self.iteration = 0
        self.total_games = 0
        self.start_time: float | None = None
        self.save_dir: str | None = None  # Will be set by train()

    def init_dual_buffer(
        self,
        bc_samples: list[ReplaySample],
        bc_ratio: float | None = None,
    ) -> None:
        """Initialize dual replay buffer with BC data for warm-start training.

        This enables the two-buffer approach that maintains a fixed BC ratio
        throughout training, preventing the "delayed catastrophic forgetting"
        that occurs when BC samples get diluted by self-play data.

        Args:
            bc_samples: List of (obs, mask, policy, value) tuples from BC dataset
            bc_ratio: Override config.bc_ratio if specified

        Example:
            # Load BC data
            bc_dataset = MCTSDataset("data/mcts_10k.jsonl.gz")
            bc_samples = [(s.state_tensor, s.action_mask, s.mcts_policy, s.value_target)
                          for s in bc_dataset.samples]

            # Initialize dual buffer
            trainer.init_dual_buffer(bc_samples, bc_ratio=0.7)
        """
        ratio = bc_ratio if bc_ratio is not None else self.config.bc_ratio

        self.dual_buffer = DualReplayBuffer(
            bc_ratio=ratio,
            bc_capacity=self.config.replay_buffer_size,
            selfplay_capacity=self.config.replay_buffer_size,
        )

        # Load BC samples
        for obs, mask, policy, value in bc_samples:
            self.dual_buffer.add_bc(obs, mask, policy, value)
            # Also update observation normalizer
            if self.obs_normalizer is not None:
                self.obs_normalizer.update(obs)

        print(f"  Initialized dual buffer with {self.dual_buffer.bc_count:,} BC samples")
        print(f"  BC ratio: {ratio:.0%} (fixed throughout training)")

    @property
    def using_dual_buffer(self) -> bool:
        """Check if dual buffer mode is active."""
        return self.dual_buffer is not None

    def self_play_game(
        self, seed: int | None = None
    ) -> tuple[list[NDArray[np.float32]], list[NDArray[np.float32]], list[NDArray[np.float32]], float]:
        """Play one self-play game using MCTS.

        Args:
            seed: Optional random seed for game

        Returns:
            (observations, masks, policies, outcome) tuple
        """
        game = PyGame()
        game.reset(seed if seed is not None else np.random.randint(0, 2**31))

        observations: list[NDArray[np.float32]] = []
        masks: list[NDArray[np.float32]] = []
        policies: list[NDArray[np.float32]] = []
        move_count = 0

        while not game.is_done():
            obs = game.observe()
            mask = game.action_mask()

            # Update normalizer with raw observation
            if self.obs_normalizer is not None:
                self.obs_normalizer.update(obs)

            # MCTS search (batched for GPU efficiency)
            action_probs = self.mcts.search_batched(
                game,
                batch_size=self.config.mcts_batch_size,
                add_noise=True,
                virtual_loss=self.config.mcts_virtual_loss,
            )

            # Store training data
            observations.append(obs)
            masks.append(mask)
            policies.append(action_probs)

            # Select action with temperature
            if move_count < self.config.temperature_moves:
                # Sample proportionally to visit counts
                action = np.random.choice(422, p=action_probs)
            else:
                # Low temperature: almost deterministic
                temp_probs = action_probs ** (1.0 / self.config.temperature_final)
                temp_probs = temp_probs / temp_probs.sum()
                action = np.random.choice(422, p=temp_probs)

            game.step(action)
            move_count += 1

        # Get game outcome from player 0's perspective
        outcome = game.get_reward(0)  # +1 if player 0 won, -1 if lost

        return observations, masks, policies, outcome

    def generate_self_play_games(self, num_games: int) -> None:
        """Generate multiple self-play games and add to replay buffer.

        Args:
            num_games: Number of games to generate
        """
        for i in range(num_games):
            observations, masks, policies, outcome = self.self_play_game(
                seed=self.total_games + i
            )
            # Add to appropriate buffer
            if self.dual_buffer is not None:
                self.dual_buffer.add_selfplay_game(observations, masks, policies, outcome)
            else:
                self.replay_buffer.add_game(observations, masks, policies, outcome)
            self.total_games += 1

    def train_step(self) -> dict[str, float]:
        """Perform one training step on a batch from replay buffer.

        Returns:
            Dictionary with loss values
        """
        self.network.train()

        # Sample batch from appropriate buffer
        if self.dual_buffer is not None:
            obs, masks, policy_targets, value_targets = self.dual_buffer.sample(
                self.config.batch_size
            )
        else:
            obs, masks, policy_targets, value_targets = self.replay_buffer.sample(
                self.config.batch_size
            )

        # Normalize observations
        if self.obs_normalizer is not None:
            obs = self.obs_normalizer.normalize(obs).astype(np.float32)

        # Convert to tensors
        obs_t = torch.tensor(obs, dtype=torch.float32, device=self.device)
        masks_t = torch.tensor(masks > 0, dtype=torch.bool, device=self.device)
        policy_targets_t = torch.tensor(policy_targets, dtype=torch.float32, device=self.device)
        value_targets_t = torch.tensor(value_targets, dtype=torch.float32, device=self.device)

        # Forward pass
        policy_logits, value_pred = self.network(obs_t, masks_t)

        # Policy loss: cross-entropy with MCTS policy targets
        policy_log_probs = torch.log_softmax(policy_logits, dim=-1)
        policy_loss = -(policy_targets_t * policy_log_probs).sum(dim=-1).mean()

        # Value loss: MSE with game outcomes
        value_loss = nn.functional.mse_loss(value_pred, value_targets_t)

        # Total loss
        total_loss = policy_loss + value_loss

        # Optimize
        self.optimizer.zero_grad()
        total_loss.backward()
        # Gradient clipping
        nn.utils.clip_grad_norm_(self.network.parameters(), max_norm=1.0)
        self.optimizer.step()

        return {
            "policy_loss": policy_loss.item(),
            "value_loss": value_loss.item(),
            "total_loss": total_loss.item(),
        }

    def train(self, num_iterations: int | None = None, save_dir: str | None = None) -> dict[str, Any]:
        """Main AlphaZero training loop.

        Args:
            num_iterations: Override config num_iterations
            save_dir: Directory to save checkpoints (for periodic saves)

        Returns:
            Training statistics
        """
        num_iterations = num_iterations or self.config.num_iterations
        self.save_dir = save_dir
        self.start_time = time.time()

        print(f"Starting AlphaZero training for {num_iterations} iterations...")
        print(f"  Games per iteration: {self.config.games_per_iteration}")
        print(f"  Simulations per move: {self.config.num_simulations}")
        print(f"  Device: {self.device}")
        if self.dual_buffer is not None:
            print(f"  Mode: Dual buffer (BC ratio: {self.dual_buffer.bc_ratio:.0%})")
            print(f"  BC samples: {self.dual_buffer.bc_count:,}")
        if self.config.checkpoint_interval > 0:
            print(f"  Checkpoint interval: every {self.config.checkpoint_interval} iterations")

        all_losses: list[float] = []
        loss_info: dict[str, float] = {}

        try:
            for iteration in range(1, num_iterations + 1):
                self.iteration = iteration
                iter_start = time.time()

                # Generate self-play games
                print(f"\nIteration {iteration}/{num_iterations}")
                print(f"  Generating {self.config.games_per_iteration} self-play games...")
                self.generate_self_play_games(self.config.games_per_iteration)

                # Print buffer stats
                if self.dual_buffer is not None:
                    stats = self.dual_buffer.stats()
                    print(f"  Buffer: {stats['bc_samples']:,} BC + {stats['selfplay_samples']:,} self-play (sampling {stats['effective_bc_ratio']:.0%} BC)")
                else:
                    print(f"  Replay buffer size: {len(self.replay_buffer)}")

                # Determine if we have enough samples to train
                if self.dual_buffer is not None:
                    # In dual buffer mode, we can train as soon as we have BC data
                    can_train = self.dual_buffer.bc_count >= self.config.min_replay_size
                else:
                    can_train = len(self.replay_buffer) >= self.config.min_replay_size

                # Training
                if can_train:
                    print(f"  Training for {self.config.training_steps_per_iteration} steps...")
                    iter_losses: list[float] = []
                    for _ in range(self.config.training_steps_per_iteration):
                        loss_info = self.train_step()
                        iter_losses.append(loss_info["total_loss"])
                        all_losses.append(loss_info["total_loss"])

                    mean_loss = np.mean(iter_losses)
                    print(f"  Mean loss: {mean_loss:.4f}")

                    if self.writer is not None:
                        self.writer.add_scalar("loss/total", mean_loss, iteration)
                        self.writer.add_scalar("loss/policy", loss_info["policy_loss"], iteration)
                        self.writer.add_scalar("loss/value", loss_info["value_loss"], iteration)

                # Evaluation
                if iteration % self.config.eval_interval == 0:
                    win_rate_greedy = self.evaluate_vs_greedy(self.config.eval_games)
                    win_rate_random = self.evaluate_vs_random(self.config.eval_games)
                    print(f"  Eval vs Greedy: {win_rate_greedy:.1%}")
                    print(f"  Eval vs Random: {win_rate_random:.1%}")

                    if self.writer is not None:
                        self.writer.add_scalar("eval/win_rate_vs_greedy", win_rate_greedy, iteration)
                        self.writer.add_scalar("eval/win_rate_vs_random", win_rate_random, iteration)

                # Periodic checkpoint saving
                if (
                    self.config.checkpoint_interval > 0
                    and iteration % self.config.checkpoint_interval == 0
                    and self.save_dir is not None
                ):
                    from pathlib import Path
                    checkpoint_path = Path(self.save_dir) / f"checkpoint_iter_{iteration}.pt"
                    self.save(str(checkpoint_path))
                    print(f"  [Checkpoint saved: {checkpoint_path.name}]")

                iter_time = time.time() - iter_start
                print(f"  Iteration time: {iter_time:.1f}s")

        except KeyboardInterrupt:
            print("\n\n[INTERRUPTED] Training stopped by user")
            if self.save_dir is not None:
                from pathlib import Path
                interrupt_checkpoint = Path(self.save_dir) / f"checkpoint_interrupted_iter_{self.iteration}.pt"
                self.save(str(interrupt_checkpoint))
                print(f"[CHECKPOINT SAVED] Progress saved to: {interrupt_checkpoint.name}")
            print(f"Completed {self.iteration}/{num_iterations} iterations")
            # Re-raise to allow outer handler
            raise

        assert self.start_time is not None  # Set at beginning of train()
        total_time = time.time() - self.start_time
        print(f"\nTraining complete in {total_time:.1f}s")

        # Final evaluation
        final_win_rate = self.evaluate_vs_greedy(100)
        print(f"Final win rate vs Greedy: {final_win_rate:.1%}")

        return {
            "iterations": num_iterations,
            "total_games": self.total_games,
            "final_loss": np.mean(all_losses[-100:]) if all_losses else 0,
            "final_win_rate": final_win_rate,
        }

    def _evaluate_vs_opponent(self, opponent: str, num_games: int, seed_offset: int) -> float:
        """Evaluate network against an opponent with random deck matchups.

        Uses direct policy evaluation (no MCTS) for speed.

        Args:
            opponent: Opponent type ("greedy" or "random")
            num_games: Number of evaluation games
            seed_offset: Seed offset for reproducibility

        Returns:
            Win rate (0.0 to 1.0)
        """
        from essence_wars._core import PyGame

        self.network.eval()
        wins = 0
        all_decks = PyGame.list_decks()
        rng = np.random.default_rng(seed_offset)

        for game_idx in range(num_games):
            deck1 = rng.choice(all_decks)
            deck2 = rng.choice(all_decks)
            env = EssenceWarsEnv(deck1=deck1, deck2=deck2, opponent=opponent)
            obs, info = env.reset(seed=game_idx + seed_offset)
            done = False

            while not done:
                if self.obs_normalizer is not None:
                    obs_norm = self.obs_normalizer.normalize(obs).astype(np.float32)
                else:
                    obs_norm = obs

                obs_t = torch.tensor(obs_norm, dtype=torch.float32, device=self.device).unsqueeze(0)
                mask_t = torch.tensor(info["action_mask"] > 0, dtype=torch.bool, device=self.device).unsqueeze(0)

                with torch.no_grad():
                    policy, _ = self.network.evaluate(obs_t, mask_t)

                action = int(policy.argmax(dim=-1).item())

                obs, reward, terminated, truncated, info = env.step(action)
                done = terminated or truncated

                if done and float(reward) > 0:
                    wins += 1

        return wins / num_games

    def evaluate_vs_greedy(self, num_games: int = 100) -> float:
        """Evaluate network against GreedyBot with random decks.

        Args:
            num_games: Number of evaluation games

        Returns:
            Win rate (0.0 to 1.0)
        """
        return self._evaluate_vs_opponent("greedy", num_games, seed_offset=50000)

    def evaluate_vs_random(self, num_games: int = 100) -> float:
        """Evaluate network against RandomBot with random decks."""
        return self._evaluate_vs_opponent("random", num_games, seed_offset=60000)

    def evaluate_with_mcts(self, num_games: int = 50) -> float:
        """Evaluate using full MCTS search (slower but stronger).

        Args:
            num_games: Number of evaluation games

        Returns:
            Win rate (0.0 to 1.0)
        """
        wins = 0

        for game_idx in range(num_games):
            game = PyGame()
            game.reset(seed=game_idx + 70000)

            # Create a greedy opponent by forking
            while not game.is_done():
                if game.current_player() == 0:
                    # Our turn: use MCTS
                    action_probs = self.mcts.search(game, add_noise=False)
                    action = int(action_probs.argmax())
                else:
                    # Opponent's turn: use greedy action
                    action = game.greedy_action()

                game.step(int(action))

            if game.get_reward(0) > 0:
                wins += 1

        return wins / num_games

    def save(self, path: str) -> None:
        """Save model checkpoint."""
        checkpoint = {
            "network_state_dict": self.network.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "iteration": self.iteration,
            "total_games": self.total_games,
            "config": self.config,
        }
        # Save normalizer state if enabled
        if self.obs_normalizer is not None:
            checkpoint["obs_normalizer"] = {
                "mean": self.obs_normalizer.mean,
                "var": self.obs_normalizer.var,
                "count": self.obs_normalizer.count,
            }
        torch.save(checkpoint, path)
        print(f"Saved checkpoint to {path}")

    def load(self, path: str) -> None:
        """Load model checkpoint."""
        checkpoint = torch.load(path, map_location=self.device, weights_only=False)
        self.network.load_state_dict(checkpoint["network_state_dict"])
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        self.iteration = checkpoint.get("iteration", 0)
        self.total_games = checkpoint.get("total_games", 0)
        # Restore normalizer state if present
        if "obs_normalizer" in checkpoint and self.obs_normalizer is not None:
            self.obs_normalizer.mean = checkpoint["obs_normalizer"]["mean"]
            self.obs_normalizer.var = checkpoint["obs_normalizer"]["var"]
            self.obs_normalizer.count = checkpoint["obs_normalizer"]["count"]
        print(f"Loaded checkpoint from {path}")

    def get_network(self) -> AlphaZeroNetwork:
        """Get the trained network."""
        return self.network
