import {
  AircraftController,
  type AircraftControllerSnapshot,
  type AircraftGeographicPoint,
} from "../aircraft/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 CaliforniaFlightAction {
  throttle: number;
  yaw: number;
  pitch: number;
  roll: number;
}

export interface CaliforniaFlightObservation {
  lat: number;
  lng: number;
  altitudeM: number;
  headingDeg: number;
  pitchDeg: number;
  rollDeg: number;
  speedMps: number;
  verticalSpeedMps: number;
  goalLat: number;
  goalLng: number;
  goalAltitudeM: number;
  distanceToGoalM: number;
  bearingToGoalDeg: number;
  altitudeErrorM: number;
  envelopeContact: boolean;
}

export type CaliforniaFlightReward = Record<
  "progress" | "success" | "time" | "altitude" | "control" | "safety",
  number
>;

interface FlightScenarioParameters {
  startLat: number;
  startLng: number;
  startAltitudeM: number;
  startHeadingDeg: number;
  goalLat: number;
  goalLng: number;
  goalAltitudeM: number;
}

interface FlightSimulationSnapshot {
  controller: AircraftControllerSnapshot;
  previousGoalDistanceM: number;
}

const EARTH_RADIUS_M = 6_371_000;
const FIXED_STEP = 0.1;
const MAX_STEPS = 500;
const SUCCESS_RADIUS_M = 90;

const DEFINITIONS = [
  {
    id: "train-la-east-leg",
    split: "train" as const,
    parameters: {
      startLat: 34.0522, startLng: -118.2437, startAltitudeM: 1_200, startHeadingDeg: 0,
      goalLat: 34.0522, goalLng: -118.2325, goalAltitudeM: 1_200,
    },
  },
  {
    id: "train-bay-west-climb",
    split: "train" as const,
    parameters: {
      startLat: 37.70, startLng: -122.30, startAltitudeM: 1_350, startHeadingDeg: 15,
      goalLat: 37.70, goalLng: -122.313, goalAltitudeM: 1_450,
    },
  },
  {
    id: "dev-socal-southeast",
    split: "dev" as const,
    parameters: {
      startLat: 34.20, startLng: -118.40, startAltitudeM: 1_500, startHeadingDeg: 330,
      goalLat: 34.194, goalLng: -118.391, goalAltitudeM: 1_380,
    },
  },
  {
    id: "dev-bay-northeast",
    split: "dev" as const,
    parameters: {
      startLat: 37.62, startLng: -122.35, startAltitudeM: 1_100, startHeadingDeg: 280,
      goalLat: 37.628, goalLng: -122.341, goalAltitudeM: 1_260,
    },
  },
] as const;

export const CALIFORNIA_FLIGHT_SCENARIOS = new ArenaScenarioRegistry<FlightScenarioParameters>(
  "california-flight-v1",
  DEFINITIONS,
  (base, random) => {
    const northJitter = random.between(-0.00015, 0.00015);
    const eastJitter = random.between(-0.00015, 0.00015);
    return {
      ...base,
      startLat: base.startLat + northJitter,
      startLng: base.startLng + eastJitter,
      startAltitudeM: base.startAltitudeM + random.between(-12, 12),
      startHeadingDeg: base.startHeadingDeg + random.between(-3, 3),
      goalLat: base.goalLat - northJitter,
      goalLng: base.goalLng - eastJitter,
      goalAltitudeM: base.goalAltitudeM + random.between(-12, 12),
    };
  },
);

export const CALIFORNIA_FLIGHT_MANIFEST: ArenaManifest = Object.freeze({
  apiVersion: ARENA_API_VERSION,
  id: "california-flight-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: "California electric-flight waypoint",
  description: "Manual fixed-wing waypoint control over Tera's renderer-neutral aircraft simulator.",
  simulator: "AircraftController",
  fixedStepSeconds: FIXED_STEP,
  maxSteps: MAX_STEPS,
  actionFields: ["throttle", "yaw", "pitch", "roll"],
  observationFields: [
    "lat", "lng", "altitudeM", "headingDeg", "pitchDeg", "rollDeg", "speedMps",
    "verticalSpeedMps", "goalLat", "goalLng", "goalAltitudeM", "distanceToGoalM",
    "bearingToGoalDeg", "altitudeErrorM", "envelopeContact",
  ],
  actionSpace: [
    { name: "throttle", kind: "float", low: 0, high: 1, unit: "fraction" },
    { name: "yaw", kind: "float", low: -1, high: 1, unit: "fraction" },
    { name: "pitch", kind: "float", low: -1, high: 1, unit: "fraction" },
    { name: "roll", kind: "float", low: -1, high: 1, unit: "fraction" },
  ] satisfies readonly ArenaFieldSpec[],
  // Latitude and longitude are bounded to California rather than to the globe.
  // A ±90/±180 box would put every scenario in this environment inside four
  // decimal places of the same normalized value, which is a constant input
  // wearing a coordinate's clothes.
  observationSpace: [
    { name: "lat", kind: "float", low: 32, high: 42.2, unit: "deg north" },
    { name: "lng", kind: "float", low: -124.5, high: -114, unit: "deg east" },
    { name: "altitudeM", kind: "float", low: 0, high: 4000, unit: "m" },
    { name: "headingDeg", kind: "float", low: 0, high: 360, unit: "deg true" },
    { name: "pitchDeg", kind: "float", low: -90, high: 90, unit: "deg" },
    { name: "rollDeg", kind: "float", low: -90, high: 90, unit: "deg" },
    { name: "speedMps", kind: "float", low: 0, high: 260, unit: "m/s" },
    { name: "verticalSpeedMps", kind: "float", low: -60, high: 60, unit: "m/s" },
    { name: "goalLat", kind: "float", low: 32, high: 42.2, unit: "deg north" },
    { name: "goalLng", kind: "float", low: -124.5, high: -114, unit: "deg east" },
    { name: "goalAltitudeM", kind: "float", low: 0, high: 4000, unit: "m" },
    { name: "distanceToGoalM", kind: "float", low: 0, high: 20000, unit: "m" },
    { name: "bearingToGoalDeg", kind: "float", low: 0, high: 360, unit: "deg true" },
    { name: "altitudeErrorM", kind: "float", low: -2000, high: 2000, unit: "m" },
    { name: "envelopeContact", kind: "bool" },
  ] satisfies readonly ArenaFieldSpec[],
  rewardComponents: {
    progress: "Reduction in three-dimensional waypoint distance.",
    success: "Sparse arrival bonus.",
    time: "Per-step pressure; straight-ahead inaction misses lateral goals.",
    altitude: "Counterweight on altitude error.",
    control: "Small cost on throttle and surface demand.",
    safety: "Terminal California flight-envelope penalty.",
  },
  safetyTerminals: ["flight-envelope-contact"],
  scenarioIds: {
    train: CALIFORNIA_FLIGHT_SCENARIOS.ids("train"),
    dev: CALIFORNIA_FLIGHT_SCENARIOS.ids("dev"),
  },
  baselines: {
    inaction: "Zero surfaces and throttle fly straight past the lateral waypoint and remain below zero return.",
    scripted: "Proportional bearing, bank, yaw and altitude guidance completes every public scenario.",
  },
});

export const CALIFORNIA_FLIGHT_INACTION: Readonly<CaliforniaFlightAction> = Object.freeze({
  throttle: 0, yaw: 0, pitch: 0, roll: 0,
});

function horizontalDistanceM(a: AircraftGeographicPoint, b: AircraftGeographicPoint): number {
  const mean = (a.lat + b.lat) / 2 * Math.PI / 180;
  const north = (b.lat - a.lat) * Math.PI / 180 * EARTH_RADIUS_M;
  const east = (b.lng - a.lng) * Math.PI / 180 * Math.cos(mean) * EARTH_RADIUS_M;
  return Math.sqrt(north * north + east * east);
}

function distance3dM(
  a: AircraftGeographicPoint & { altitudeM: number },
  b: AircraftGeographicPoint & { altitudeM: number },
): number {
  const horizontal = horizontalDistanceM(a, b);
  const vertical = b.altitudeM - a.altitudeM;
  return Math.sqrt(horizontal * horizontal + vertical * vertical);
}

function bearingDeg(a: AircraftGeographicPoint, b: AircraftGeographicPoint): number {
  const mean = (a.lat + b.lat) / 2 * Math.PI / 180;
  return (Math.atan2((b.lng - a.lng) * Math.cos(mean), b.lat - a.lat) * 180 / Math.PI + 360) % 360;
}

function signedAngleDeg(from: number, to: number): number {
  return ((to - from + 540) % 360) - 180;
}

export function californiaFlightScriptedBaseline(
  observation: CaliforniaFlightObservation,
): CaliforniaFlightAction {
  const headingError = signedAngleDeg(observation.headingDeg, observation.bearingToGoalDeg);
  return {
    throttle: 0.62,
    yaw: Math.max(-0.6, Math.min(0.6, headingError / 80)),
    pitch: Math.max(-0.7, Math.min(0.7, observation.altitudeErrorM / 220)),
    roll: Math.max(-1, Math.min(1, headingError / 42)),
  };
}

export class CaliforniaFlightEnvironment extends BaseArenaEnvironment<
  CaliforniaFlightAction,
  CaliforniaFlightObservation,
  CaliforniaFlightReward,
  FlightSimulationSnapshot,
  FlightScenarioParameters
> {
  readonly manifest = CALIFORNIA_FLIGHT_MANIFEST;
  protected readonly registry = CALIFORNIA_FLIGHT_SCENARIOS;
  protected readonly sourceHashes = ARENA_SOURCE_HASHES["california-flight-v1"]!;
  private controller: AircraftController | null = null;
  private previousGoalDistanceM = 0;

  protected resetSimulation(scenario: ArenaScenario<FlightScenarioParameters>): CaliforniaFlightObservation {
    const parameters = scenario.parameters;
    const altitudeMin = Math.min(parameters.startAltitudeM, parameters.goalAltitudeM);
    const altitudeMax = Math.max(parameters.startAltitudeM, parameters.goalAltitudeM);
    this.controller = new AircraftController({
      initialPosition: { lat: parameters.startLat, lng: parameters.startLng },
      initialAltitudeM: parameters.startAltitudeM,
      initialHeadingDeg: parameters.startHeadingDeg,
      initialSpeedMps: 38,
      mode: "manual",
      fixedStepSeconds: FIXED_STEP,
      envelope: {
        minLat: Math.min(parameters.startLat, parameters.goalLat) - 0.02,
        maxLat: Math.max(parameters.startLat, parameters.goalLat) + 0.02,
        minLng: Math.min(parameters.startLng, parameters.goalLng) - 0.02,
        maxLng: Math.max(parameters.startLng, parameters.goalLng) + 0.02,
        minAltitudeM: altitudeMin - 260,
        maxAltitudeM: altitudeMax + 260,
      },
    });
    this.previousGoalDistanceM = this.goalDistance();
    return this.observation();
  }

  protected normalizeAction(action: CaliforniaFlightAction): CaliforniaFlightAction {
    const axis = (value: number) => Math.max(-1, Math.min(1, Number.isFinite(value) ? value : 0));
    return {
      throttle: Math.max(0, Math.min(1, Number.isFinite(action?.throttle) ? action.throttle : 0)),
      yaw: axis(action?.yaw),
      pitch: axis(action?.pitch),
      roll: axis(action?.roll),
    };
  }

  protected advanceSimulation(
    action: CaliforniaFlightAction,
  ): SimulationTransition<CaliforniaFlightObservation, CaliforniaFlightReward> {
    const controller = this.requireController();
    controller.stepFixed({ ...action, modeRequest: "manual", 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 safety = state.envelopeContact;
    const altitudeError = this.currentScenario().parameters.goalAltitudeM - state.altitudeM;
    return {
      observation: this.observation(),
      rewardComponents: {
        progress: progress * 0.0025,
        success: success ? 3.5 : 0,
        time: -0.006,
        altitude: -0.000002 * altitudeError ** 2,
        control: -0.0005 * (action.throttle ** 2 + action.yaw ** 2 +
          action.pitch ** 2 + action.roll ** 2),
        safety: safety ? -4 : 0,
      },
      terminated: success || safety,
      terminalReason: success ? "goal" : safety ? "flight-envelope-contact" : undefined,
    };
  }

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

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

  private observation(): CaliforniaFlightObservation {
    const state = this.requireController().state();
    const goal = this.currentScenario().parameters;
    return {
      lat: state.lat,
      lng: state.lng,
      altitudeM: state.altitudeM,
      headingDeg: state.headingDeg,
      pitchDeg: state.pitchDeg,
      rollDeg: state.rollDeg,
      speedMps: state.speedMps,
      verticalSpeedMps: state.verticalSpeedMps,
      goalLat: goal.goalLat,
      goalLng: goal.goalLng,
      goalAltitudeM: goal.goalAltitudeM,
      distanceToGoalM: this.goalDistance(),
      bearingToGoalDeg: bearingDeg(state, { lat: goal.goalLat, lng: goal.goalLng }),
      altitudeErrorM: goal.goalAltitudeM - state.altitudeM,
      envelopeContact: state.envelopeContact,
    };
  }

  private goalDistance(): number {
    const state = this.requireController().state();
    const goal = this.currentScenario().parameters;
    return distance3dM(state, {
      lat: goal.goalLat, lng: goal.goalLng, altitudeM: goal.goalAltitudeM,
    });
  }

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