"""Competition stdio agent: direct policy sampling over a trained equinox checkpoint.

Self-contained competition bot. It speaks the wire protocol in
`generals-bots/competition/protocol.py` (handshake, then one observation frame
per turn, one action line per turn, EOF on stdin = game over) and, on each
turn, samples an action from the policy network exactly as the model is used
during training/evaluation — no tree search.

Unlike `scripts/mcts_agent.py`, this module does **not** import the
`general_bots_training` training package. It inlines the policy-value network
and the observation encoding so the script can be dropped into the competition
sandbox with just the engine (`generals`), `equinox`, `jax`, and `numpy`
available.

This is the search-free counterpart to `scripts/competition_puct.py`: it shares
the same network and observation encoding, but selects an action by sampling
from the masked policy (as in `PolicyValueNetwork.__call__`) instead of running
particle PUCT.

Run it directly through the bundled matchup driver, e.g.:

    PYTHONPATH=src:generals-bots .venv/bin/python \
      generals-bots/competition/matchup.py \
      scripts/competition.py --checkpoint ppo_model.eqx \
      generals-bots/competition/agents/expander_python/run.sh \
      --mode competition

or wrap it in a `run.sh` that passes the checkpoint path.
"""

import argparse
import os
import sys
import time

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jrandom
import numpy as np
from generals.core import game
from generals.core.action import compute_valid_move_mask
from generals.core.game import GameState
from generals.core.observation import Observation
from generals.modifiers import build_castles, deathtouch

os.environ.setdefault("JAX_PLATFORMS", "cpu")

PASS_ACTION = np.array([1, 0, 0, 0, 0], dtype=np.int32)


# ---------------------------------------------------------------------------
# Policy-value network (inlined from general_bots_training.network)
# ---------------------------------------------------------------------------


def obs_to_tensor(obs: Observation) -> jnp.ndarray:
    """Encode an Observation into a (C, H, W) float32 tensor for the network.

    Armies and army-counts are log-normalized; land counts are divided by the
    number of cells; scalar values are broadcast to spatial planes so a plain
    conv stack can consume them. Grid-size agnostic.
    """
    H, W = obs.armies.shape
    armies = jnp.log1p(obs.armies.astype(jnp.float32)) / jnp.log(50.0)

    def bcast(scalar):
        return jnp.broadcast_to(scalar.astype(jnp.float32), (H, W))

    own_land = bcast(obs.owned_land_count / (H * W))
    own_army = bcast(jnp.log1p(obs.owned_army_count.astype(jnp.float32)) / jnp.log(50.0))
    opp_land = bcast(obs.opponent_land_count / (H * W))
    opp_army = bcast(jnp.log1p(obs.opponent_army_count.astype(jnp.float32)) / jnp.log(50.0))
    timestep = bcast(jnp.log1p(obs.timestep.astype(jnp.float32)) / jnp.log(1201.0))

    return jnp.stack(
        [
            armies,
            obs.generals.astype(jnp.float32),
            obs.castles.astype(jnp.float32),
            obs.mountains.astype(jnp.float32),
            obs.neutral_cells.astype(jnp.float32),
            obs.owned_cells.astype(jnp.float32),
            obs.opponent_cells.astype(jnp.float32),
            obs.fog_cells.astype(jnp.float32),
            obs.structures_in_fog.astype(jnp.float32),
            own_land,
            own_army,
            opp_land,
            opp_army,
            timestep,
        ],
        axis=0,
    )


class PolicyValueNetwork(eqx.Module):
    """Conv policy-value network.

    Action layout: 4 full-move directions and 4 half-move (split) directions
    per source cell, plus one global pass action. Invalid moves are masked to
    -1e9; pass is always available.

    The value head uses global average pooling over the spatial grid, so the
    network works for any board size without reshaping linear layers.
    """

    conv1: eqx.nn.Conv2d
    conv2: eqx.nn.Conv2d
    conv3: eqx.nn.Conv2d
    conv4: eqx.nn.Conv2d
    policy_conv: eqx.nn.Conv2d
    value_conv: eqx.nn.Conv2d
    value_linear1: eqx.nn.Linear
    value_linear2: eqx.nn.Linear

    def __init__(self, key, in_channels: int = 14, channels=(32, 32, 32, 16)):
        keys = jrandom.split(key, 8)
        self.conv1 = eqx.nn.Conv2d(in_channels, channels[0], kernel_size=3, padding=1, key=keys[0])
        self.conv2 = eqx.nn.Conv2d(channels[0], channels[1], kernel_size=3, padding=1, key=keys[1])
        self.conv3 = eqx.nn.Conv2d(channels[1], channels[2], kernel_size=3, padding=1, key=keys[2])
        self.conv4 = eqx.nn.Conv2d(channels[2], channels[3], kernel_size=3, padding=1, key=keys[3])
        # 9 = 4 dirs (full) + 4 dirs (half) + 1 pass
        self.policy_conv = eqx.nn.Conv2d(channels[3], 9, kernel_size=1, key=keys[4])
        self.value_conv = eqx.nn.Conv2d(channels[3], 4, kernel_size=1, key=keys[5])
        self.value_linear1 = eqx.nn.Linear(4, 64, key=keys[6])
        self.value_linear2 = eqx.nn.Linear(64, 1, key=keys[7])

    def _features(self, obs):
        x = jax.nn.relu(self.conv1(obs))
        x = jax.nn.relu(self.conv2(x))
        x = jax.nn.relu(self.conv3(x))
        x = jax.nn.relu(self.conv4(x))
        return x

    def _value_from_features(self, feat):
        v = jax.nn.relu(self.value_conv(feat))  # (4, H, W)
        v = v.mean(axis=(1, 2))  # (4,) global average pool
        v = jax.nn.relu(self.value_linear1(v))
        return self.value_linear2(v)[0]

    def policy_value(self, obs, mask):
        """Return deterministic masked policy logits and the critic value."""
        features = self._features(obs)
        return self._policy_logits(features, mask), self._value_from_features(features)

    def _policy_logits(self, feat, mask):
        """Return 8*H*W move logits plus one global pass logit."""
        logits = self.policy_conv(feat)  # (9, H, W)
        mask_t = jnp.transpose(mask, (2, 0, 1))  # (4, H, W)
        penalty = (1.0 - mask_t) * -1e9
        move_penalty = jnp.concatenate([penalty, penalty], axis=0)
        move_logits = (logits[:8] + move_penalty).reshape(-1)
        pass_logit = jnp.mean(logits[8])[None]
        return jnp.concatenate([move_logits, pass_logit])


# ---------------------------------------------------------------------------
# Action selection (matches training/evaluation rollout)
# ---------------------------------------------------------------------------


@jax.jit
def competition_move_mask(obs: Observation) -> jnp.ndarray:
    """Legal move mask, identical to the one used during training/evaluation."""
    return compute_valid_move_mask(obs.armies, obs.owned_cells, obs.mountains)


def _decode_policy_index(index: int, height: int, width: int) -> np.ndarray:
    cells = height * width
    if index == 8 * cells:
        return PASS_ACTION.copy()
    encoded_direction, position = divmod(index, cells)
    row, col = divmod(position, width)
    split = int(encoded_direction >= 4)
    direction = encoded_direction - 4 if split else encoded_direction
    return np.array([0, row, col, direction, split], dtype=np.int32)


def observation_from_wire(
    timestep: int,
    owned_land_count: int,
    owned_army_count: int,
    opponent_land_count: int,
    opponent_army_count: int,
    type_grid,
    owner_grid,
    army_grid,
) -> Observation:
    """Convert a competition wire frame into the network's Observation type."""
    types = jnp.asarray(type_grid, dtype=jnp.int32)
    owners = jnp.asarray(owner_grid, dtype=jnp.int32)
    armies = jnp.asarray(army_grid, dtype=jnp.int32)
    visible = (types != 0) & (types != 5)
    return Observation(
        armies=armies,
        generals=types == 4,
        castles=types == 3,
        mountains=types == 2,
        neutral_cells=visible & (owners == 0) & (types != 2),
        owned_cells=owners == 1,
        opponent_cells=owners == 2,
        fog_cells=types == 0,
        structures_in_fog=types == 5,
        owned_land_count=jnp.int32(owned_land_count),
        owned_army_count=jnp.int32(owned_army_count),
        opponent_land_count=jnp.int32(opponent_land_count),
        opponent_army_count=jnp.int32(opponent_army_count),
        timestep=jnp.int32(timestep),
    )


@jax.jit
def competition_step(state: GameState, actions: jnp.ndarray):
    state, actions = build_castles.apply_build_actions(state, actions)
    return deathtouch.step(state, actions, turn=800)


class PolicyAgent:
    """Direct policy agent that samples actions as in training/evaluation.

    On each turn it computes the masked policy logits over all legal moves and
    samples one action with `jrandom.categorical`, exactly mirroring
    `PolicyValueNetwork.__call__` used during rollout collection.
    """

    def __init__(
        self,
        network: PolicyValueNetwork,
        seed: int = 0,
    ):
        self.network = network
        self.key = jrandom.PRNGKey(seed)
        self._infer = eqx.filter_jit(
            lambda network, tensor, mask: network.policy_value(tensor, mask)
        )

    def warmup(self, height: int, width: int) -> None:
        """Compile shape-dependent network and competition transition kernels."""
        grid = jnp.zeros((height, width), dtype=jnp.int32)
        grid = grid.at[0, 0].set(1).at[height - 1, width - 1].set(2)
        state = game.create_initial_state(grid)
        state = state._replace(
            armies=state.armies.at[0, 0].set(100).at[height - 1, width - 1].set(100),
            time=jnp.int32(100),
        )
        obs = game.get_observation(state, 0)
        self.policy_value(obs)
        actions = jnp.stack([jnp.asarray(PASS_ACTION), jnp.asarray(PASS_ACTION)])
        warmed_state, _ = competition_step(state, actions)
        jax.block_until_ready(warmed_state)

    def policy_value(self, obs: Observation):
        mask = competition_move_mask(obs)
        logits, value = self._infer(self.network, obs_to_tensor(obs), mask)
        return np.asarray(logits), float(value), mask

    def select(self, obs: Observation) -> tuple[np.ndarray, dict[str, float]]:
        """Sample one action from the masked policy and report statistics."""
        started_at = time.perf_counter()
        logits, value, _ = self.policy_value(obs)
        self.key, sample_key = jrandom.split(self.key)
        index = int(jrandom.categorical(sample_key, jnp.asarray(logits)))
        height, width = obs.armies.shape
        action = _decode_policy_index(index, height, width)
        elapsed_ms = (time.perf_counter() - started_at) * 1000
        return action, {
            "value": value,
            "elapsed_ms": elapsed_ms,
        }


# ---------------------------------------------------------------------------
# Stdio driver
# ---------------------------------------------------------------------------


def _read_grid(stream, height: int):
    return [[int(value) for value in stream.readline().split()] for _ in range(height)]


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("checkpoint", help="Equinox policy checkpoint")
    # Accepted for compatibility with the original run.sh; the direct policy
    # agent does no search, so these are ignored.
    parser.add_argument("--time-budget-ms", type=float, default=125.0)
    parser.add_argument("--max-simulations", type=int, default=128)
    parser.add_argument("--rollout-depth", type=int, default=2)
    parser.add_argument("--top-k", type=int, default=20)
    parser.add_argument("--seed", type=int, default=0)
    return parser.parse_args()


def main():
    args = parse_args()
    handshake = sys.stdin.readline()
    if not handshake:
        return
    player_index, height, width = (int(value) for value in handshake.split())

    network = PolicyValueNetwork(jrandom.PRNGKey(args.seed))
    network = eqx.tree_deserialise_leaves(args.checkpoint, network)
    agent = PolicyAgent(network, seed=args.seed)
    agent.warmup(height, width)
    print(f"[competition] warmup complete for {height}x{width}", file=sys.stderr, flush=True)

    while True:
        scalar_line = sys.stdin.readline()
        if not scalar_line:
            return
        timestep, own_land, own_army, opponent_land, opponent_army = (
            int(value) for value in scalar_line.split()
        )
        type_grid = _read_grid(sys.stdin, height)
        owner_grid = _read_grid(sys.stdin, height)
        army_grid = _read_grid(sys.stdin, height)
        observation = observation_from_wire(
            timestep,
            own_land,
            own_army,
            opponent_land,
            opponent_army,
            type_grid,
            owner_grid,
            army_grid,
        )

        started_at = time.perf_counter()
        action, stats = agent.select(observation)
        elapsed_ms = (time.perf_counter() - started_at) * 1000
        print(
            f"[competition] turn={timestep} value={stats['value']:+.3f} "
            f"elapsed_ms={elapsed_ms:.1f}",
            file=sys.stderr,
            flush=True,
        )
        print(" ".join(str(int(value)) for value in action), flush=True)


if __name__ == "__main__":
    main()