/**
 * The observation and action *shapes*, and the encoder every trainer was
 * writing by hand.
 *
 * `ArenaManifest` has always published `observationFields` — a list of names.
 * A name list answers "what does this environment see"; it does not answer the
 * only question a trainer asks before it can allocate anything, which is "how
 * many numbers is that, and what are they". So every consumer wrote its own
 * encoder against the environment's TypeScript source: a fork of the
 * observation contract, kept outside this repository, silently invalidated by
 * any field this repository adds.
 *
 * `ArenaFieldSpec` (in `types.ts`) closes that. Each field declares how it
 * becomes numbers, `flattenObservation` performs exactly that encoding, and the
 * width is a function of the manifest rather than of a comment.
 *
 * ### Raw values, declared bounds
 *
 * A float lands in the vector as its **clamped raw value**, not as a 0..1
 * normalisation of it. Two reasons. Normalisation is a modelling decision — a
 * layer norm, a running mean, a tanh squash are all defensible and the trainer
 * is the only party that knows which it is using — and baking one in here would
 * make it impossible to recover the metre or the decibel on the other side.
 * And `low`/`high` are *declared* bounds rather than guarantees: they are what
 * an author believes the field spans, published so a consumer can normalise,
 * and quietly wrong bounds should show up as a saturated input rather than as
 * a silently rescaled one.
 *
 * ### Why this module imports the manifests directly
 *
 * `index.ts` re-exports this file, so importing `ARENA_MANIFESTS` from there
 * would close a cycle. Six named imports instead, which is the same list
 * `index.ts` itself keeps, and `src/test/arena/spaces.test.ts` asserts the two
 * catalogues hold the same ids so the duplication cannot drift.
 */

import { deriveArenaSeed } from "./random.ts";
import { CALIFORNIA_FLIGHT_MANIFEST } from "./californiaFlight.ts";
import { CROW_NAV_MANIFEST } from "./crowNav.ts";
import { DRIVE_101_MANIFEST } from "./drive101.ts";
import { OFFICE_JOBS_MANIFEST } from "./officeJobs.ts";
import { OFFICE_NAV_MANIFEST } from "./officeNav.ts";
import { STUDIO_OPS_MANIFEST } from "./studioOps.ts";
import type { ArenaFieldSpec, ArenaManifest } from "./types.ts";

const BY_ID: ReadonlyMap<string, ArenaManifest> = new Map(
  [
    DRIVE_101_MANIFEST,
    OFFICE_NAV_MANIFEST,
    OFFICE_JOBS_MANIFEST,
    CROW_NAV_MANIFEST,
    CALIFORNIA_FLIGHT_MANIFEST,
    STUDIO_OPS_MANIFEST,
  ].map((manifest) => [manifest.id, manifest]),
);

/** Every environment id this package can encode for, in catalogue order. */
export function arenaEnvironmentIds(): readonly string[] {
  return [...BY_ID.keys()];
}

/** The manifest for an id. Throws rather than returning `undefined`: a caller
 * that mistyped an env id wants to know now, not to receive a zero-width
 * vector fifty thousand steps into a run. */
export function arenaManifest(envId: string): ArenaManifest {
  const manifest = BY_ID.get(envId);
  if (!manifest) throw new RangeError(`unknown arena environment: ${envId}`);
  return manifest;
}

/**
 * How many slots one field occupies.
 *
 * This is also where a malformed spec is caught, because it is on the path of
 * every width and every encode: a float without finite ordered bounds and an
 * enum without a usable vocabulary are authoring mistakes that would otherwise
 * surface as a vector of the wrong length or a slot that is always zero.
 */
export function arenaFieldWidth(spec: ArenaFieldSpec): number {
  switch (spec.kind) {
    case "float": {
      const { low, high } = spec;
      if (
        typeof low !== "number" || typeof high !== "number" ||
        !Number.isFinite(low) || !Number.isFinite(high) || !(low < high)
      ) throw new RangeError(`arena float field ${spec.name} needs finite low < high`);
      return 1;
    }
    case "bool":
    case "id":
      return 1;
    case "enum": {
      const values = spec.values;
      if (!Array.isArray(values) || values.length === 0) {
        throw new RangeError(`arena enum field ${spec.name} needs a non-empty vocabulary`);
      }
      if (new Set(values).size !== values.length) {
        throw new RangeError(`arena enum field ${spec.name} has a duplicate value`);
      }
      return values.length;
    }
  }
}

/** Total slots in a flattened observation for this environment. */
export function observationWidth(envId: string): number {
  return spaceWidth(arenaManifest(envId).observationSpace);
}

/** Total slots in a flattened action for this environment. */
export function actionWidth(envId: string): number {
  return spaceWidth(arenaManifest(envId).actionSpace);
}

function spaceWidth(space: readonly ArenaFieldSpec[]): number {
  let width = 0;
  for (const spec of space) width += arenaFieldWidth(spec);
  return width;
}

/**
 * One observation as a fixed-width vector of finite numbers.
 *
 * Fixed-width is the whole contract and it holds unconditionally: a field the
 * observation does not carry, carries as `undefined`, or carries as the wrong
 * type still occupies its slots. That is not leniency, it is the property a
 * trainer depends on — a vector whose length changes at step 400 because a
 * station id went null is a crash a long way from its cause, and a batch that
 * silently absorbs a `NaN` is worse than one that never sees the value at all.
 *
 * What an unusable value encodes as is chosen to be *inert* rather than
 * plausible: a float falls back to its declared `low`, a bool to 0, an enum to
 * all-zeros, an id to 0. None of those can be mistaken for a measurement, and
 * `observationFields` plus the environment's own tests are what keep the case
 * from arising.
 */
export function flattenObservation(envId: string, observation: unknown): number[] {
  const space = arenaManifest(envId).observationSpace;
  const record = isRecord(observation) ? observation : {};
  const vector: number[] = [];
  for (const spec of space) encodeField(spec, record[spec.name], vector);
  return vector;
}

/** The same encoding applied to an action struct, for a trainer logging what
 * it did as well as what it saw. */
export function flattenAction(envId: string, action: unknown): number[] {
  const space = arenaManifest(envId).actionSpace;
  const record = isRecord(action) ? action : {};
  const vector: number[] = [];
  for (const spec of space) encodeField(spec, record[spec.name], vector);
  return vector;
}

/**
 * The inverse of `flattenAction`: a policy's output vector as the struct
 * `step()` takes.
 *
 * A harness that can read an observation vector but cannot emit an action is
 * half a harness, and the half it is missing is the one every consumer has to
 * hand-write against a field order it read out of a source file.
 *
 * Only `float` and `bool` are accepted, because that is what the six action
 * spaces contain and because the alternatives are worse than a refusal: an
 * `enum` action would need an argmax whose tie-break this module would be
 * inventing, and an `id` action cannot be inverted from a hash at all. If an
 * environment ever wants a categorical action, this function should learn about
 * it deliberately rather than by falling through to a default.
 *
 * Floats are clamped into their declared bounds and a bool is `value >= 0.5`,
 * so a raw network output needs no squashing before it gets here.
 */
export function structureAction(
  envId: string,
  vector: readonly number[],
): Record<string, number | boolean> {
  const space = arenaManifest(envId).actionSpace;
  const width = spaceWidth(space);
  if (vector.length !== width) {
    throw new RangeError(`arena ${envId} action needs exactly ${width} values`);
  }
  const action: Record<string, number | boolean> = {};
  let cursor = 0;
  for (const spec of space) {
    const value = vector[cursor] ?? 0;
    cursor += arenaFieldWidth(spec);
    if (spec.kind === "float") {
      action[spec.name] = Number.isFinite(value)
        ? clamp(value, spec.low as number, spec.high as number)
        : (spec.low as number);
    } else if (spec.kind === "bool") {
      action[spec.name] = Number.isFinite(value) && value >= 0.5;
    } else {
      throw new RangeError(`arena ${envId} action field ${spec.name} is not invertible`);
    }
  }
  return action;
}

function encodeField(spec: ArenaFieldSpec, value: unknown, into: number[]): void {
  // Validated first, always: `arenaFieldWidth` is where a malformed spec is
  // caught, and reading `spec.low` before asking whether it is a number would
  // be the one path on which a bad manifest reached a vector instead of a throw.
  arenaFieldWidth(spec);
  switch (spec.kind) {
    case "float": {
      const low = spec.low as number;
      const high = spec.high as number;
      into.push(typeof value === "number" && Number.isFinite(value) ? clamp(value, low, high) : low);
      return;
    }
    case "bool":
      into.push(value === true ? 1 : 0);
      return;
    case "enum": {
      for (const member of spec.values as readonly string[]) {
        into.push(value === member ? 1 : 0);
      }
      return;
    }
    case "id":
      into.push(typeof value === "string" && value.length > 0 ? idHash(value) : 0);
      return;
  }
}

/**
 * A string as a stable number in `[0, 1)`.
 *
 * `deriveArenaSeed` rather than a second hash written here: it is already the
 * package's string mixer, it is already exercised by every scenario
 * materialization, and two hashes over identifiers is two things to keep in
 * step. The constant is an arbitrary non-zero base — it exists so that this
 * stream is not the same one scenario seeds are drawn from.
 */
function idHash(value: string): number {
  return deriveArenaSeed(0x9e3779b9, value) / 4_294_967_296;
}

function clamp(value: number, low: number, high: number): number {
  return Math.min(high, Math.max(low, value));
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}
