#!/usr/bin/env python3
"""Generate match states and their observations so TypeScript can be checked against JAX.

The two observation builders must agree to the bit. A mismatch does not throw --
the ONNX policy consumes a wrong vector quite happily and plays worse for reasons
nobody can see. So the port is not trusted; it is checked.

States are drawn from real random play so the fixture covers the awkward parts:
mid-round, both seats to move, past a round roll-over, with a non-zero carry.
"""

from __future__ import annotations

import json
from pathlib import Path

import jax
import numpy as np

from lost_cities_jax.match import MatchState, match_reset_from, match_step
from lost_cities_jax.match_obs import MATCH_OBS_DIM, match_observation
from lost_cities_jax.opponents import random_legal_action
from lost_cities_jax.types import N_CARDS

OUTPUT = Path(__file__).resolve().parents[1] / "web" / "src" / "game" / "match-parity-fixture.json"
N_ROUNDS = 3


def match_json(match: MatchState) -> dict:
    round_state = match.round
    return {
        "round": {
            "deckOrder": np.asarray(round_state.deck_order).astype(int).tolist(),
            "drawPtr": int(round_state.draw_ptr),
            "cardLoc": np.asarray(round_state.card_loc).astype(int).tolist(),
            "handPublic": np.asarray(round_state.hand_public).astype(bool).tolist(),
            "colTop": np.asarray(round_state.col_top).astype(int).tolist(),
            "colHandshakes": np.asarray(round_state.col_hs).astype(int).tolist(),
            "colLength": np.asarray(round_state.col_len).astype(int).tolist(),
            "piles": [
                np.asarray(round_state.pile[color, : int(round_state.pile_len[color])])
                .astype(int)
                .tolist()
                for color in range(5)
            ],
            "toMove": int(round_state.to_move),
            "stepCount": int(round_state.step_count),
            "done": bool(round_state.done),
        },
        "deckOrders": np.asarray(match.deck_orders).astype(int).tolist(),
        "coinFlips": np.asarray(match.coin_flips).astype(int).tolist(),
        "roundIdx": int(match.round_idx),
        "carry": np.asarray(match.carry).astype(int).tolist(),
        "done": bool(match.done),
    }


def main() -> None:
    rng = np.random.default_rng(20260715)
    key = jax.random.PRNGKey(7)
    rows = []

    for match_index in range(6):
        decks = np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)])
        coins = rng.integers(0, 2, size=(N_ROUNDS,))
        match = match_reset_from(decks.astype(np.int8), coins.astype(np.int8))

        # Sample the opening position and then every 17th ply, which lands in all
        # three rounds and on both seats without hand-picking anything.
        ply = 0
        while not bool(match.done) and ply < 400:
            if ply % 17 == 0 or ply == 0:
                for player in (0, 1):
                    obs = np.asarray(match_observation(match, player), dtype=np.float64)
                    assert obs.shape == (MATCH_OBS_DIM,)
                    rows.append(
                        {
                            "match": match_json(match),
                            "player": player,
                            "observation": [round(float(v), 7) for v in obs],
                        }
                    )
            key, step_key = jax.random.split(key)
            action = int(random_legal_action(match.round, match.round.to_move, step_key))
            match, _, _ = match_step(match, action)
            ply += 1
        del match_index

    OUTPUT.write_text(
        json.dumps({"format": "jax-web-match-parity-v1", "obsDim": MATCH_OBS_DIM, "rows": rows})
        + "\n"
    )
    print(f"wrote {len(rows)} rows -> {OUTPUT}")


if __name__ == "__main__":
    main()