import type { Vec2, PolyMap, PolyWarp } from '../types';
import { leastSquares } from './linalg';

/** Number of monomial terms of total degree ≤ d in two variables. */
function termCount(degree: number): number {
  return ((degree + 1) * (degree + 2)) / 2;
}

/** Monomials x^i·y^j with i+j ≤ degree, in a fixed order. */
function monomials(degree: number, x: number, y: number): number[] {
  const terms: number[] = [];
  for (let i = 0; i <= degree; i++) {
    for (let j = 0; j <= degree - i; j++) terms.push(x ** i * y ** j);
  }
  return terms;
}

/** Center+scale so normalized inputs sit roughly in [-1,1] (conditions the fit). */
function computeNorm(pts: Vec2[]): { off: Vec2; scl: Vec2 } {
  const xs = pts.map((p) => p[0]);
  const ys = pts.map((p) => p[1]);
  const off: Vec2 = [xs.reduce((a, b) => a + b, 0) / xs.length, ys.reduce((a, b) => a + b, 0) / ys.length];
  const sclx = Math.max(...xs.map((x) => Math.abs(x - off[0]))) || 1;
  const scly = Math.max(...ys.map((y) => Math.abs(y - off[1]))) || 1;
  return { off, scl: [sclx, scly] };
}

function applyNorm(norm: { off: Vec2; scl: Vec2 }, p: Vec2): Vec2 {
  return [(p[0] - norm.off[0]) / norm.scl[0], (p[1] - norm.off[1]) / norm.scl[1]];
}

/** Fit one polynomial map src → dst of the given degree. */
function fitPolyMap(src: Vec2[], dst: Vec2[], degree: number): PolyMap {
  if (src.length !== dst.length) throw new Error('fitPolyMap: src/dst length mismatch');
  const need = termCount(degree);
  if (src.length < need) throw new Error(`Need at least ${need} points for degree ${degree}`);
  const norm = computeNorm(src);
  const rows = src.map((p) => {
    const [x, y] = applyNorm(norm, p);
    return monomials(degree, x, y);
  });
  const cu = leastSquares(rows, dst.map((d) => d[0]));
  const cv = leastSquares(rows, dst.map((d) => d[1]));
  return { degree, norm, cu, cv };
}

/** Evaluate a fitted polynomial map at a point. */
export function applyPolyMap(m: PolyMap, p: Vec2): Vec2 {
  const [x, y] = applyNorm(m.norm, p);
  const mon = monomials(m.degree, x, y);
  let u = 0;
  let v = 0;
  for (let i = 0; i < mon.length; i++) {
    u += m.cu[i]! * mon[i]!;
    v += m.cv[i]! * mon[i]!;
  }
  return [u, v];
}

/**
 * Fit a bidirectional warp from matching machine↔image point pairs.
 * Defaults to degree 3: barrel distortion is a cubic effect on position
 * (the radial term r² multiplies the coordinate), so degree 3 is what actually
 * compensates a wide-angle lens. Degree 2 is a fallback for sparse calibration.
 */
export function estimatePolyWarp(machine: Vec2[], image: Vec2[], degree = 3): PolyWarp {
  return {
    degree,
    fwd: fitPolyMap(machine, image, degree),
    inv: fitPolyMap(image, machine, degree),
  };
}

/** machine-mm → normalized image. */
export function applyPolyWarp(w: PolyWarp, machine: Vec2): Vec2 {
  return applyPolyMap(w.fwd, machine);
}

/** normalized image → machine-mm. */
export function applyPolyWarpInverse(w: PolyWarp, image: Vec2): Vec2 {
  return applyPolyMap(w.inv, image);
}