"""
PyTorch neural network for Phutball.

The game state is already stored in canonical orientation by PhutballGame:
the current player's goal is at the bottom row. This file therefore does not
take a player argument or flip boards internally.

Action logits use the same layout as PhutballGame:
  0 .. rows*cols-1 : placements
  rows*cols .. +7  : jump directions
  rows*cols + 8    : end jump sequence
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Tuple

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F

from phutball_game_in_python import BALL, MAN, PhutballState


def encode_state(state: PhutballState, bfs_reachable: np.ndarray | None = None) -> np.ndarray:
    """
    Convert a canonical PhutballState into NN input planes.

    Returns float32 array with shape (5, rows, cols), or (6, rows, cols)
    when bfs_reachable is provided:
      0: men
      1: ball
      2: mid-jump flag, broadcast across the board
      3: BFS reachable squares, if provided
      -2: row coordinate, -1 at opponent goal and +1 at current player's goal
      -1: column coordinate, -1 at left edge and +1 at right edge
    """
    board = state.board
    rows, cols = board.shape
    use_bfs_input = bfs_reachable is not None
    input_channels = 6 if use_bfs_input else 5

    features = np.zeros((input_channels, rows, cols), dtype=np.float32)
    features[0] = board == MAN
    features[1] = board == BALL
    features[2].fill(1.0 if state.is_mid_jump else 0.0)
    coord_start = 4 if use_bfs_input else 3
    if bfs_reachable is not None:
        bfs_array = np.asarray(bfs_reachable, dtype=np.float32)
        if bfs_array.shape != (rows, cols):
            raise ValueError(f"Expected BFS plane shape {(rows, cols)}, got {bfs_array.shape}")
        features[3] = bfs_array

    if rows == 1:
        row_coord = np.zeros((rows,), dtype=np.float32)
    else:
        row_coord = np.linspace(-1.0, 1.0, rows, dtype=np.float32)

    if cols == 1:
        col_coord = np.zeros((cols,), dtype=np.float32)
    else:
        col_coord = np.linspace(-1.0, 1.0, cols, dtype=np.float32)

    features[coord_start] = row_coord[:, None]
    features[coord_start + 1] = col_coord[None, :]
    return features


@dataclass(frozen=True)
class PhutballNetConfig:
    rows: int = 19
    cols: int = 19
    blocks: int = 10
    channels: int = 128
    use_batch_norm: bool = True
    value_head_hidden: int = 256
    jump_head_hidden: int = 128
    gpool_interval: int = 3       # place a GlobalPool residual block every N blocks (0 = disabled)
    gpool_channels: int = 0       # 0 = full-width global pooling, otherwise bottleneck pooled features
    value_conv_channels: int = 32  # channels after the value head 1x1 conv before global pooling
    value_output: str = "tanh"     # "tanh" for [-1, 1] MSE, "logit" for BCE-with-logits
    predict_moves_remaining: bool = False
    predict_next_move: bool = False
    predict_reachable_ball: bool = False
    predict_jump_final_state: bool = False
    use_bfs_input: bool = False
    bfs_node_limit: int = 10_000
    bfs_ordering: str = "novelty"

    @property
    def input_channels(self) -> int:
        return 6 if self.use_bfs_input else 5


class ConvNormRelu(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        *,
        use_batch_norm: bool,
        kernel_size: int = 3,
        padding: int = 1,
    ):
        super().__init__()
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            padding=padding,
            bias=not use_batch_norm,
        )
        self.norm = nn.BatchNorm2d(out_channels) if use_batch_norm else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.relu(self.norm(self.conv(x)))


class InputStem(nn.Module):
    """
    Initial feature mixer.

    Board planes use a normal 3x3 convolution. Coordinate planes use only a
    1x1 convolution, so the net sees each square's own row/column coordinates
    without convolving over neighboring coordinate values.
    """

    def __init__(self, channels: int, *, use_batch_norm: bool, input_channels: int = 5):
        super().__init__()
        if input_channels not in (5, 6):
            raise ValueError("input_channels must be 5 or 6")
        self.input_channels = input_channels
        board_channels = input_channels - 2
        self.board_conv = nn.Conv2d(
            board_channels,
            channels,
            kernel_size=3,
            padding=1,
            bias=not use_batch_norm,
        )
        self.coord_conv = nn.Conv2d(
            2,
            channels,
            kernel_size=1,
            bias=False,
        )
        self.norm = nn.BatchNorm2d(channels) if use_batch_norm else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        board_planes = x[:, : self.input_channels - 2]
        coord_planes = x[:, self.input_channels - 2 :]
        x = self.board_conv(board_planes) + self.coord_conv(coord_planes)
        return F.relu(self.norm(x))


class ResidualBlock(nn.Module):
    def __init__(self, channels: int, *, use_batch_norm: bool):
        super().__init__()
        self.conv1 = nn.Conv2d(
            channels,
            channels,
            kernel_size=3,
            padding=1,
            bias=not use_batch_norm,
        )
        self.norm1 = nn.BatchNorm2d(channels) if use_batch_norm else nn.Identity()
        self.conv2 = nn.Conv2d(
            channels,
            channels,
            kernel_size=3,
            padding=1,
            bias=not use_batch_norm,
        )
        self.norm2 = nn.BatchNorm2d(channels) if use_batch_norm else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = F.relu(self.norm1(self.conv1(x)))
        x = self.norm2(self.conv2(x))
        return F.relu(x + residual)


class GlobalPoolBias(nn.Module):
    """
    Pools a spatial feature map globally and projects the result to per-channel
    bias scalars via a linear layer.

    If gpool_channels is positive and smaller than in_channels, the spatial map
    is first bottlenecked with a 1x1 convolution before mean/max pooling. This
    keeps the global context path closer to KataGo-style ratios for wide trunks.

    The caller adds the returned (batch, out_channels) vector to its spatial
    feature map after broadcasting over the spatial dimensions.
    """

    def __init__(self, in_channels: int, out_channels: int, *, gpool_channels: int = 0):
        super().__init__()
        if gpool_channels < 0:
            raise ValueError("gpool_channels must be non-negative")
        self.gpool_channels = gpool_channels
        if 0 < gpool_channels < in_channels:
            self.reduce = nn.Conv2d(in_channels, gpool_channels, kernel_size=1)
            pooled_channels = gpool_channels
        else:
            self.reduce = None
            pooled_channels = in_channels
        self.fc = nn.Linear(2 * pooled_channels, out_channels)

    def forward(self, g: torch.Tensor) -> torch.Tensor:
        if self.reduce is not None:
            g = F.relu(self.reduce(g))
        mean = g.mean(dim=(2, 3))
        max_ = g.amax(dim=(2, 3))
        return self.fc(torch.cat([mean, max_], dim=1))


class ResidualBlockWithGPool(nn.Module):
    """
    Residual block with a global pooling bias injection after the second conv.

    After the second conv+norm, the activated feature map is pooled globally and
    projected to per-channel biases, which are added back before the skip
    connection. This gives every spatial cell access to board-wide context even
    though the convolutions themselves are local.
    """

    def __init__(self, channels: int, *, use_batch_norm: bool, gpool_channels: int = 0):
        super().__init__()
        self.conv1 = nn.Conv2d(
            channels, channels, kernel_size=3, padding=1, bias=not use_batch_norm
        )
        self.norm1 = nn.BatchNorm2d(channels) if use_batch_norm else nn.Identity()
        self.conv2 = nn.Conv2d(
            channels, channels, kernel_size=3, padding=1, bias=not use_batch_norm
        )
        self.norm2 = nn.BatchNorm2d(channels) if use_batch_norm else nn.Identity()
        self.gpool = GlobalPoolBias(channels, channels, gpool_channels=gpool_channels)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = F.relu(self.norm1(self.conv1(x)))
        x = self.norm2(self.conv2(x))
        bias = self.gpool(F.relu(x))
        x = x + bias[:, :, None, None]
        return F.relu(x + residual)


class PhutballNet(nn.Module):
    """
    AlphaZero-style residual network for Phutball with global pooling.

    forward(x) expects x with shape (batch, input_channels, rows, cols) and returns:
      policy_logits: (batch, rows*cols + 9)
      value:         (batch,)

    The network intentionally returns raw logits. Legal-action masking belongs
    in MCTS and training loss code, not inside the model.
    """

    def __init__(self, config: PhutballNetConfig | None = None):
        super().__init__()
        self.config = config or PhutballNetConfig()
        cfg = self.config

        self.input = InputStem(
            cfg.channels,
            use_batch_norm=cfg.use_batch_norm,
            input_channels=cfg.input_channels,
        )

        blocks = []
        for i in range(cfg.blocks):
            use_gpool = cfg.gpool_interval > 0 and (i + 1) % cfg.gpool_interval == 0
            if use_gpool:
                blocks.append(
                    ResidualBlockWithGPool(
                        cfg.channels,
                        use_batch_norm=cfg.use_batch_norm,
                        gpool_channels=cfg.gpool_channels,
                    )
                )
            else:
                blocks.append(
                    ResidualBlock(cfg.channels, use_batch_norm=cfg.use_batch_norm)
                )
        self.trunk = nn.Sequential(*blocks)

        # Placement head: inject global context then 1x1 conv
        self.placement_gpool = GlobalPoolBias(
            cfg.channels,
            cfg.channels,
            gpool_channels=cfg.gpool_channels,
        )
        self.placement_head = nn.Conv2d(cfg.channels, 1, kernel_size=1)

        # Jump head: already receives global context via _jump_features
        jump_in = cfg.channels * 2
        self.jump_head = nn.Sequential(
            nn.Linear(jump_in, cfg.jump_head_hidden),
            nn.ReLU(),
            nn.Linear(cfg.jump_head_hidden, 9),
        )

        # Value head: 1x1 conv → global pool → MLP
        self.value_conv = nn.Conv2d(cfg.channels, cfg.value_conv_channels, kernel_size=1)
        self.value_fc1 = nn.Linear(2 * cfg.value_conv_channels, cfg.value_head_hidden)
        self.value_fc2 = nn.Linear(cfg.value_head_hidden, 1)
        if cfg.predict_moves_remaining:
            self.moves_remaining_fc = nn.Linear(2 * cfg.value_conv_channels, 1)
            nn.init.constant_(self.moves_remaining_fc.bias, 7.0)

        if cfg.predict_next_move:
            self.next_move_placement_gpool = GlobalPoolBias(
                cfg.channels,
                cfg.channels,
                gpool_channels=cfg.gpool_channels,
            )
            self.next_move_placement_head = nn.Conv2d(cfg.channels, 1, kernel_size=1)
            self.next_move_jump_head = nn.Sequential(
                nn.Linear(jump_in, cfg.jump_head_hidden),
                nn.ReLU(),
                nn.Linear(cfg.jump_head_hidden, 9),
            )

        if cfg.predict_reachable_ball:
            self.reachable_ball_gpool = GlobalPoolBias(
                cfg.channels,
                cfg.channels,
                gpool_channels=cfg.gpool_channels,
            )
            self.reachable_ball_head = nn.Conv2d(cfg.channels, 1, kernel_size=1)

        if cfg.predict_jump_final_state:
            self.jump_final_gpool = GlobalPoolBias(
                cfg.channels,
                cfg.channels,
                gpool_channels=cfg.gpool_channels,
            )
            self.jump_final_ball_head = nn.Conv2d(cfg.channels, 1, kernel_size=1)
            self.jump_final_men_head = nn.Conv2d(cfg.channels, 1, kernel_size=1)

    @property
    def action_size(self) -> int:
        return self.config.rows * self.config.cols + 9

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
        cfg = self.config
        if x.ndim != 4:
            raise ValueError(
                f"Expected input shape (batch, {cfg.input_channels}, rows, cols), got {tuple(x.shape)}"
            )
        if x.shape[1:] != (cfg.input_channels, cfg.rows, cfg.cols):
            raise ValueError(
                f"Expected input shape (batch, {cfg.input_channels}, {cfg.rows}, {cfg.cols}), "
                f"got {tuple(x.shape)}"
            )
        return self._forward_impl(x)

    def _forward_impl(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
        cfg = self.config
        features = self.trunk(self.input(x))

        gpool_bias = self.placement_gpool(features)
        placement_logits = self.placement_head(
            features + gpool_bias[:, :, None, None]
        ).flatten(start_dim=1)

        jump_features = self._jump_features(features, x[:, 1])
        jump_logits = self.jump_head(jump_features)

        policy_logits = torch.cat([placement_logits, jump_logits], dim=1)

        v = F.relu(self.value_conv(features))
        value_features = torch.cat([v.mean(dim=(2, 3)), v.amax(dim=(2, 3))], dim=1)
        value_hidden = F.relu(self.value_fc1(value_features))
        value_raw = self.value_fc2(value_hidden)
        if cfg.value_output == "tanh":
            value = torch.tanh(value_raw).squeeze(-1)
        elif cfg.value_output == "logit":
            value = value_raw.squeeze(-1)
        else:
            raise ValueError(f"Unknown value_output: {cfg.value_output!r}")

        outputs: list[torch.Tensor] = [policy_logits, value]
        if cfg.predict_moves_remaining:
            sqrt_moves_remaining = self.moves_remaining_fc(value_features).squeeze(-1)
            outputs.append(sqrt_moves_remaining)

        if cfg.predict_next_move:
            next_gpool_bias = self.next_move_placement_gpool(features)
            next_placement_logits = self.next_move_placement_head(
                features + next_gpool_bias[:, :, None, None]
            ).flatten(start_dim=1)
            next_jump_logits = self.next_move_jump_head(jump_features)
            outputs.append(torch.cat([next_placement_logits, next_jump_logits], dim=1))

        if cfg.predict_reachable_ball:
            reachable_gpool_bias = self.reachable_ball_gpool(features)
            reachable_logits = self.reachable_ball_head(
                features + reachable_gpool_bias[:, :, None, None]
            ).flatten(start_dim=1)
            outputs.append(reachable_logits)

        if cfg.predict_jump_final_state:
            jump_final_gpool_bias = self.jump_final_gpool(features)
            jump_final_features = features + jump_final_gpool_bias[:, :, None, None]
            jump_final_ball_logits = self.jump_final_ball_head(jump_final_features).flatten(start_dim=1)
            jump_final_men_logits = self.jump_final_men_head(jump_final_features).flatten(start_dim=1)
            outputs.extend([jump_final_ball_logits, jump_final_men_logits])

        return tuple(outputs)

    def _jump_features(self, features: torch.Tensor, ball_plane: torch.Tensor) -> torch.Tensor:
        batch, channels, rows, cols = features.shape
        flat_features = features.flatten(start_dim=2)

        global_features = flat_features.mean(dim=2)

        ball_flat = ball_plane.reshape(batch, rows * cols)
        ball_index = ball_flat.argmax(dim=1)
        gather_index = ball_index[:, None, None].expand(batch, channels, 1)
        ball_features = flat_features.gather(dim=2, index=gather_index).squeeze(2)

        return torch.cat([global_features, ball_features], dim=1)


if __name__ == "__main__":
    from phutball_game_in_python import PhutballGame

    game = PhutballGame(rows=19, cols=19, max_rows=19, max_cols=19)
    state = game.initial_state()
    x = torch.from_numpy(encode_state(state)).unsqueeze(0)
    net = PhutballNet()
    policy, value = net(x)[:2]
    print("policy", tuple(policy.shape))
    print("value", tuple(value.shape))
