import {
  isPlainCell,
  isPlainStruct,
  isPlainTensor,
  type PlainStruct,
  type PlainValue,
} from "../runner/plain";
import type { Field2D, FieldFrames } from "../fields/field";
import type {
  MethodPlan,
  PlateResult,
  ResultInfo,
  RigidResult,
  StaticResult,
  SupportsResult,
} 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`);
}

/** A numeric vector that may be absent or empty. */
function vecOpt(s: PlainStruct, name: string): Float64Array {
  const v = s.fields[name];
  if (v === undefined) return new Float64Array(0);
  if (typeof v === "number") return Float64Array.of(v);
  if (isPlainTensor(v)) return v.data;
  return new Float64Array(0);
}

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 field2dOpt(s: PlainStruct, name: string): Field2D | undefined {
  const v = s.fields[name];
  if (!isPlainTensor(v) || v.shape.length !== 2) return undefined;
  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 };
}

/** Rows of an m x n column-major matrix as m separate arrays. */
function rows(s: PlainStruct, name: string, m: number): Float64Array[] {
  const v = s.fields[name];
  if (m === 0 || v === undefined || !isPlainTensor(v)) return [];
  const [mm, n] = [v.shape[0] ?? 0, v.shape[1] ?? 0];
  if (mm !== m) throw new ResultShapeError(`field "${name}" should have ${m} rows`);
  const out: Float64Array[] = [];
  for (let k = 0; k < m; k++) {
    const r = new Float64Array(n);
    for (let j = 0; j < n; j++) r[j] = v.data[k + j * m];
    out.push(r);
  }
  return out;
}

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 parseRigid(s: PlainStruct, framesStruct: PlainStruct): RigidResult {
  const nr = num(s, "nr");
  const ed = vec(s, "ed");
  const eh = vec(s, "eh");
  return {
    nr,
    names: strings(s, "names"),
    series: rows(s, "series", nr),
    frames: rows(framesStruct, "rigid", nr),
    pivot: { x: num(s, "pivot_x"), y: num(s, "pivot_y") },
    ed: [ed[0] ?? 1, ed[1] ?? 0],
    eh: [eh[0] ?? 0, eh[1] ?? 1],
    hz: Array.from(vecOpt(s, "hz")),
    modeNames: strings(s, "mode_names"),
  };
}

function parseSupports(s: PlainStruct): SupportsResult {
  const n = num(s, "n");
  return {
    names: strings(s, "names"),
    units: strings(s, "units"),
    kinds: strings(s, "kinds"),
    series: rows(s, "series", n),
    static: Array.from(vecOpt(s, "static")).slice(0, n),
  };
}

function parseStatic(s: PlainStruct): StaticResult {
  return {
    w0: field2d(s, "w0"),
    sag: num(s, "sag"),
    gx: num(s, "gx"),
    gy: num(s, "gy"),
    gz: num(s, "gz"),
    g: num(s, "g"),
    mass: num(s, "mass"),
    xc: num(s, "xc"),
    yc: num(s, "yc"),
    held: field2d(s, "held"),
    Nxx: field2dOpt(s, "Nxx"),
    Nyy: field2dOpt(s, "Nyy"),
    Nxy: field2dOpt(s, "Nxy"),
  };
}

/** 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");
  return {
    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"),
      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"), w: vec(spec, "w"), vm: vec(spec, "vm") },
    audio: {
      fs: num(audio, "fs"),
      p: vec(audio, "p"),
      t0: num(audio, "t0"),
      model: str(audio, "model"),
    },
    rigid: parseRigid(struct(v, "rigid"), frames),
    supports: parseSupports(struct(v, "supports")),
    static: parseStatic(struct(v, "static")),
    info: parseInfo(struct(v, "info")),
  };
}

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: num(v, "dt_max"),
    dt: num(v, "dt"),
    nsteps: num(v, "nsteps"),
    ncells: num(v, "ncells"),
    nunknowns: num(v, "nunknowns"),
    cell_steps: num(v, "cell_steps"),
    rigid_hz: Array.from(vecOpt(v, "rigid_hz")),
    rigid_names: strings(v, "rigid_names"),
    rigid_mode_names: strings(v, "rigid_mode_names"),
    sag: num(v, "sag"),
    gx: num(v, "gx"),
    gy: num(v, "gy"),
    gz: num(v, "gz"),
    mass: num(v, "mass"),
    xc: num(v, "xc"),
    yc: num(v, "yc"),
    tension: Array.from(vecOpt(v, "tension")),
    notes: strings(v, "notes"),
  };
}