"""MCTS implementation for AlphaZero.

Provides neural network guided Monte Carlo Tree Search with:
- PUCT selection formula for exploration-exploitation balance
- Dirichlet noise at root for exploration
- Batched GPU inference for efficiency
"""

from __future__ import annotations

import math
from typing import TYPE_CHECKING, Any

import numpy as np
import torch

if TYPE_CHECKING:
    from numpy.typing import NDArray

    from essence_wars._core import PyGame
    from essence_wars.agents.networks import AlphaZeroNetwork

    from ..utils import RunningMeanStd
    from .config import AlphaZeroConfig


class MCTSNode:
    """Node in the MCTS tree.

    Each node stores:
    - Visit count N(s,a) for each action
    - Total value W(s,a) for each action
    - Prior probability P(s,a) from neural network
    - Children nodes

    Args:
        prior: Prior probability from policy network
        parent: Parent node in the tree
        action: Action that led to this node
    """

    def __init__(
        self,
        prior: float = 0.0,
        parent: MCTSNode | None = None,
        action: int | None = None,
    ) -> None:
        self.prior = prior
        self.parent = parent
        self.action = action

        self.visit_count = 0
        self.value_sum = 0.0

        # Children: action -> MCTSNode
        self.children: dict[int, MCTSNode] = {}

        # Whether this node has been expanded
        self.is_expanded = False

    @property
    def value(self) -> float:
        """Mean value Q(s,a) = W(s,a) / N(s,a)."""
        if self.visit_count == 0:
            return 0.0
        return self.value_sum / self.visit_count

    def expand(
        self,
        legal_actions: list[int],
        priors: NDArray[np.float32],
    ) -> None:
        """Expand this node with children for each legal action.

        Args:
            legal_actions: List of legal action indices
            priors: Prior probabilities for all actions (422,)
        """
        self.is_expanded = True

        # Normalize priors over legal actions only
        legal_priors = priors[legal_actions]
        legal_priors = legal_priors / (legal_priors.sum() + 1e-8)

        for i, action in enumerate(legal_actions):
            self.children[action] = MCTSNode(
                prior=legal_priors[i],
                parent=self,
                action=action,
            )

    def select_child(self, c_puct: float) -> tuple[int, MCTSNode]:
        """Select child with highest UCB score.

        UCB(s,a) = Q(s,a) + c_puct * P(s,a) * sqrt(N(s)) / (1 + N(s,a))

        Args:
            c_puct: Exploration constant

        Returns:
            (action, child_node) tuple
        """
        best_score = -float("inf")
        best_action = -1
        best_child = None

        sqrt_total = math.sqrt(self.visit_count + 1)

        for action, child in self.children.items():
            # UCB score
            q_value = child.value
            u_value = c_puct * child.prior * sqrt_total / (1 + child.visit_count)
            score = q_value + u_value

            if score > best_score:
                best_score = score
                best_action = action
                best_child = child

        assert best_child is not None
        return best_action, best_child

    def add_exploration_noise(
        self,
        dirichlet_alpha: float,
        epsilon: float,
    ) -> None:
        """Add Dirichlet noise to priors at root for exploration.

        Args:
            dirichlet_alpha: Dirichlet distribution parameter
            epsilon: Weight of noise (1-epsilon is weight of original prior)
        """
        actions = list(self.children.keys())
        noise = np.random.dirichlet([dirichlet_alpha] * len(actions))

        for i, action in enumerate(actions):
            child = self.children[action]
            child.prior = (1 - epsilon) * child.prior + epsilon * noise[i]

    def apply_virtual_loss(self, virtual_loss: float = 3.0) -> None:
        """Apply virtual loss to discourage other threads from selecting same path."""
        self.visit_count += 1
        self.value_sum -= virtual_loss

    def remove_virtual_loss(self, virtual_loss: float = 3.0) -> None:
        """Remove virtual loss after evaluation is complete."""
        self.visit_count -= 1
        self.value_sum += virtual_loss


class NeuralMCTS:
    """MCTS guided by neural network policy and value.

    Key features:
    - Uses neural network for prior probabilities and leaf evaluation
    - PUCT formula for exploration-exploitation balance
    - Dirichlet noise at root for exploration

    Args:
        network: AlphaZeroNetwork for policy and value
        config: AlphaZeroConfig with MCTS parameters
        obs_normalizer: Optional observation normalizer

    Example:
        network = AlphaZeroNetwork()
        mcts = NeuralMCTS(network, config)

        game = PyGame()
        game.reset(seed=42)

        action_probs = mcts.search(game)
        action = np.random.choice(422, p=action_probs)
    """

    def __init__(
        self,
        network: AlphaZeroNetwork,
        config: AlphaZeroConfig,
        obs_normalizer: RunningMeanStd | None = None,
    ) -> None:
        self.config = config
        self.device = torch.device(config.device)
        self.network = network.to(self.device)
        self.obs_normalizer = obs_normalizer

    def search(
        self,
        game: PyGame,
        add_noise: bool = True,
    ) -> NDArray[np.float32]:
        """Run MCTS search from current game state.

        Args:
            game: Current game state (will be forked, not modified)
            add_noise: Whether to add Dirichlet noise at root

        Returns:
            Action probabilities based on visit counts (422,)
        """
        # Create root node
        root = MCTSNode()

        # Get initial evaluation
        obs = game.observe()
        mask = game.action_mask()
        legal_actions = np.where(mask > 0)[0].tolist()

        # Evaluate root position
        policy, _ = self._evaluate(obs, mask)
        root.expand(legal_actions, policy)

        # Add exploration noise at root
        if add_noise and len(legal_actions) > 1:
            root.add_exploration_noise(
                self.config.dirichlet_alpha,
                self.config.dirichlet_epsilon,
            )

        # Run simulations
        for _ in range(self.config.num_simulations):
            # Fork game for simulation
            sim_game = game.fork()
            node = root
            search_path = [node]

            # Selection: traverse tree using UCB
            while node.is_expanded and not sim_game.is_done():
                action, node = node.select_child(self.config.c_puct)
                sim_game.step(action)
                search_path.append(node)

            # Get value
            if sim_game.is_done():
                # Terminal node: use actual game result
                # Reward is from player 0's perspective
                value = sim_game.get_reward(0)
            else:
                # Non-terminal: expand and evaluate with network
                obs = sim_game.observe()
                mask = sim_game.action_mask()
                legal_actions = np.where(mask > 0)[0].tolist()

                if legal_actions:
                    policy, value = self._evaluate(obs, mask)
                    node.expand(legal_actions, policy)
                else:
                    # No legal actions (shouldn't happen in normal game)
                    value = 0.0

            # Backup: propagate value up the tree
            # Note: value alternates sign as we go up (opponent's perspective)
            sim_game.current_player()
            for node in reversed(search_path):
                # Flip value for opponent's nodes
                node.visit_count += 1
                node.value_sum += value
                value = -value  # Flip for parent (opponent's perspective)

        # Return visit count distribution as action probabilities
        return self._get_action_probs(root)

    def _evaluate(
        self,
        obs: NDArray[np.float32],
        mask: NDArray[np.float32],
    ) -> tuple[NDArray[np.float32], float]:
        """Evaluate position with neural network.

        Args:
            obs: Observation array (STATE_TENSOR_SIZE,)
            mask: Action mask (422,)

        Returns:
            (policy, value) tuple
        """
        # Normalize observation if enabled
        if self.obs_normalizer is not None:
            obs = self.obs_normalizer.normalize(obs).astype(np.float32)

        obs_t = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0)
        mask_t = torch.tensor(mask > 0, dtype=torch.bool, device=self.device).unsqueeze(0)

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

        return policy.squeeze(0).cpu().numpy(), value.item()

    def _get_action_probs(self, root: MCTSNode) -> NDArray[np.float32]:
        """Get action probabilities from root visit counts.

        Args:
            root: Root node of MCTS tree

        Returns:
            Action probabilities (422,)
        """
        probs = np.zeros(422, dtype=np.float32)
        total_visits = sum(child.visit_count for child in root.children.values())

        if total_visits > 0:
            for action, child in root.children.items():
                probs[action] = child.visit_count / total_visits

        return probs

    def _evaluate_batch(
        self,
        obs_batch: NDArray[np.float32],
        mask_batch: NDArray[np.float32],
    ) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
        """Evaluate a batch of positions with neural network.

        Args:
            obs_batch: Observations array (batch_size, STATE_TENSOR_SIZE)
            mask_batch: Action masks (batch_size, 422)

        Returns:
            (policies, values) tuple - policies (batch_size, 422), values (batch_size,)
        """
        # Normalize observations if enabled
        if self.obs_normalizer is not None:
            obs_batch = np.stack([
                self.obs_normalizer.normalize(obs).astype(np.float32)
                for obs in obs_batch
            ])

        obs_t = torch.tensor(obs_batch, dtype=torch.float32, device=self.device)
        mask_t = torch.tensor(mask_batch > 0, dtype=torch.bool, device=self.device)

        self.network.eval()
        with torch.no_grad():
            policies, values = self.network.evaluate(obs_t, mask_t)

        return policies.cpu().numpy(), values.cpu().numpy().flatten()

    def search_batched(
        self,
        game: PyGame,
        batch_size: int = 32,
        add_noise: bool = True,
        virtual_loss: float = 3.0,
    ) -> NDArray[np.float32]:
        """Run batched MCTS search for GPU efficiency.

        Collects multiple leaf nodes before doing a single batched GPU inference,
        using virtual loss to encourage exploration of different paths.

        Args:
            game: Current game state (will be forked, not modified)
            batch_size: Number of leaves to collect before GPU inference
            add_noise: Whether to add Dirichlet noise at root
            virtual_loss: Virtual loss to apply during selection

        Returns:
            Action probabilities based on visit counts (422,)
        """
        # Create root node
        root = MCTSNode()

        # Get initial evaluation
        obs = game.observe()
        mask = game.action_mask()
        legal_actions = np.where(mask > 0)[0].tolist()

        # Evaluate root position
        policy, _ = self._evaluate(obs, mask)
        root.expand(legal_actions, policy)

        # Add exploration noise at root
        if add_noise and len(legal_actions) > 1:
            root.add_exploration_noise(
                self.config.dirichlet_alpha,
                self.config.dirichlet_epsilon,
            )

        # Run simulations in batches
        remaining_sims = self.config.num_simulations
        while remaining_sims > 0:
            # Collect leaves for this batch
            current_batch = min(batch_size, remaining_sims)
            leaves_to_evaluate: list[dict[str, Any]] = []

            for _ in range(current_batch):
                # Fork game for simulation
                sim_game = game.fork()
                node = root
                search_path = [node]

                # Selection with virtual loss
                while node.is_expanded and not sim_game.is_done():
                    action, node = node.select_child(self.config.c_puct)
                    node.apply_virtual_loss(virtual_loss)
                    sim_game.step(action)
                    search_path.append(node)

                # Check if terminal
                if sim_game.is_done():
                    # Terminal: use actual game result, remove virtual loss and backup
                    value = sim_game.get_reward(0)
                    for path_node in reversed(search_path[1:]):
                        path_node.remove_virtual_loss(virtual_loss)
                        path_node.visit_count += 1
                        path_node.value_sum += value
                        value = -value
                    root.visit_count += 1
                else:
                    # Non-terminal: collect for batch evaluation
                    leaf_obs = sim_game.observe()
                    leaf_mask = sim_game.action_mask()
                    leaf_legal = np.where(leaf_mask > 0)[0].tolist()

                    if leaf_legal:
                        leaves_to_evaluate.append({
                            "obs": leaf_obs,
                            "mask": leaf_mask,
                            "legal_actions": leaf_legal,
                            "node": node,
                            "search_path": search_path,
                        })
                    else:
                        # No legal actions: backup with 0 value
                        value = 0.0
                        for path_node in reversed(search_path[1:]):
                            path_node.remove_virtual_loss(virtual_loss)
                            path_node.visit_count += 1
                            path_node.value_sum += value
                            value = -value
                        root.visit_count += 1

            # Batch evaluate all collected leaves
            if leaves_to_evaluate:
                obs_batch = np.stack([leaf["obs"] for leaf in leaves_to_evaluate])
                mask_batch = np.stack([leaf["mask"] for leaf in leaves_to_evaluate])

                policies, values = self._evaluate_batch(obs_batch, mask_batch)

                # Expand and backup each leaf
                for i, leaf in enumerate(leaves_to_evaluate):
                    node = leaf["node"]
                    search_path = leaf["search_path"]

                    # Expand leaf
                    node.expand(leaf["legal_actions"], policies[i])

                    # Remove virtual loss and backup
                    value = values[i]
                    for path_node in reversed(search_path[1:]):
                        path_node.remove_virtual_loss(virtual_loss)
                        path_node.visit_count += 1
                        path_node.value_sum += value
                        value = -value
                    root.visit_count += 1

            remaining_sims -= current_batch

        return self._get_action_probs(root)
