"""Policy-value network and observation encoding for Generals.io PPO."""

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jrandom


def obs_to_tensor(obs) -> 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 value(self, obs):
        return self._value_from_features(self._features(obs))

    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])

    def __call__(self, obs, mask, key, action=None):
        """
        Args:
            obs: (C, H, W) tensor.
            mask: (H, W, 4) valid-move mask.
            key: PRNG key (only used when action is None).
            action: if given, evaluate logprob of this [pass,row,col,dir,split];
                otherwise sample an action.

        Returns:
            (action, value, logprob, entropy)
        """
        logits, value = self.policy_value(obs, mask)

        H, W = mask.shape[:2]
        cells = H * W

        if action is None:
            idx = jrandom.categorical(key, logits)
        else:
            is_pass, row, col, direction, is_half = action
            encoded_dir = jnp.where(is_half > 0, direction + 4, direction)
            move_idx = encoded_dir * cells + row * W + col
            idx = jnp.where(is_pass > 0, 8 * cells, move_idx)

        log_probs = jax.nn.log_softmax(logits)
        logprob = log_probs[idx]
        probs = jax.nn.softmax(logits)
        entropy = -jnp.sum(probs * log_probs)

        if action is None:
            is_pass = idx == 8 * cells
            move_idx = jnp.minimum(idx, 8 * cells - 1)
            direction = move_idx // cells
            position = move_idx % cells
            row = jnp.where(is_pass, 0, position // W)
            col = jnp.where(is_pass, 0, position % W)
            is_half = (~is_pass) & (direction >= 4)
            actual_dir = jnp.where(is_pass, 0, jnp.where(is_half, direction - 4, direction))
            action = jnp.array(
                [is_pass.astype(jnp.int32), row, col, actual_dir, is_half.astype(jnp.int32)],
                dtype=jnp.int32,
            )

        return action, value, logprob, entropy