/**
 * A cymbal as a thin closed body built from quadrilateral patches. A port of
 * fmm3dbie_gpu/cymbal.py from the CUDA proof of concept; lengths in meters.
 *
 * fmm3dbie's integral equations need a closed surface, so the cymbal is not a
 * zero-thickness shell but a thin body around an axisymmetric midsurface
 * z = F(rho), rho = x^2 + y^2, with thickness T(rho). The midsurface has a raised
 * bell at the center and a shallow bow outside it; the thickness decreases from
 * the center to the rim, with zero slope at the rim. The top and bottom faces are
 * the offset surfaces M +- (T/2) N, with M the midsurface and N its upward unit
 * normal, and a half-torus of diameter T closes the body at the rim. Every patch
 * lies entirely on one face or on the rim, so every patch is analytic.
 *
 * The faces are meshed with an O-grid in the (x, y) plane: an equiangular central
 * square of m_c x m_c patches, then rings of patches, graded toward the rim and
 * refined across the bell. Neighboring rings need not have the same number of
 * patches around (the mesh is not conforming). The node derivatives dx/du, dx/dv
 * are computed with forward-mode dual numbers, exactly to rounding. The surface
 * carries no analytic chart for the solver: the integrals use the Legendre
 * interpolant of the node data, as in fmm3dbie.
 */
import { discNodes, Surface } from './surface.ts';

// ------------------------------------------------------------ dual numbers

/** a + b eps with eps^2 = 0: a value and its derivative along one direction. */
class D {
  readonly v: number;
  readonly d: number;
  constructor(v: number, d: number) {
    this.v = v;
    this.d = d;
  }
}
type DN = D | number;
const lift = (a: DN): D => (typeof a === 'number' ? new D(a, 0) : a);
const add = (a: DN, b: DN) => {
  const x = lift(a);
  const y = lift(b);
  return new D(x.v + y.v, x.d + y.d);
};
const sub = (a: DN, b: DN) => {
  const x = lift(a);
  const y = lift(b);
  return new D(x.v - y.v, x.d - y.d);
};
const mul = (a: DN, b: DN) => {
  const x = lift(a);
  const y = lift(b);
  return new D(x.v * y.v, x.d * y.v + x.v * y.d);
};
const div = (a: DN, b: DN) => {
  const x = lift(a);
  const y = lift(b);
  return new D(x.v / y.v, (x.d * y.v - x.v * y.d) / (y.v * y.v));
};
const dsqrt = (a: D) => {
  const s = Math.sqrt(a.v);
  return new D(s, a.d / (2 * s));
};
const dexp = (a: D) => {
  const e = Math.exp(a.v);
  return new D(e, e * a.d);
};
const dcos = (a: D) => new D(Math.cos(a.v), -Math.sin(a.v) * a.d);
const dsin = (a: D) => new D(Math.sin(a.v), Math.cos(a.v) * a.d);
const dtan = (a: D) => {
  const t = Math.tan(a.v);
  return new D(t, (1 + t * t) * a.d);
};

// ------------------------------------------------------------------ shape

export interface CymbalParams {
  /** outer radius, including the rim */
  radius: number;
  /** thickness at the rim */
  tRim: number;
  /** thickness at the center / tRim */
  taper: number;
  bellRadius: number;
  bellHeight: number;
  bellWidth: number;
  /** fall of the midsurface from the center to the rim, apart from the bell */
  bowHeight: number;
}

export const CYMBAL_DEFAULTS: CymbalParams = {
  radius: 0.2,
  tRim: 0.01,
  taper: 2.0,
  bellRadius: 0.06,
  bellHeight: 0.025,
  bellWidth: 0.01,
  bowHeight: 0.015,
};

export class CymbalShape {
  readonly par: CymbalParams;
  readonly rMid: number;

  constructor(par: CymbalParams = CYMBAL_DEFAULTS) {
    this.par = par;
    this.rMid = par.radius - par.tRim / 2;
  }

  private bellArg(rho: DN): D {
    const { bellRadius: rb, bellWidth: w } = this.par;
    return div(sub(rho, rb * rb), 2 * rb * w);
  }

  /** midsurface height as a function of rho */
  F(rho: DN): D {
    const { bowHeight, bellHeight } = this.par;
    const R2 = this.rMid ** 2;
    return add(mul(bowHeight, sub(1, div(rho, R2))), div(bellHeight, add(1, dexp(this.bellArg(rho)))));
  }

  /** dF/drho */
  dF(rho: DN): D {
    const { bowHeight, bellHeight, bellRadius: rb, bellWidth: w } = this.par;
    const e = dexp(this.bellArg(rho));
    const ope = add(1, e);
    return sub(-bowHeight / this.rMid ** 2, div(mul(bellHeight / (2 * rb * w), e), mul(ope, ope)));
  }

  /** thickness as a function of rho */
  T(rho: DN): D {
    const { tRim, taper } = this.par;
    const s = sub(1, div(rho, this.rMid ** 2));
    return mul(tRim, add(1, mul(taper - 1, mul(s, s))));
  }

  /** point of the top (side = 1) or bottom (side = -1) face over (x, y) of the midsurface */
  facePoint(x: DN, y: DN, side: number): [D, D, D] {
    const rho = add(mul(x, x), mul(y, y));
    const g = mul(2, this.dF(rho));
    const q = dsqrt(add(1, mul(mul(g, g), rho)));
    const s = div(mul(side / 2, this.T(rho)), q);
    const sg = mul(s, g);
    return [sub(x, mul(sg, x)), sub(y, mul(sg, y)), add(this.F(rho), s)];
  }

  /** point of the rim at azimuth theta; phi from -pi/2 (bottom face) to pi/2 (top face) */
  rimPoint(theta: DN, phi: DN): [D, D, D] {
    const R = this.rMid;
    const rho = R * R;
    const fp = 2 * R * this.dF(rho).v;
    const q = Math.sqrt(1 + fp * fp);
    const h = this.T(rho).v / 2;
    const a = mul(h / q, dcos(lift(phi)));
    const b = mul(h / q, dsin(lift(phi)));
    const rr = sub(add(R, a), mul(b, fp));
    const th = lift(theta);
    return [mul(rr, dcos(th)), mul(rr, dsin(th)), add(add(this.F(rho).v, mul(a, fp)), b)];
  }

  /** closed outline of the body in the (r, z) half plane, counterclockwise: [r[], z[]] */
  crossSection(n = 400): [Float64Array, Float64Array] {
    const r: number[] = [];
    const z: number[] = [];
    for (let i = 0; i < n; i++) {
      const p = this.facePoint((this.rMid * i) / (n - 1), 0, -1);
      r.push(p[0].v);
      z.push(p[2].v);
    }
    const nr = Math.floor(n / 4);
    for (let i = 0; i < nr; i++) {
      const p = this.rimPoint(0, -Math.PI / 2 + (Math.PI * i) / (nr - 1));
      r.push(p[0].v);
      z.push(p[2].v);
    }
    for (let i = n - 1; i >= 0; i--) {
      const p = this.facePoint((this.rMid * i) / (n - 1), 0, 1);
      r.push(p[0].v);
      z.push(p[2].v);
    }
    return [Float64Array.from(r), Float64Array.from(z)];
  }

  /** volume by Pappus' theorem on the cross section (a check) */
  volume(n = 20000): number {
    const [r, z] = this.crossSection(n);
    let s = 0;
    for (let i = 0; i < r.length; i++) {
      const j = (i + 1) % r.length;
      s += ((r[i] * r[i] + r[i] * r[j] + r[j] * r[j]) / 6) * (z[j] - z[i]);
    }
    return 2 * Math.PI * s;
  }
}

/**
 * A test for points inside the body, or within margin of its surface. Every
 * vertical line meets the cross section in one interval, so the outline, offset
 * outward by margin, splits at its outermost point into a lower and an upper
 * chain, and a point is inside when z lies between them at its radius.
 */
export function insideTest(shape: CymbalShape, margin = 0): (x: ArrayLike<number>) => boolean {
  const [r0, z0] = shape.crossSection(4000);
  const n = r0.length;
  const r = new Float64Array(n);
  const z = new Float64Array(n);
  for (let i = 0; i < n; i++) {
    const a = (i + n - 1) % n;
    const b = (i + 1) % n;
    let tr = r0[b] - r0[a];
    let tz = z0[b] - z0[a];
    const l = Math.hypot(tr, tz);
    tr /= l;
    tz /= l;
    r[i] = Math.max(0, r0[i] + margin * tz);
    z[i] = z0[i] - margin * tr;
  }
  let im = 0;
  for (let i = 1; i < n; i++) if (r[i] > r[im]) im = i;
  const chain = (idx: number[]) => {
    const cr: number[] = [];
    const cz: number[] = [];
    let m = -Infinity;
    for (const i of idx) {
      m = Math.max(m, r[i]);
      cr.push(m);
      cz.push(z[i]);
    }
    return { cr, cz };
  };
  const lo = chain(Array.from({ length: im + 1 }, (_, i) => i));
  const up = chain(Array.from({ length: n - im }, (_, i) => n - 1 - i));
  const interp = (c: { cr: number[]; cz: number[] }, x: number) => {
    const { cr, cz } = c;
    if (x <= cr[0]) return cz[0];
    let a = 0;
    let b = cr.length - 1;
    if (x >= cr[b]) return cz[b];
    while (b - a > 1) {
      const m = (a + b) >> 1;
      if (cr[m] <= x) a = m;
      else b = m;
    }
    const t = cr[b] > cr[a] ? (x - cr[a]) / (cr[b] - cr[a]) : 0;
    return cz[a] + t * (cz[b] - cz[a]);
  };
  const rmax = r[im];
  return (x) => {
    const rr = Math.hypot(x[0], x[1]);
    return rr <= rmax && x[2] > interp(lo, rr) && x[2] < interp(up, rr);
  };
}

// ----------------------------------------------------------------- the mesh

export interface RingLayout {
  mC: number;
  c: number;
  radii: number[];
  counts: number[];
  nrim: number;
}

/**
 * Radii and patch counts of the O-grid for a target patch size h: the central
 * square has half-width c and mC x mC patches (mC >= 2); ring i lies between
 * radii[i] and radii[i + 1] (radii[0] = c means the square) and has counts[i]
 * patches around. The ring around the square reaches r = 2c, so that it is 0.59c
 * wide at the square's corners.
 */
export function ringLayout(shape: CymbalShape, h: number, grade = 1.5, cFrac = 0.12): RingLayout {
  const R = shape.rMid;
  const { tRim, bellRadius, bellWidth } = shape.par;
  const c = cFrac * R;
  const mC = Math.max(2, Math.ceil((2 * c) / h));
  const first = 2 * c;
  const b0 = bellRadius - 2.5 * bellWidth;
  const b1 = bellRadius + 2.5 * bellWidth;
  const wmax = (edge: number) => {
    let w = Math.min(h, tRim + (grade - 1) * (R - edge));
    if (edge > b0 && edge - w < b1) w = Math.min(w, 2 * bellWidth);
    return w;
  };
  const out = [R];
  for (;;) {
    const last = out[out.length - 1];
    const w = wmax(last);
    if (last - w < first + 0.5 * w) break;
    out.push(last - w);
  }
  const radii = [c, first, ...out.reverse()];
  const counts: number[] = [];
  for (let i = 0; i < radii.length - 1; i++) {
    const width = radii[i + 1] - radii[i];
    const arc = Math.min(h, 2.5 * width);
    counts.push(4 * Math.max(mC, Math.ceil((2 * Math.PI * radii[i + 1]) / (4 * arc))));
  }
  counts[0] = 4 * mC;
  for (let i = 1; i < counts.length; i++) counts[i] = Math.max(counts[i], counts[i - 1]);
  return { mC, c, radii, counts, nrim: counts[counts.length - 1] };
}

type Chart = (u: D, v: D) => [D, D];
type PatchMap = (u: D, v: D) => [D, D, D];

/** planar charts (u, v) in [-1, 1]^2 -> (x, y) of the face patches, oriented so that d/du x d/dv points up */
function faceCharts(L: RingLayout): Chart[] {
  const charts: Chart[] = [];
  const q = Math.PI / 4;
  const { mC, c, radii, counts } = L;
  for (let i = 0; i < mC; i++) {
    for (let j = 0; j < mC; j++) {
      const a0 = -1 + (2 * i) / mC;
      const a1 = -1 + (2 * (i + 1)) / mC;
      const b0 = -1 + (2 * j) / mC;
      const b1 = -1 + (2 * (j + 1)) / mC;
      charts.push((u, v) => {
        const s = add(a0, mul(0.5 * (a1 - a0), add(u, 1)));
        const t = add(b0, mul(0.5 * (b1 - b0), add(v, 1)));
        return [mul(c, dtan(mul(q, s))), mul(c, dtan(mul(q, t)))];
      });
    }
  }
  for (let ir = 0; ir < radii.length - 1; ir++) {
    const r0 = radii[ir];
    const r1 = radii[ir + 1];
    const n = counts[ir];
    const dth = (2 * Math.PI) / n;
    for (let k = 0; k < n; k++) {
      const th0 = -q + k * dth;
      if (ir === 0) {
        charts.push((u, v) => {
          const th = add(th0, mul(0.5 * dth, add(v, 1)));
          const sec = Math.floor((th.v + q) / (2 * q));
          const rin = div(c, dcos(sub(th, sec * 2 * q)));
          const a = mul(0.5, add(u, 1));
          const rad = add(mul(sub(1, a), rin), mul(a, r1));
          return [mul(rad, dcos(th)), mul(rad, dsin(th))];
        });
      } else {
        charts.push((u, v) => {
          const th = add(th0, mul(0.5 * dth, add(v, 1)));
          const rad = add(r0, mul(0.5 * (r1 - r0), add(u, 1)));
          return [mul(rad, dcos(th)), mul(rad, dsin(th))];
        });
      }
    }
  }
  return charts;
}

/** all patch maps (u, v) -> xyz with outward orientation: top faces, bottom faces, rim */
function patchMaps(shape: CymbalShape, h: number, grade: number): { maps: PatchMap[]; layout: RingLayout } {
  const layout = ringLayout(shape, h, grade);
  const charts = faceCharts(layout);
  const maps: PatchMap[] = [];
  for (const ch of charts) maps.push((u, v) => shape.facePoint(...ch(u, v), 1));
  for (const ch of charts) maps.push((u, v) => shape.facePoint(...ch(v, u), -1)); // swapping u and v flips the normal
  const { nrim } = layout;
  for (let k = 0; k < nrim; k++) {
    const th0 = -Math.PI / 4 + (2 * Math.PI * k) / nrim;
    maps.push((u, v) => shape.rimPoint(add(th0, mul(Math.PI / nrim, add(u, 1))), mul(0.5 * Math.PI, v)));
  }
  return { maps, layout };
}

/** Node data (12, npts) in the fmm3dbie iptype 11 convention, and the layout. */
export function cymbalSrcvals(shape: CymbalShape, h: number, norder: number, grade = 1.5): { srcvals: Float64Array; layout: RingLayout } {
  const { uv } = discNodes(norder);
  const npols = (norder + 1) ** 2;
  const { maps, layout } = patchMaps(shape, h, grade);
  const npts = maps.length * npols;
  const sv = new Float64Array(12 * npts);
  for (const [j, f] of maps.entries()) {
    for (let k = 0; k < npols; k++) {
      const i = j * npols + k;
      const pu = f(new D(uv[k], 1), new D(uv[npols + k], 0));
      const pv = f(new D(uv[k], 0), new D(uv[npols + k], 1));
      const ru = [pu[0].d, pu[1].d, pu[2].d];
      const rv = [pv[0].d, pv[1].d, pv[2].d];
      const nx = ru[1] * rv[2] - ru[2] * rv[1];
      const ny = ru[2] * rv[0] - ru[0] * rv[2];
      const nz = ru[0] * rv[1] - ru[1] * rv[0];
      const l = Math.hypot(nx, ny, nz);
      const vals = [pu[0].v, pu[1].v, pu[2].v, ...ru, ...rv, nx / l, ny / l, nz / l];
      for (let c = 0; c < 12; c++) sv[c * npts + i] = vals[c];
    }
  }
  return { srcvals: sv, layout };
}

/** The cymbal as a Surface with the interpolated geometry, for a target patch size h (m). */
export function cymbal(shape: CymbalShape, h: number, norder: number, grade = 1.5): Surface & { layout: RingLayout } {
  const { srcvals, layout } = cymbalSrcvals(shape, h, norder, grade);
  const s = new Surface(norder, srcvals, null) as Surface & { layout: RingLayout };
  s.layout = layout;
  return s;
}