import type { Anchor, BeadSystem, Excitation } from '../../sim/types';
import { dampingFromT60 } from '../bead-string/model';

// A round drumhead: beads on a triangular (hex) lattice in the x-z plane,
// cut to a circle, held by a rim of fixed lattice points just outside it.
// Transverse vibration is along y, as on the string.

export type DrumExcitation = 'mallet' | 'strike' | 'force' | 'press';

export interface DrumParams {
  // geometry
  diameter: number; // m, rim to rim
  rings: number; // beads from centre to edge (lattice spacing a = R / (rings + 1))
  // membrane properties (continuum values; the lattice is chosen to match them)
  density: number; // kg/m^2, surface density
  tension: number; // N/m, isotropic membrane tension
  inPlane: number; // N/m, in-plane stiffness E t
  // damping
  t60Low: number;
  fLow: number;
  t60High: number;
  fHigh: number;
  // excitation
  excitation: DrumExcitation;
  excR: number; // fraction of radius, 0 = centre
  excAngle: number; // degrees, in the x-z plane from +x
  excWidth: number; // m, Gaussian sigma of the contact patch (0 = one bead)
  malletMass: number; // kg
  malletSpeed: number; // m/s
  feltK: number; // N/m^p
  feltP: number;
  strikeVel: number; // m/s
  forcePeak: number; // N
  forceMs: number; // ms
  pressAmp: number; // m
  // pickup
  pickupMode: 'bead' | 'sum';
  pickupR: number;
  pickupAngle: number;
  pickupAxis: 'y' | 'x' | 'z';
  // simulation
  duration: number;
  fs: number;
  cfl: number;
  frameRate: number;
}

export interface Lattice {
  a: number; // spacing (m)
  pos: [number, number][]; // bead positions (x, z)
  rim: [number, number][]; // anchor positions (x, z)
  edges: [number, number][]; // node pairs; node >= pos.length is rim[node - n]
}

const key = (i: number, j: number) => `${i},${j}`;

export function hexLattice(R: number, rings: number): Lattice {
  const a = R / (rings + 1);
  const at = (i: number, j: number): [number, number] => [a * (i + 0.5 * j), a * (Math.sqrt(3) / 2) * j];
  const inside = (i: number, j: number) => Math.hypot(...at(i, j)) <= R - 0.5 * a + 1e-9;
  const index = new Map<string, number>();
  const pos: [number, number][] = [];
  const M = rings + 3;
  for (let j = -M; j <= M; j++)
    for (let i = -M - Math.abs(j); i <= M + Math.abs(j); i++)
      if (inside(i, j)) {
        index.set(key(i, j), pos.length);
        pos.push(at(i, j));
      }
  const n = pos.length;
  const rimIndex = new Map<string, number>();
  const rim: [number, number][] = [];
  const edges: [number, number][] = [];
  const nb = [[1, 0], [0, 1], [-1, 1]];
  const all = [[1, 0], [0, 1], [-1, 1], [-1, 0], [0, -1], [1, -1]];
  for (const [k, b] of index) {
    const [i, j] = k.split(',').map(Number);
    // bead-bead springs: each pair once
    for (const [di, dj] of nb) {
      const o = index.get(key(i + di, j + dj));
      if (o !== undefined) edges.push([b, o]);
    }
    // bead-rim springs: every missing neighbour is a rim point
    for (const [di, dj] of all) {
      const kk = key(i + di, j + dj);
      if (index.has(kk)) continue;
      let r = rimIndex.get(kk);
      if (r === undefined) {
        r = rim.length;
        rimIndex.set(kk, r);
        rim.push(at(i + di, j + dj));
      }
      edges.push([b, n + r]);
    }
  }
  return { a, pos, rim, edges };
}

export interface DrumDerived {
  lattice: Lattice;
  n: number;
  beadMass: number;
  springK: number;
  springTension: number;
  restLength: number;
  prestretch: number;
  waveSpeedT: number;
  waveSpeedP: number; // in-plane longitudinal (P) wave speed
  idealF01: number; // ideal clamped membrane fundamental
  sigma0: number;
  sigma1: number;
  dashpot: number;
  warnings: string[];
}

// Triangular spring lattice, spacing a, spring stiffness k, spring tension f:
//   membrane tension tau = sqrt(3) f / a, 2D Young's modulus E t = 2 k / sqrt(3)
//   (Poisson ratio 1/3), area per bead sqrt(3)/2 a^2.
export function deriveDrum(p: DrumParams): DrumDerived {
  const R = p.diameter / 2;
  const lattice = hexLattice(R, Math.max(1, Math.round(p.rings)));
  const a = lattice.a;
  const beadMass = p.density * (Math.sqrt(3) / 2) * a * a;
  const springK = (Math.sqrt(3) / 2) * p.inPlane;
  const springTension = (p.tension * a) / Math.sqrt(3);
  const restLength = a - springTension / springK;
  const warnings: string[] = [];
  if (!(restLength > 0)) warnings.push('Tension too high for this in-plane stiffness: the springs would need a negative rest length.');
  const { sigma0, sigma1, note } = dampingFromT60(p);
  if (note) warnings.push(note);
  const waveSpeedT = Math.sqrt(Math.max(p.tension, 0) / p.density);
  const lambdaPlus2Mu = (3 * Math.sqrt(3) * springK) / 4;
  return {
    lattice,
    n: lattice.pos.length,
    beadMass,
    springK,
    springTension,
    restLength,
    prestretch: a / restLength - 1,
    waveSpeedT,
    waveSpeedP: Math.sqrt(lambdaPlus2Mu / p.density),
    idealF01: (2.404826 * waveSpeedT) / (2 * Math.PI * R),
    sigma0,
    sigma1,
    // each spring's transverse stiffness is f / a, so this damps transverse modes by sigma1 omega^2
    dashpot: 2 * sigma1 * (springTension > 0 ? springTension / a : springK),
    warnings,
  };
}

export function buildDrum(d: DrumDerived): BeadSystem {
  const { pos, rim, edges } = d.lattice;
  const n = pos.length;
  const rest = new Float64Array(3 * n);
  pos.forEach(([x, z], i) => {
    rest[3 * i] = x;
    rest[3 * i + 2] = z;
  });
  const anchors: Anchor[] = rim.map(([x, z]) => ({ pos: [x, 0, z], slideBead: -1 }));
  const ns = edges.length;
  return {
    n,
    rest,
    mass: new Float64Array(n).fill(d.beadMass),
    pinned: new Uint8Array(n),
    anchors,
    si: Int32Array.from(edges.map((e) => e[0])),
    sj: Int32Array.from(edges.map((e) => e[1])),
    sk: new Float64Array(ns).fill(d.springK),
    sL0: new Float64Array(ns).fill(d.restLength),
    sc: new Float64Array(ns).fill(d.dashpot),
    sigma0: d.sigma0,
  };
}

const polar = (R: number, frac: number, deg: number): [number, number] => {
  const th = (deg * Math.PI) / 180;
  return [R * frac * Math.cos(th), R * frac * Math.sin(th)];
};

export function nearestBead(d: DrumDerived, x: number, z: number): number {
  let best = 0;
  let bd = Infinity;
  d.lattice.pos.forEach(([px, pz], i) => {
    const q = (px - x) ** 2 + (pz - z) ** 2;
    if (q < bd) {
      bd = q;
      best = i;
    }
  });
  return best;
}

// Gaussian patch around the excitation point; width 0 = the nearest bead
function patch(p: DrumParams, d: DrumDerived): Float64Array {
  const [cx, cz] = polar(p.diameter / 2, p.excR, p.excAngle);
  const n = d.n;
  const w = new Float64Array(n);
  if (!(p.excWidth > 0)) {
    w[nearestBead(d, cx, cz)] = 1;
    return w;
  }
  d.lattice.pos.forEach(([x, z], i) => {
    const r2 = ((x - cx) ** 2 + (z - cz) ** 2) / (p.excWidth * p.excWidth);
    w[i] = r2 < 9 ? Math.exp(-0.5 * r2) : 0;
  });
  if (!w.some((v) => v > 0)) w[nearestBead(d, cx, cz)] = 1;
  return w;
}

export function excitationPoint(p: DrumParams): [number, number] {
  return polar(p.diameter / 2, p.excR, p.excAngle);
}

export function buildDrumExcitation(p: DrumParams, d: DrumDerived): Excitation {
  const n = d.n;
  const u0 = new Float64Array(3 * n);
  const v0 = new Float64Array(3 * n);
  const w = patch(p, d);
  let max = 0;
  let sum = 0;
  for (let i = 0; i < n; i++) {
    max = Math.max(max, w[i]);
    sum += w[i];
  }
  // everything acts downward (-y), onto the head from above
  if (p.excitation === 'press') {
    for (let i = 0; i < n; i++) u0[3 * i + 1] = (-p.pressAmp * w[i]) / max;
    return { u0, v0 };
  }
  if (p.excitation === 'strike') {
    for (let i = 0; i < n; i++) v0[3 * i + 1] = (-p.strikeVel * w[i]) / max;
    return { u0, v0 };
  }
  const weights = w.map((v) => v / sum);
  if (p.excitation === 'force') {
    const shape = new Float64Array(3 * n);
    for (let i = 0; i < n; i++) shape[3 * i + 1] = -weights[i];
    return { u0, v0, force: { shape, F0: p.forcePeak, tau: p.forceMs / 1000 } };
  }
  return {
    u0,
    v0,
    mallet: { weights, normal: [0, -1, 0], vel0: [0, -p.malletSpeed, 0], mass: p.malletMass, K: p.feltK, p: p.feltP, mu: 0 },
  };
}

export function buildDrumPickup(p: DrumParams, d: DrumDerived): Float64Array {
  const w = new Float64Array(3 * d.n);
  const axis = p.pickupAxis === 'x' ? 0 : p.pickupAxis === 'y' ? 1 : 2;
  if (p.pickupMode === 'sum') for (let i = 0; i < d.n; i++) w[3 * i + axis] = 1;
  else w[3 * nearestBead(d, ...polar(p.diameter / 2, p.pickupR, p.pickupAngle)) + axis] = 1;
  return w;
}