import { cardColor } from "./cards";
import { boardScore } from "./engine";
import { N_ROUNDS, type MatchState } from "./match";
import { observation } from "./observation";
import {
  CARDS_PER_COLOR,
  LOC_DECK,
  LOC_DISCARD,
  LOC_P0_HAND,
  N_CARDS,
  N_COLORS,
  OBS_DIM,
} from "./types";

/**
 * Player-view observation for a three-round match: the mirror of
 * src/lost_cities_jax/match_obs.py.
 *
 * These two must agree to the bit. A mismatch does not throw -- the model happily
 * consumes a wrong vector and plays worse for reasons nobody can see. That is why
 * matchObservation.test.ts checks this against fixtures generated from the Python
 * side rather than trusting the port.
 *
 * On top of the single-round observation this adds the four things a match policy
 * cannot play without: carry (scalar *and* binned, because round three is a
 * threshold problem); which round it is; whose turn it is; and the deck clock,
 * since a round ends on the last deck draw and players bend that parity by drawing
 * from discard piles. Live points per colour are split by hand / discard pile /
 * unseen, because a discard pile is public and recoverable.
 */

/** Packed tightly around zero: that is where the round-three decision flips. */
const CARRY_BIN_EDGES = [-60, -30, -12, -1, 1, 12, 30, 60];
export const N_CARRY_BINS = CARRY_BIN_EDGES.length + 1;

/** A typical round margin, not the theoretical 780 maximum the old obs divided by. */
const CARRY_SCALE = 75;

/** 2+3+...+10, the most one expedition can be worth before multipliers. */
const MAX_COLOR_POINTS = 54;

const N_MATCH_SCALARS =
  1 + N_CARRY_BINS + N_ROUNDS + 1 + 1 + 1 + 1 + 2 * N_COLORS * 3;
export const MATCH_OBS_DIM = OBS_DIM + N_MATCH_SCALARS;

function cardRank(card: number): number {
  const slot = card % CARDS_PER_COLOR;
  return slot >= 3 ? slot - 1 : 0;
}

/** numpy.digitize: the count of edges strictly below `value`. */
function digitize(value: number, edges: readonly number[]): number {
  let index = 0;
  while (index < edges.length && value >= edges[index]) index += 1;
  return index;
}

/**
 * Points still reachable for `subject`, split by where the card sits:
 * [in hand, in a discard pile, unseen] per colour.
 *
 * Seen through `viewer`'s eyes -- a card in the opponent's hand only counts as "in
 * hand" if it is public, otherwise it is unseen.
 */
function livePoints(state: MatchState, viewer: number, subject: number): number[] {
  const round = state.round;
  const isMine = subject === viewer;
  const subjectHandLoc = LOC_P0_HAND + subject;
  const out = Array.from({ length: N_COLORS }, () => [0, 0, 0]);

  for (let card = 0; card < N_CARDS; card += 1) {
    const color = cardColor(card);
    const rank = cardRank(card);
    // An ascending column can only take cards above its current top.
    if (rank <= round.colTop[subject][color]) continue;

    const loc = round.cardLoc[card];
    const inSubjectHand = loc === subjectHandLoc;
    if (inSubjectHand && (isMine || round.handPublic[card])) {
      out[color][0] += rank;
    } else if (loc === LOC_DISCARD) {
      out[color][1] += rank;
    } else if (loc === LOC_DECK || (inSubjectHand && !isMine && !round.handPublic[card])) {
      // Cards in the *other* player's hidden hand are unseen to the viewer too,
      // but the subject cannot reach them, so they are deliberately excluded.
      out[color][2] += rank;
    }
  }
  return out.flat().map((points) => points / MAX_COLOR_POINTS);
}

export function matchObservation(state: MatchState, player: number): Float32Array {
  const opponent = 1 - player;
  const round = state.round;
  const base = observation(round, player);

  const board = boardScore(round);
  const lead =
    state.carry[player] + board[player] - (state.carry[opponent] + board[opponent]);

  const carryScaled = Math.min(2, Math.max(-2, lead / CARRY_SCALE));
  const carryBins = Array.from({ length: N_CARRY_BINS }, (_, i) =>
    Number(i === digitize(lead, CARRY_BIN_EDGES)),
  );

  const roundOneHot = Array.from({ length: N_ROUNDS }, (_, i) => Number(i === state.roundIdx));
  const roundsLeft = (N_ROUNDS - 1 - state.roundIdx) / (N_ROUNDS - 1);

  const toMove = round.toMove;
  const myTurn = Number(toMove === player);
  // toMove flips every ply, so the round's opener is recoverable from parity.
  const roundOpener = toMove ^ (round.stepCount & 1);
  const iOpened = Number(roundOpener === player);

  // If both players drew from the deck from here, the last deck card falls to
  // whoever is on move after `remaining - 1` more plies.
  const remaining = Math.max(N_CARDS - round.drawPtr, 1);
  const lastDrawer = toMove ^ ((remaining - 1) & 1);
  const iTakeLast = Number(lastDrawer === player);

  const extra = [
    carryScaled,
    ...carryBins,
    ...roundOneHot,
    roundsLeft,
    myTurn,
    iOpened,
    iTakeLast,
    ...livePoints(state, player, player),
    ...livePoints(state, player, opponent),
  ];

  if (extra.length !== N_MATCH_SCALARS) {
    throw new Error(`match scalars ${extra.length} != ${N_MATCH_SCALARS}`);
  }
  const out = new Float32Array(MATCH_OBS_DIM);
  out.set(base, 0);
  out.set(Float32Array.from(extra), OBS_DIM);
  return out;
}
