import {
  isPlainCell,
  isPlainStruct,
  isPlainTensor,
  type PlainStruct,
  type PlainValue,
} from "../runner/plain";
import type { Field2D, FieldFrames } from "../fields/field";
import type { MethodId } from "../spec/types";
import type { ContactResult, MethodPlan, ModesResult, 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 numOpt(s: PlainStruct, name: string, fallback: number): number {
  return s.fields[name] === undefined ? fallback : num(s, name);
}

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 [];
}

function methodId(s: string): MethodId {
  if (s === "rr" || s === "fem") return s;
  throw new ResultShapeError(`unknown method "${s}"`);
}

const INFO_STANDARD = new Set([
  "method",
  "ndof",
  "nmodes",
  "n_audible",
  "n_rigid",
  "fs",
  "nsamples",
  "nframes",
  "frame_span",
  "plate_mass",
  "thread_k",
  "critical_frequency",
  "f_max",
  "size_note",
  "modes_wall_s",
  "wall_s",
  "strike_wall_s",
  "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 {
    method: methodId(str(s, "method")),
    ndof: num(s, "ndof"),
    nmodes: num(s, "nmodes"),
    n_audible: num(s, "n_audible"),
    n_rigid: num(s, "n_rigid"),
    fs: num(s, "fs"),
    nsamples: num(s, "nsamples"),
    nframes: num(s, "nframes"),
    frame_span: num(s, "frame_span"),
    plate_mass: num(s, "plate_mass"),
    thread_k: num(s, "thread_k"),
    critical_frequency: num(s, "critical_frequency"),
    f_max: num(s, "f_max"),
    size_note: str(s, "size_note"),
    modes_wall_s: num(s, "modes_wall_s"),
    wall_s: num(s, "wall_s"),
    strike_wall_s: numOpt(s, "strike_wall_s", 0),
    notes: strings(s, "notes"),
    extra,
  };
}

function parseModes(s: PlainStruct): ModesResult {
  return {
    n: num(s, "n"),
    freq: vec(s, "freq"),
    q: vec(s, "q"),
    zeta: vec(s, "zeta"),
    hit: vec(s, "hit"),
    support: vec(s, "support"),
    radiation: vec(s, "radiation"),
    radiation_factor: vec(s, "radiation_factor"),
    acoustic: vec(s, "acoustic"),
    audible: Array.from(vec(s, "audible"), (x) => x !== 0),
    q0: vec(s, "q0"),
    v0: vec(s, "v0"),
    shapes: frames3d(s, "shapes"),
  };
}

function parseContact(s: PlainStruct): ContactResult {
  return {
    model: str(s, "model"),
    tc: num(s, "tc"),
    impulse: num(s, "impulse"),
    peak_force: num(s, "peak_force"),
    t: vec(s, "t"),
    force: vec(s, "force"),
    hammer_x: vec(s, "hammer_x"),
    hammer_v: vec(s, "hammer_v"),
    plate_w: vec(s, "plate_w"),
    plate_v: vec(s, "plate_v"),
    hammer_exit_v: num(s, "hammer_exit_v"),
    plate_exit_v: num(s, "plate_exit_v"),
    relative_separation_v: num(s, "relative_separation_v"),
    effective_restitution: num(s, "effective_restitution"),
    hertz_duration: num(s, "hertz_duration"),
    hertz_impulse: num(s, "hertz_impulse"),
    hertz_peak_force: num(s, "hertz_peak_force"),
    contact_k: num(s, "contact_k"),
  };
}

/** Validate and type the strike'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");
  return {
    method: methodId(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"),
      w_total: frames3d(frames, "w_total"),
      vm: frames3d(frames, "vm"),
      sp: frames3d(frames, "sp"),
    },
    envelope: { vm: field2d(env, "vm"), sp: field2d(env, "sp") },
    spectrum: { f: vec(spec, "f"), p: vec(spec, "p"), w: vec(spec, "w") },
    audio: {
      fs: num(audio, "fs"),
      p: vec(audio, "p"),
      peak_pa: num(audio, "peak_pa"),
      gain: num(audio, "gain"),
      normalize_dbfs: num(audio, "normalize_dbfs"),
      model: str(audio, "model"),
    },
    modes: parseModes(struct(v, "modes")),
    contact: parseContact(struct(v, "contact")),
    info: parseInfo(struct(v, "info")),
  };
}

export function parseMethodPlan(v: PlainValue): MethodPlan {
  if (!isPlainStruct(v)) throw new ResultShapeError("the plan variable is not a struct");
  const opt = (name: string) => (v.fields[name] === undefined ? undefined : num(v, name));
  return {
    method: methodId(str(v, "method")),
    ndof: num(v, "ndof"),
    work: num(v, "work"),
    size_note: str(v, "size_note"),
    plate_mass: num(v, "plate_mass"),
    thread_k: num(v, "thread_k"),
    pendulum_hz: num(v, "pendulum_hz"),
    f_max: num(v, "f_max"),
    modes_expected: num(v, "modes_expected"),
    modes_cap: num(v, "modes_cap"),
    hertz_duration: num(v, "hertz_duration"),
    hertz_impulse: num(v, "hertz_impulse"),
    hertz_peak_force: num(v, "hertz_peak_force"),
    contact_k: num(v, "contact_k"),
    critical_frequency: num(v, "critical_frequency"),
    nsamples: num(v, "nsamples"),
    nframes: num(v, "nframes"),
    notes: strings(v, "notes"),
    nx: opt("nx"),
    ny: opt("ny"),
    dx: opt("dx"),
    dy: opt("dy"),
  };
}

/** Sanity check of a modes struct read back from stage 1, before it is kept. */
export function checkModesValue(v: PlainValue): { method: MethodId; nmodes: number } {
  if (!isPlainStruct(v)) throw new ResultShapeError("the modes variable is not a struct");
  const method = methodId(str(v, "method"));
  const freqs = vec(v, "freqs");
  return { method, nmodes: freqs.length };
}