import { boardScore } from "./engine";
import {
  CARDS_PER_COLOR,
  DECK_DRAWS,
  HAND_SIZE,
  LOC_DECK,
  LOC_DISCARD,
  LOC_P0_BOARD,
  LOC_P0_HAND,
  MAX_ABS_SCORE,
  MAX_STEPS,
  N_CARDS,
  N_COLORS,
  OBS_DIM,
  type GameState,
} from "./types";
import { cardColor } from "./cards";

export function observation(state: GameState, player: number): Float32Array {
  const opponent = 1 - player;
  const myHandLoc = LOC_P0_HAND + player;
  const opponentHandLoc = LOC_P0_HAND + opponent;
  const myBoardLoc = LOC_P0_BOARD + player;
  const opponentBoardLoc = LOC_P0_BOARD + opponent;
  const pileTops = state.piles.map((pile) => pile.at(-1) ?? -1);
  const values: number[] = [];
  let hiddenOpponentCards = 0;

  for (let card = 0; card < N_CARDS; card += 1) {
    const loc = state.cardLoc[card];
    const isDiscard = loc === LOC_DISCARD;
    const isDiscardTop = isDiscard && pileTops[cardColor(card)] === card;
    const isOpponentHand = loc === opponentHandLoc;
    const isOpponentPublic = isOpponentHand && state.handPublic[card];
    if (isOpponentHand && !state.handPublic[card]) hiddenOpponentCards += 1;
    values.push(
      Number(loc === myHandLoc),
      Number(loc === myBoardLoc),
      Number(loc === opponentBoardLoc),
      Number(isDiscardTop),
      Number(isDiscard && !isDiscardTop),
      Number(isOpponentPublic),
      Number(loc === LOC_DECK || (isOpponentHand && !state.handPublic[card])),
    );
  }

  const scores = boardScore(state);
  values.push(
    (N_CARDS - state.drawPtr) / DECK_DRAWS,
    hiddenOpponentCards / HAND_SIZE,
    state.stepCount / MAX_STEPS,
    ...state.colTop[player].map((value) => value / 10),
    ...state.colTop[opponent].map((value) => value / 10),
    ...state.colHandshakes[player].map((value) => value / 3),
    ...state.colHandshakes[opponent].map((value) => value / 3),
    ...state.colLength[player].map((value) => value / CARDS_PER_COLOR),
    ...state.colLength[opponent].map((value) => value / CARDS_PER_COLOR),
    (scores[player] - scores[opponent]) / MAX_ABS_SCORE,
  );
  if (values.length !== OBS_DIM) throw new Error(`observation length ${values.length} != ${OBS_DIM}`);
  return Float32Array.from(values);
}
