import {
  ActorController,
  type ActorControllerSnapshot,
} from "../actors/controller.ts";
import { BaseArenaEnvironment, type SimulationTransition } from "./base.ts";
import { ARENA_SOURCE_HASHES } from "./sourceHashes.ts";
import { ArenaScenarioRegistry } from "./scenarios.ts";
import {
  ARENA_API_VERSION,
  type ArenaFieldSpec,
  type ArenaManifest,
  type ArenaScenario,
} from "./types.ts";

export interface CrowNavAction {
  forward: number;
  turn: number;
  pitch: number;
  climb: number;
  glide: boolean;
}

export interface CrowNavObservation {
  x: number;
  y: number;
  z: number;
  yaw: number;
  pitch: number;
  speedMps: number;
  verticalSpeedMps: number;
  goalX: number;
  goalY: number;
  goalZ: number;
  deltaX: number;
  deltaY: number;
  deltaZ: number;
  distanceToGoalM: number;
  altitudeBoundContact: "none" | "minimum" | "maximum";
}

export type CrowNavReward = Record<
  "progress" | "success" | "time" | "energy" | "heading" | "safety",
  number
>;

interface CrowScenarioParameters {
  startX: number;
  startY: number;
  startZ: number;
  startYaw: number;
  goalX: number;
  goalY: number;
  goalZ: number;
}

interface CrowSimulationSnapshot {
  controller: ActorControllerSnapshot;
  previousGoalDistanceM: number;
}

const FIXED_STEP = 0.1;
const MAX_STEPS = 300;
const SUCCESS_RADIUS_M = 2.5;
const BOUNDS = Object.freeze({ minX: -120, maxX: 120, minZ: -120, maxZ: 120 });

const DEFINITIONS = [
  {
    id: "train-east-crosswind",
    split: "train" as const,
    parameters: { startX: 0, startY: 12, startZ: 0, startYaw: 0, goalX: 60, goalY: 12, goalZ: 0 },
  },
  {
    id: "train-north-climb",
    split: "train" as const,
    parameters: { startX: 5, startY: 10, startZ: 20, startYaw: -1, goalX: 5, goalY: 16, goalZ: -48 },
  },
  {
    id: "dev-west-return",
    split: "dev" as const,
    parameters: { startX: 25, startY: 15, startZ: -10, startYaw: 0.4, goalX: -48, goalY: 13, goalZ: -10 },
  },
  {
    id: "dev-south-descent",
    split: "dev" as const,
    parameters: { startX: -15, startY: 18, startZ: -40, startYaw: -0.8, goalX: -5, goalY: 11, goalZ: 38 },
  },
] as const;

export const CROW_NAV_SCENARIOS = new ArenaScenarioRegistry<CrowScenarioParameters>(
  "crow-nav-v1",
  DEFINITIONS,
  (base, random) => ({
    ...base,
    startX: base.startX + random.between(-1, 1),
    startY: base.startY + random.between(-0.4, 0.4),
    startZ: base.startZ + random.between(-1, 1),
    startYaw: base.startYaw + random.between(-0.08, 0.08),
    goalX: base.goalX + random.between(-1, 1),
    goalY: base.goalY + random.between(-0.4, 0.4),
    goalZ: base.goalZ + random.between(-1, 1),
  }),
);

export const CROW_NAV_MANIFEST: ArenaManifest = Object.freeze({
  apiVersion: ARENA_API_VERSION,
  id: "crow-nav-v1",
  // 2 rather than 1: `ArenaScenarioRegistry` now selects a bare `{ split }`
  // request by hashing the seed against each scenario id instead of indexing
  // definition order, so a seed run before that change may resolve to a
  // different scenario after it. `version` is what a snapshot, a trace and a
  // results table are pinned to, and a selection change that nothing recorded
  // is exactly the silent remap the new selector exists to prevent.
  version: 2,
  title: "Crow waypoint navigation",
  description: "Three-dimensional waypoint control over Tera's deterministic crow flight controller.",
  simulator: "ActorController(kind=crow, mode=flight)",
  fixedStepSeconds: FIXED_STEP,
  maxSteps: MAX_STEPS,
  actionFields: ["forward", "turn", "pitch", "climb", "glide"],
  observationFields: [
    "x", "y", "z", "yaw", "pitch", "speedMps", "verticalSpeedMps",
    "goalX", "goalY", "goalZ", "deltaX", "deltaY", "deltaZ",
    "distanceToGoalM", "altitudeBoundContact",
  ],
  actionSpace: [
    { name: "forward", kind: "float", low: -1, high: 1, unit: "fraction" },
    { name: "turn", kind: "float", low: -1, high: 1, unit: "fraction" },
    { name: "pitch", kind: "float", low: -1, high: 1, unit: "fraction" },
    { name: "climb", kind: "float", low: -1, high: 1, unit: "fraction" },
    { name: "glide", kind: "bool" },
  ] satisfies readonly ArenaFieldSpec[],
  // The horizontal bounds are `BOUNDS` above and the altitude band is the
  // controller's own 2..40 m envelope, both stated here rather than restated:
  // a policy normalising against a wider box than the simulator enforces learns
  // a state distribution the environment never produces.
  observationSpace: [
    { name: "x", kind: "float", low: -120, high: 120, unit: "m" },
    { name: "y", kind: "float", low: 0, high: 40, unit: "m" },
    { name: "z", kind: "float", low: -120, high: 120, unit: "m" },
    { name: "yaw", kind: "float", low: -Math.PI, high: Math.PI, unit: "rad" },
    { name: "pitch", kind: "float", low: -Math.PI / 2, high: Math.PI / 2, unit: "rad" },
    { name: "speedMps", kind: "float", low: 0, high: 30, unit: "m/s" },
    { name: "verticalSpeedMps", kind: "float", low: -20, high: 20, unit: "m/s" },
    { name: "goalX", kind: "float", low: -120, high: 120, unit: "m" },
    { name: "goalY", kind: "float", low: 0, high: 40, unit: "m" },
    { name: "goalZ", kind: "float", low: -120, high: 120, unit: "m" },
    { name: "deltaX", kind: "float", low: -240, high: 240, unit: "m" },
    { name: "deltaY", kind: "float", low: -40, high: 40, unit: "m" },
    { name: "deltaZ", kind: "float", low: -240, high: 240, unit: "m" },
    { name: "distanceToGoalM", kind: "float", low: 0, high: 350, unit: "m" },
    { name: "altitudeBoundContact", kind: "enum", values: ["none", "minimum", "maximum"] },
  ] satisfies readonly ArenaFieldSpec[],
  rewardComponents: {
    progress: "Reduction in 3D distance to the waypoint.",
    success: "Sparse waypoint completion bonus.",
    time: "Per-step pressure that makes minimum-power inaction negative.",
    energy: "Counterweight on powered speed and control demand.",
    heading: "Small cost for pointing away from the target.",
    safety: "Terminal flight-envelope contact penalty.",
  },
  safetyTerminals: ["flight-envelope-contact"],
  scenarioIds: {
    train: CROW_NAV_SCENARIOS.ids("train"),
    dev: CROW_NAV_SCENARIOS.ids("dev"),
  },
  baselines: {
    inaction: "Minimum-power straight flight (forward=-1) misses the lateral waypoint and stays below zero return.",
    scripted: "Proportional yaw and climb guidance reaches every public waypoint.",
  },
});

export const CROW_NAV_INACTION: Readonly<CrowNavAction> = Object.freeze({
  forward: -1, turn: 0, pitch: 0, climb: 0.03, glide: false,
});

function wrapAngle(value: number): number {
  return ((value + Math.PI) % (Math.PI * 2) + Math.PI * 2) % (Math.PI * 2) - Math.PI;
}

export function crowNavScriptedBaseline(observation: CrowNavObservation): CrowNavAction {
  const desiredYaw = Math.atan2(-observation.deltaX, -observation.deltaZ);
  const headingError = wrapAngle(desiredYaw - observation.yaw);
  return {
    forward: 0.45,
    turn: Math.max(-1, Math.min(1, headingError / 0.55)),
    pitch: 0,
    climb: Math.max(-0.8, Math.min(0.8, observation.deltaY / 10 - 0.2)),
    glide: false,
  };
}

export class CrowNavEnvironment extends BaseArenaEnvironment<
  CrowNavAction,
  CrowNavObservation,
  CrowNavReward,
  CrowSimulationSnapshot,
  CrowScenarioParameters
> {
  readonly manifest = CROW_NAV_MANIFEST;
  protected readonly registry = CROW_NAV_SCENARIOS;
  protected readonly sourceHashes = ARENA_SOURCE_HASHES["crow-nav-v1"]!;
  private controller: ActorController | null = null;
  private previousGoalDistanceM = 0;

  protected resetSimulation(scenario: ArenaScenario<CrowScenarioParameters>): CrowNavObservation {
    const parameters = scenario.parameters;
    this.controller = new ActorController({
      kind: "crow",
      mode: "flight",
      identity: {
        id: "arena-crow", displayName: "Arena Crow", authenticated: false, profile: {},
      },
      position: { x: parameters.startX, y: parameters.startY, z: parameters.startZ },
      yaw: parameters.startYaw,
      groundY: 0,
      minFlightAltitude: 2,
      maxFlightAltitude: 40,
      horizontalBounds: BOUNDS,
      fixedStepSeconds: FIXED_STEP,
    });
    this.previousGoalDistanceM = this.goalDistance();
    return this.observation();
  }

  protected normalizeAction(action: CrowNavAction): CrowNavAction {
    const axis = (value: number) => Math.max(-1, Math.min(1, Number.isFinite(value) ? value : 0));
    return {
      forward: axis(action?.forward),
      turn: axis(action?.turn),
      pitch: axis(action?.pitch),
      climb: axis(action?.climb),
      glide: action?.glide === true,
    };
  }

  protected advanceSimulation(action: CrowNavAction): SimulationTransition<CrowNavObservation, CrowNavReward> {
    const controller = this.requireController();
    controller.stepFixed({
      ...action,
      right: 0,
      sprint: false,
      modeRequest: "none",
      kindRequest: "none",
      reset: false,
    });
    const state = controller.state();
    const distance = this.goalDistance();
    const progress = this.previousGoalDistanceM - distance;
    this.previousGoalDistanceM = distance;
    const success = distance <= SUCCESS_RADIUS_M;
    const horizontalContact =
      state.x <= BOUNDS.minX || state.x >= BOUNDS.maxX ||
      state.z <= BOUNDS.minZ || state.z >= BOUNDS.maxZ;
    const safety = horizontalContact || state.altitudeBoundContact !== "none";
    const observation = this.observation();
    const desiredYaw = Math.atan2(-observation.deltaX, -observation.deltaZ);
    const headingError = Math.abs(wrapAngle(desiredYaw - observation.yaw));
    return {
      observation,
      rewardComponents: {
        progress: progress * 0.045,
        success: success ? 3 : 0,
        time: -0.008,
        energy: -0.001 * ((action.forward + 1) / 2 + Math.abs(action.turn) +
          Math.abs(action.pitch) + Math.abs(action.climb)),
        heading: -0.0008 * (headingError / Math.PI) ** 2,
        safety: safety ? -3 : 0,
      },
      terminated: success || safety,
      terminalReason: success ? "goal" : safety ? "flight-envelope-contact" : undefined,
    };
  }

  protected simulationSnapshot(): CrowSimulationSnapshot {
    return {
      controller: this.requireController().snapshot(),
      previousGoalDistanceM: this.previousGoalDistanceM,
    };
  }

  protected restoreSimulation(snapshot: CrowSimulationSnapshot): CrowNavObservation {
    this.resetSimulation(this.currentScenario());
    this.requireController().restore(snapshot.controller);
    this.previousGoalDistanceM = snapshot.previousGoalDistanceM;
    return this.observation();
  }

  private observation(): CrowNavObservation {
    const state = this.requireController().state();
    const goal = this.currentScenario().parameters;
    return {
      x: state.x,
      y: state.y,
      z: state.z,
      yaw: state.yaw,
      pitch: state.pitch,
      speedMps: state.speedMps,
      verticalSpeedMps: state.verticalSpeedMps,
      goalX: goal.goalX,
      goalY: goal.goalY,
      goalZ: goal.goalZ,
      deltaX: goal.goalX - state.x,
      deltaY: goal.goalY - state.y,
      deltaZ: goal.goalZ - state.z,
      distanceToGoalM: this.goalDistance(),
      altitudeBoundContact: state.altitudeBoundContact,
    };
  }

  private goalDistance(): number {
    const state = this.requireController().state();
    const goal = this.currentScenario().parameters;
    const dx = goal.goalX - state.x;
    const dy = goal.goalY - state.y;
    const dz = goal.goalZ - state.z;
    return Math.sqrt(dx * dx + dy * dy + dz * dz);
  }

  private requireController(): ActorController {
    if (!this.controller) throw new Error("crow environment has not been reset");
    return this.controller;
  }
}
