import { isPlainCell, isPlainStruct, isPlainTensor, type PlainStruct, type PlainValue } from "../runner/plain";
import type { Field2D, FieldFrames } from "../fields/field";
import type { CompareResult, MethodPlan, ModeTable, PlateResult, ResultInfo } from "./types";

export class ResultShapeError extends Error {}

function field(s: PlainStruct, name: string): PlainValue {
  const v = s.fields[name];
  if (v === undefined) throw new ResultShapeError(`result is missing field "${name}"`);
  return v;
}

function struct(s: PlainStruct, name: string): PlainStruct {
  const v = field(s, name);
  if (!isPlainStruct(v)) throw new ResultShapeError(`field "${name}" should be a struct`);
  return v;
}

function num(s: PlainStruct, name: string): number {
  const v = field(s, name);
  if (typeof v === "number") return v;
  if (typeof v === "boolean") return v ? 1 : 0;
  if (isPlainTensor(v) && v.data.length === 1) return v.data[0];
  throw new ResultShapeError(`field "${name}" should be a scalar`);
}

function bool(s: PlainStruct, name: string): boolean {
  return num(s, name) !== 0;
}

function str(s: PlainStruct, name: string): string {
  const v = field(s, name);
  if (typeof v === "string") return v;
  throw new ResultShapeError(`field "${name}" should be a string`);
}

function vec(s: PlainStruct, name: string): Float64Array {
  const v = field(s, name);
  if (typeof v === "number") return Float64Array.of(v);
  if (isPlainTensor(v)) return v.data;
  throw new ResultShapeError(`field "${name}" should be a numeric vector`);
}

function field2d(s: PlainStruct, name: string): Field2D {
  const v = field(s, name);
  if (!isPlainTensor(v) || v.shape.length !== 2) throw new ResultShapeError(`field "${name}" should be a 2-D array`);
  return { ny: v.shape[0], nx: v.shape[1], data: v.data };
}

function frames3d(s: PlainStruct, name: string): FieldFrames {
  const v = field(s, name);
  if (!isPlainTensor(v)) throw new ResultShapeError(`field "${name}" should be a 3-D array`);
  const [ny, nx, nf] = [v.shape[0] ?? 0, v.shape[1] ?? 0, v.shape[2] ?? 1];
  if (v.shape.length < 2 || v.shape.length > 3) throw new ResultShapeError(`field "${name}" should be a 3-D array`);
  return { ny, nx, nf, data: v.data };
}

function strings(s: PlainStruct, name: string): string[] {
  const v = s.fields[name];
  if (v === undefined) return [];
  if (typeof v === "string") return [v];
  if (isPlainCell(v)) return v.items.map((x) => (typeof x === "string" ? x : JSON.stringify(x)));
  return [];
}

const INFO_STANDARD = new Set(["nsteps", "ncells", "dt", "fs", "wall_s", "stable", "diverged_step", "notes"]);

function parseInfo(s: PlainStruct): ResultInfo {
  const extra: Record<string, number | string | boolean> = {};
  for (const [k, v] of Object.entries(s.fields)) {
    if (INFO_STANDARD.has(k)) continue;
    if (typeof v === "number" || typeof v === "string" || typeof v === "boolean") extra[k] = v;
    else if (isPlainTensor(v) && v.data.length === 1) extra[k] = v.data[0];
  }
  return {
    nsteps: num(s, "nsteps"),
    ncells: num(s, "ncells"),
    dt: num(s, "dt"),
    fs: num(s, "fs"),
    wall_s: num(s, "wall_s"),
    stable: bool(s, "stable"),
    diverged_step: num(s, "diverged_step"),
    notes: strings(s, "notes"),
    extra,
  };
}

function parseModes(s: PlainStruct): ModeTable {
  return {
    m: vec(s, "m"),
    n: vec(s, "n"),
    f_hz: vec(s, "f_hz"),
    sigma: vec(s, "sigma"),
    G: vec(s, "G"),
    truncated: bool(s, "truncated"),
  };
}

/** Validate and type the solver's result struct. Throws ResultShapeError with a specific message. */
export function parsePlateResult(v: PlainValue): PlateResult {
  if (!isPlainStruct(v)) throw new ResultShapeError("the result variable is not a struct");
  const grid = struct(v, "grid");
  const probe = struct(v, "probe");
  const frames = struct(v, "frames");
  const env = struct(v, "envelope");
  const spec = struct(v, "spectrum");
  const audio = struct(v, "audio");
  const out: PlateResult = {
    method: str(v, "method"),
    grid: { x: vec(grid, "x"), y: vec(grid, "y") },
    t: vec(v, "t"),
    probe: {
      x: num(probe, "x"),
      y: num(probe, "y"),
      w: vec(probe, "w"),
      sxx: vec(probe, "sxx"),
      syy: vec(probe, "syy"),
      txy: vec(probe, "txy"),
      vm: vec(probe, "vm"),
      sp: vec(probe, "sp"),
    },
    frames: { t: vec(frames, "t"), w: frames3d(frames, "w"), vm: frames3d(frames, "vm"), sp: frames3d(frames, "sp") },
    envelope: { vm: field2d(env, "vm"), sp: field2d(env, "sp") },
    spectrum: { f: vec(spec, "f"), w: vec(spec, "w"), vm: vec(spec, "vm") },
    audio: { fs: num(audio, "fs"), p: vec(audio, "p"), t0: num(audio, "t0") },
    info: parseInfo(struct(v, "info")),
  };
  if (isPlainStruct(v.fields.modes)) out.modes = parseModes(v.fields.modes);
  return out;
}

export function parseMethodPlan(v: PlainValue): MethodPlan {
  if (!isPlainStruct(v)) throw new ResultShapeError("the plan variable is not a struct");
  return {
    nx: num(v, "nx"),
    ny: num(v, "ny"),
    dx: num(v, "dx"),
    dy: num(v, "dy"),
    dt_max: v.fields.dt_max !== undefined ? num(v, "dt_max") : undefined,
    dt: num(v, "dt"),
    nsteps: num(v, "nsteps"),
    ncells: num(v, "ncells"),
    cell_steps: num(v, "cell_steps"),
  };
}

export function parseCompareResult(v: PlainValue): CompareResult {
  if (!isPlainStruct(v)) throw new ResultShapeError("the comparison variable is not a struct");
  const fd = v.fields.frames_diff;
  const framesDiff = isPlainTensor(fd) && fd.data.length > 0 ? frames3d(v, "frames_diff") : null;
  return {
    t: vec(v, "t"),
    w_a: vec(v, "w_a"),
    w_b: vec(v, "w_b"),
    vm_a: vec(v, "vm_a"),
    vm_b: vec(v, "vm_b"),
    w_rel_l2: num(v, "w_rel_l2"),
    vm_rel_l2: num(v, "vm_rel_l2"),
    w_peak_a: num(v, "w_peak_a"),
    w_peak_b: num(v, "w_peak_b"),
    vm_peak_a: num(v, "vm_peak_a"),
    vm_peak_b: num(v, "vm_peak_b"),
    f_peak_a: num(v, "f_peak_a"),
    f_peak_b: num(v, "f_peak_b"),
    frames_t: vec(v, "frames_t"),
    frames_rel_l2: vec(v, "frames_rel_l2"),
    frames_diff: framesDiff,
  };
}

/** @deprecated use parseMethodPlan */
export const parseFdPlan = parseMethodPlan;