"""Competition stdio agent using particle PUCT and a trained checkpoint."""

import argparse
import os
import sys
import time

import equinox as eqx
import jax.random as jrandom

from general_bots_training.mcts import MCTSConfig, ParticlePUCT, observation_from_wire
from general_bots_training.network import PolicyValueNetwork

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


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")
    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)
    search = ParticlePUCT(
        network,
        player_index,
        MCTSConfig(
            time_budget_ms=args.time_budget_ms,
            max_simulations=args.max_simulations,
            rollout_depth=args.rollout_depth,
            top_k=args.top_k,
        ),
        seed=args.seed,
    )
    search.warmup(height, width)
    print(f"[mcts] 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 = search.search(observation)
        elapsed_ms = (time.perf_counter() - started_at) * 1000
        print(
            f"[mcts] turn={timestep} simulations={int(stats['simulations'])} "
            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()