"""Replay buffers for AlphaZero training.

Provides replay buffers for storing and sampling training data:
- ReplayBuffer: Standard single-buffer for self-play training
- DualReplayBuffer: Two-buffer system for BC warm-start with fixed ratio
"""

from __future__ import annotations

import random
from collections import deque
from dataclasses import dataclass, field
from typing import Any

import numpy as np
from numpy.typing import NDArray

# Type alias for replay buffer samples
ReplaySample = tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.float32], float]


@dataclass
class ReplayBuffer:
    """Replay buffer for AlphaZero training.

    Stores (observation, action_mask, policy_target, value_target) tuples
    from self-play games.

    Args:
        capacity: Maximum number of samples to store

    Example:
        buffer = ReplayBuffer(capacity=100_000)
        buffer.add_game(observations, masks, policies, outcome)
        obs, masks, policies, values = buffer.sample(batch_size=256)
    """

    capacity: int = 100_000
    buffer: deque[ReplaySample] = field(default_factory=deque)

    def __post_init__(self) -> None:
        self.buffer = deque(maxlen=self.capacity)

    def add(
        self,
        obs: NDArray[np.float32],
        mask: NDArray[np.float32],
        policy: NDArray[np.float32],
        value: float,
    ) -> None:
        """Add a training sample to the buffer."""
        self.buffer.append((obs.copy(), mask.copy(), policy.copy(), value))

    def add_game(
        self,
        observations: list[NDArray[np.float32]],
        masks: list[NDArray[np.float32]],
        policies: list[NDArray[np.float32]],
        outcome: float,
    ) -> None:
        """Add all positions from a game to the buffer.

        Args:
            observations: List of observations from game
            masks: List of action masks from game
            policies: List of MCTS policy targets from game
            outcome: Game outcome from player 0's perspective (+1/-1)
        """
        # Assign values: each position gets the outcome from its player's perspective
        # Positions alternate between players, so values alternate
        for i, (obs, mask, policy) in enumerate(zip(observations, masks, policies, strict=False)):
            # Even indices are player 0, odd are player 1
            value = outcome if i % 2 == 0 else -outcome
            self.add(obs, mask, policy, value)

    def sample(
        self, batch_size: int
    ) -> tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.float32], NDArray[np.float32]]:
        """Sample a random batch from the buffer.

        Returns:
            (observations, masks, policies, values) tensors
        """
        indices = np.random.choice(len(self.buffer), size=batch_size, replace=False)
        batch = [self.buffer[i] for i in indices]

        obs = np.stack([b[0] for b in batch])
        masks = np.stack([b[1] for b in batch])
        policies = np.stack([b[2] for b in batch])
        values = np.array([b[3] for b in batch], dtype=np.float32)

        return obs, masks, policies, values

    def __len__(self) -> int:
        return len(self.buffer)


@dataclass
class DualReplayBuffer:
    """Dual replay buffer for BC warm-start training.

    Maintains separate BC and self-play buffers with fixed sampling ratio.
    This prevents BC ratio from dropping as self-play data accumulates,
    solving the "delayed catastrophic forgetting" problem.

    The key insight is that mixing BC and self-play data in a single buffer
    causes the BC ratio to drop over time, eventually leading to collapse.
    By keeping them separate and sampling with a fixed ratio, we maintain
    stable training indefinitely.

    Args:
        bc_ratio: Ratio of BC samples in training batches (0.7 = 70% BC)
        bc_capacity: Maximum BC samples to store
        selfplay_capacity: Maximum self-play samples to store

    Example:
        buffer = DualReplayBuffer(bc_ratio=0.7)

        # Load BC data during initialization
        for sample in bc_dataset:
            buffer.add_bc(sample.obs, sample.mask, sample.policy, sample.value)

        # Add self-play during training
        buffer.add_selfplay_game(observations, masks, policies, outcome)

        # Sample with guaranteed 70% BC ratio
        batch = buffer.sample(256)
    """

    bc_ratio: float = 0.7
    bc_capacity: int = 100_000
    selfplay_capacity: int = 100_000

    def __post_init__(self) -> None:
        self.bc_buffer: deque[ReplaySample] = deque(maxlen=self.bc_capacity)
        self.selfplay_buffer: deque[ReplaySample] = deque(maxlen=self.selfplay_capacity)

    def add_bc(
        self,
        obs: NDArray[np.float32],
        mask: NDArray[np.float32],
        policy: NDArray[np.float32],
        value: float,
    ) -> None:
        """Add sample to BC buffer (typically during initialization)."""
        self.bc_buffer.append((obs.copy(), mask.copy(), policy.copy(), value))

    def add_selfplay(
        self,
        obs: NDArray[np.float32],
        mask: NDArray[np.float32],
        policy: NDArray[np.float32],
        value: float,
    ) -> None:
        """Add sample to self-play buffer (during training)."""
        self.selfplay_buffer.append((obs.copy(), mask.copy(), policy.copy(), value))

    def add_selfplay_game(
        self,
        observations: list[NDArray[np.float32]],
        masks: list[NDArray[np.float32]],
        policies: list[NDArray[np.float32]],
        outcome: float,
    ) -> None:
        """Add all positions from a self-play game to the buffer.

        Args:
            observations: List of observations from game
            masks: List of action masks from game
            policies: List of MCTS policy targets from game
            outcome: Game outcome from player 0's perspective (+1/-1)
        """
        for i, (obs, mask, policy) in enumerate(zip(observations, masks, policies, strict=False)):
            value = outcome if i % 2 == 0 else -outcome
            self.add_selfplay(obs, mask, policy, value)

    def sample(
        self, batch_size: int
    ) -> tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.float32], NDArray[np.float32]]:
        """Sample batch with fixed BC ratio.

        If self-play buffer doesn't have enough samples, fills remainder
        from BC buffer to ensure full batch size.

        Returns:
            (observations, masks, policies, values) numpy arrays
        """
        bc_size = int(batch_size * self.bc_ratio)
        sp_size = batch_size - bc_size

        # Sample from BC buffer
        bc_indices = np.random.choice(len(self.bc_buffer), size=bc_size, replace=True)
        bc_samples = [self.bc_buffer[i] for i in bc_indices]

        # Sample from self-play buffer (or fallback to BC if not enough)
        if len(self.selfplay_buffer) >= sp_size:
            sp_indices = np.random.choice(
                len(self.selfplay_buffer), size=sp_size, replace=True
            )
            sp_samples = [self.selfplay_buffer[i] for i in sp_indices]
        elif len(self.selfplay_buffer) > 0:
            # Use all self-play samples + fill remainder from BC
            sp_samples = list(self.selfplay_buffer)
            extra_needed = sp_size - len(sp_samples)
            extra_bc = np.random.choice(
                len(self.bc_buffer), size=extra_needed, replace=True
            )
            sp_samples.extend([self.bc_buffer[i] for i in extra_bc])
        else:
            # No self-play yet, use all BC
            extra_bc = np.random.choice(len(self.bc_buffer), size=sp_size, replace=True)
            sp_samples = [self.bc_buffer[i] for i in extra_bc]

        # Combine and shuffle
        all_samples = bc_samples + sp_samples
        random.shuffle(all_samples)

        # Convert to arrays
        obs = np.stack([s[0] for s in all_samples])
        masks = np.stack([s[1] for s in all_samples])
        policies = np.stack([s[2] for s in all_samples])
        values = np.array([s[3] for s in all_samples], dtype=np.float32)

        return obs, masks, policies, values

    def __len__(self) -> int:
        """Total samples across both buffers."""
        return len(self.bc_buffer) + len(self.selfplay_buffer)

    @property
    def bc_count(self) -> int:
        """Number of BC samples."""
        return len(self.bc_buffer)

    @property
    def selfplay_count(self) -> int:
        """Number of self-play samples."""
        return len(self.selfplay_buffer)

    def stats(self) -> dict[str, Any]:
        """Get buffer statistics."""
        total = len(self)
        return {
            "bc_samples": self.bc_count,
            "selfplay_samples": self.selfplay_count,
            "total_samples": total,
            "effective_bc_ratio": self.bc_ratio,  # Always the configured ratio
            "actual_bc_ratio": self.bc_count / total if total > 0 else 0,
        }
