import * as THREE from 'three';
import { chargePair, carrierModel, number, plateModel, wireModel, type ExperimentId, type Parameters } from '../experiments';
import { electricField, type Charge, type Point } from '../physics/electricField';
import { forceOn, pairEnergy, potential } from '../physics/models';
import { arrow, box, colours, grid, label, line, sphere } from './primitives';
import { buildTestScene } from './test4';

export type Animate = (delta: number) => void;
const still: Animate = () => {};
function signedCharge(group: THREE.Group, charge: Charge) {
  const colour = charge.coulombs >= 0 ? colours.positive : colours.negative;
  sphere(group, charge, 0.2, colour);
  label(group, String.raw`$${charge.coulombs >= 0 ? '+' : '-'}${Math.abs(charge.coulombs * 1e9).toFixed(1)}\,\mathrm{nC}$`, { ...charge, y: charge.y + 0.55 }, colour, 1.25);
}
function compressedFieldArrow(group: THREE.Group, position: Point, direction: Point, colour = colours.field) {
  const strength = Math.hypot(direction.x, direction.y, direction.z);
  if (strength < 1e-8) return;
  arrow(group, position, direction, 0.12 + 0.32 * strength / (strength + 4), colour);
}

function fieldScene(group: THREE.Group, p: Parameters) {
  grid(group);
  const charges = chargePair(p);
  charges.forEach(charge => signedCharge(group, charge));
  for (let x = -5; x <= 5; x += 0.6) for (let z = -5; z <= 5; z += 0.6) {
    const point = { x, y: 0, z };
    if (charges.some(q => Math.hypot(q.x - x, q.z - z) < 0.42)) continue;
    compressedFieldArrow(group, point, electricField(point, charges));
  }
  const force = forceOn(charges[1], [charges[0]]);
  arrow(group, { ...charges[1], y: 0.8 }, force, 1.3, colours.gold);
  label(group, String.raw`$\vec F$ on right charge`, { x: charges[1].x, y: 1.5, z: 0 }, colours.gold, 2);
  sphere(group, { x: 0, y: 0, z: 0 }, 0.06, colours.gold);
  label(group, 'Origin', { x: 0, y: 0, z: 0.45 }, colours.muted, 0.9);
}

function fluxScene(group: THREE.Group, p: Parameters) {
  const radius = number(p, 'radius');
  const charge: Charge = { x: number(p, 'position'), y: 0, z: 0, coulombs: number(p, 'charge') * 1e-9 };
  grid(group, -3.4);
  signedCharge(group, charge);
  const surface = new THREE.Mesh(new THREE.SphereGeometry(radius, 40, 24), new THREE.MeshBasicMaterial({ color: colours.negative, wireframe: true, transparent: true, opacity: 0.11, depthWrite: false }));
  group.add(surface);
  for (let i = 0; i < 150; i++) {
    const y = 1 - 2 * (i + 0.5) / 150;
    const r = Math.sqrt(1 - y * y);
    const theta = i * Math.PI * (3 - Math.sqrt(5));
    const normal = { x: r * Math.cos(theta), y, z: r * Math.sin(theta) };
    const point = { x: normal.x * radius, y: y * radius, z: normal.z * radius };
    if (Math.hypot(point.x - charge.x, point.y, point.z) < 0.3) continue;
    const e = electricField(point, [charge]);
    const dot = e.x * normal.x + e.y * normal.y + e.z * normal.z;
    compressedFieldArrow(group, point, e, dot >= 0 ? colours.field : colours.positive);
  }
  label(group, 'Gaussian surface', { x: 0, y: radius + 0.6, z: 0 }, colours.negative, 2.5);
  line(group, [{ x: -4.7, y: 0, z: 0 }, { x: 4.7, y: 0, z: 0 }], colours.muted, 0.4);
}

function energyScene(group: THREE.Group, p: Parameters) {
  grid(group, -4.3);
  const r = number(p, 'distance');
  const q = number(p, 'testCharge') * 1e-9;
  const source: Charge = { x: -3, y: 0, z: 1.5, coulombs: 3e-9 };
  const moving: Charge = { x: -3 + r, y: 0, z: 1.5, coulombs: q };
  signedCharge(group, source);
  signedCharge(group, moving);
  line(group, [source, moving], colours.muted);
  arrow(group, { ...moving, y: 0.7 }, forceOn(moving, [source]), 1.1, colours.gold);
  const curve: Point[] = [];
  for (let distance = 0.75; distance <= 6; distance += 0.05) curve.push({ x: -3 + distance, y: pairEnergy(3e-9, q, distance) * 1e9 / 20, z: -1.5 });
  line(group, curve, colours.field);
  line(group, [{ x: -2.5, y: 0, z: -1.5 }, { x: 3.5, y: 0, z: -1.5 }], colours.muted);
  line(group, [{ x: -3, y: -5.5, z: -1.5 }, { x: -3, y: 5.5, z: -1.5 }], colours.muted, 0.5);
  sphere(group, { x: -3 + r, y: pairEnergy(3e-9, q, r) * 1e9 / 20, z: -1.5 }, 0.12, colours.gold);
  label(group, String.raw`$U(r)$ · 20 nJ per height unit`, { x: 0, y: 5.7, z: -1.5 }, colours.field, 3.9);
  label(group, '0 energy', { x: 4, y: 0, z: -1.5 }, colours.muted, 1.5);
  label(group, String.raw`$r = ${r.toFixed(2)}\,\mathrm m$`, { x: 0, y: -0.65, z: 1.5 }, colours.muted, 2);
}

function potentialScene(group: THREE.Group, p: Parameters) {
  const charges = chargePair(p);
  const geometry = new THREE.PlaneGeometry(10, 10, 64, 64);
  geometry.rotateX(-Math.PI / 2);
  const positions = geometry.getAttribute('position');
  const values: number[] = [];
  const colourValues: number[] = [];
  const negative = new THREE.Color(colours.negative);
  const positive = new THREE.Color(colours.positive);
  const neutral = new THREE.Color('#203847');
  for (let i = 0; i < positions.count; i++) {
    const point = { x: positions.getX(i), y: 0, z: positions.getZ(i) };
    const v = potential(point, charges);
    const height = Number.isFinite(v) ? Math.tanh(v / 15) : NaN;
    values.push(height);
    positions.setY(i, Number.isFinite(height) ? 2 * height : 0);
    const colour = neutral.clone().lerp(height >= 0 ? positive : negative, Math.abs(Number.isFinite(height) ? height : 0));
    colourValues.push(colour.r, colour.g, colour.b);
  }
  const indices = geometry.getIndex()!;
  const kept: number[] = [];
  for (let i = 0; i < indices.count; i += 3) {
    const a = indices.getX(i), b = indices.getX(i + 1), c = indices.getX(i + 2);
    if ([a, b, c].every(index => Number.isFinite(values[index]))) kept.push(a, b, c);
  }
  geometry.setIndex(kept);
  geometry.setAttribute('color', new THREE.Float32BufferAttribute(colourValues, 3));
  geometry.computeVertexNormals();
  group.add(new THREE.Mesh(geometry, new THREE.MeshStandardMaterial({ vertexColors: true, side: THREE.DoubleSide, roughness: 0.7, transparent: true, opacity: 0.88 })));
  for (let x = -4.5; x <= 4.5; x += 0.75) for (let z = -4.5; z <= 4.5; z += 0.75) {
    const point = { x, y: 0, z };
    if (charges.some(q => Math.hypot(q.x - x, q.z - z) < 0.6)) continue;
    const v = potential(point, charges);
    if (!Number.isFinite(v)) continue;
    compressedFieldArrow(group, { ...point, y: 2 * Math.tanh(v / 15) + 0.12 }, electricField(point, charges));
  }
  const probe = { x: number(p, 'probe'), y: 0, z: 1 };
  const v = potential(probe, charges);
  sphere(group, { ...probe, y: 2 * Math.tanh(v / 15) + 0.2 }, 0.13, colours.gold);
  label(group, 'Probe', { ...probe, y: 2 * Math.tanh(v / 15) + 0.8 }, colours.gold, 1.3);
  charges.forEach(charge => signedCharge(group, { ...charge, y: Math.sign(charge.coulombs) * 2.2 }));
  grid(group, -2.4);
}

function plateScene(group: THREE.Group, p: Parameters, dielectric: boolean) {
  const c = plateModel(p);
  const halfGap = number(p, 'gap') * 0.75;
  const side = Math.sqrt(number(p, 'area')) * 0.3;
  // Uniform magnification maps lateral metres to 30 scene units/m and gap
  // metres to 1500 units/m: gap is enlarged 50x relative to plate width.
  for (const sign of [1, -1]) {
    const x = -sign * halfGap;
    const colour = sign === 1 ? colours.positive : colours.negative;
    box(group, { x, y: 0, z: 0 }, { x: 0.05, y: side, z: side }, colour, 0.65);
    label(group, sign === 1 ? '+ plate' : '− plate', { x, y: side / 2 + 0.45, z: 0 }, colour, 1.2);
    for (let y = -0.35; y <= 0.35; y += 0.35) for (let z = -0.35; z <= 0.35; z += 0.35) {
      sphere(group, { x: x + sign * 0.035, y: y * side, z: z * side }, 0.025, colour);
    }
  }
  for (let y = -0.3; y <= 0.3; y += 0.3) for (let z = -0.3; z <= 0.3; z += 0.3) {
    const length = Math.max(0.1, 2 * halfGap - 0.1);
    // Repeated arrows are direction cues; readouts provide field magnitude.
    arrow(group, { x: -halfGap + 0.05, y: y * side, z: z * side }, { x: 1, y: 0, z: 0 }, length, colours.field);
  }
  if (dielectric) {
    box(group, { x: 0, y: 0, z: 0 }, { x: Math.max(0.04, 2 * halfGap - 0.08), y: side * 0.92, z: side * 0.92 }, colours.gold, 0.1);
    const kappa = number(p, 'permittivity');
    if (kappa > 1) {
      const boundOffset = Math.min(halfGap * 0.65, 0.12) * (kappa - 1) / kappa;
      for (let y = -0.32; y <= 0.32; y += 0.32) for (let z = -0.32; z <= 0.32; z += 0.32) {
        const center = { x: 0, y: y * side, z: z * side };
        sphere(group, { ...center, x: -boundOffset }, 0.035, colours.negative);
        sphere(group, { ...center, x: boundOffset }, 0.035, colours.positive);
        line(group, [{ ...center, x: -boundOffset }, { ...center, x: boundOffset }], colours.gold);
      }
    }
  }
  line(group, [{ x: -halfGap, y: -side, z: 0 }, { x: halfGap, y: -side, z: 0 }], colours.muted);
  label(group, String.raw`$d = ${number(p, 'gap').toFixed(1)}\,\mathrm{mm}$ (gap $\times50$)`, { x: 0, y: -side - 0.45, z: 0 }, colours.muted, 2.5);
  label(group, String.raw`$${c.voltage.toFixed(2)}\,\mathrm V$ · ${p.connection === 'isolated' ? 'isolated' : 'supply connected'}`, { x: 0, y: side + 0.8, z: 0 }, colours.field, 3);
  grid(group, -3.2);
}

function currentScene(group: THREE.Group, p: Parameters): Animate {
  grid(group, -2);
  const model = wireModel(p);
  const length = number(p, 'length') * 3;
  const radius = Math.sqrt(number(p, 'area')) * 0.4;
  const wire = new THREE.Mesh(new THREE.CylinderGeometry(radius, radius, length, 28, 1, true), new THREE.MeshStandardMaterial({ color: '#bf8158', transparent: true, opacity: 0.12, side: THREE.DoubleSide, depthWrite: false }));
  wire.rotation.z = Math.PI / 2;
  group.add(wire);
  for (let i = 0; i < 6; i++) {
    const x = length * ((i + 0.5) / 6 - 0.5);
    for (let j = 0; j < 6; j++) {
      const angle = j * Math.PI / 3;
      sphere(group, { x, y: radius * 0.8 * Math.cos(angle), z: radius * 0.8 * Math.sin(angle) }, 0.055, colours.gold);
    }
  }
  const carriers: THREE.Mesh[] = [];
  for (let i = 0; i < 32; i++) {
    const angle = i * 2.39996;
    const r = radius * 0.65 * Math.sqrt((i % 7 + 1) / 8);
    carriers.push(sphere(group, { x: length * (i / 32 - 0.5), y: r * Math.cos(angle), z: r * Math.sin(angle) }, 0.065, colours.negative));
  }
  arrow(group, { x: -Math.sign(model.current) * 1.1, y: radius + 0.8, z: 0 }, { x: model.current, y: 0, z: 0 }, 2.2, colours.field);
  label(group, model.current === 0 ? 'No net current' : 'Conventional current', { x: 0, y: radius + 1.5, z: 0 }, colours.field, 2.8);
  arrow(group, { x: -Math.sign(model.electronDrift) * 1.1, y: -radius - 0.65, z: 0 }, { x: model.electronDrift, y: 0, z: 0 }, 2.2, colours.negative);
  label(group, String.raw`Electron drift · displayed $\times600$`, { x: 0, y: -radius - 1.1, z: 0 }, colours.negative, 3.3);
  const velocity = model.electronDrift * 600 * 3;
  return delta => {
    for (const electron of carriers) {
      electron.position.x = ((electron.position.x + length / 2 + delta * velocity) % length + length) % length - length / 2;
    }
  };
}

function semiconductorScene(group: THREE.Group, p: Parameters): Animate {
  grid(group, -3.5);
  const s = carrierModel(p);
  // 3 display units per eV; approximate intrinsic level at midgap.
  const bandEdge = 1.12 / 2 * 3;
  box(group, { x: 0, y: bandEdge + 0.4, z: 0 }, { x: 7, y: 0.8, z: 3 }, colours.negative, 0.12);
  box(group, { x: 0, y: -bandEdge - 0.4, z: 0 }, { x: 7, y: 0.8, z: 3 }, colours.positive, 0.12);
  line(group, [{ x: -3.5, y: 0, z: 1.6 }, { x: 3.5, y: 0, z: 1.6 }], colours.muted, 0.7);
  line(group, [{ x: -3.6, y: s.fermiShift * 3, z: 1.65 }, { x: 3.6, y: s.fermiShift * 3, z: 1.65 }], colours.gold);
  label(group, String.raw`$E_F$`, { x: 4, y: s.fermiShift * 3, z: 1.65 }, colours.gold, 0.7);
  label(group, String.raw`$E_i\approx\text{midgap}$`, { x: -4.4, y: 0, z: 1.65 }, colours.muted, 1.7);
  label(group, 'Conduction band', { x: 0, y: bandEdge + 1.1, z: 0 }, colours.negative, 2.8);
  label(group, 'Valence band', { x: 0, y: -bandEdge - 1.1, z: 0 }, colours.positive, 2.8);
  label(group, String.raw`$E_g=1.12\,\mathrm{eV}$`, { x: 4.1, y: 0, z: -1.3 }, colours.muted, 1.8);
  line(group, [{ x: 3.8, y: -bandEdge, z: -1.3 }, { x: 3.8, y: bandEdge, z: -1.3 }], colours.muted);
  const particles: { mesh: THREE.Mesh; base: THREE.Vector3; phase: number }[] = [];
  for (const [density, sign] of [[s.electrons, 1], [s.holes, -1]]) {
    const count = Math.round(3 + Math.max(0, Math.log10(density) - 8) * 2.5);
    for (let i = 0; i < count; i++) {
      const x = 6.4 * ((i * 0.618034) % 1 - 0.5);
      const z = 2.5 * ((i * 0.414214) % 1 - 0.5);
      const mesh = sphere(group, { x, y: sign * (bandEdge + 0.4), z }, 0.075, sign === 1 ? colours.negative : colours.positive);
      if (sign === -1) {
        (mesh.material as THREE.MeshStandardMaterial).wireframe = true;
      }
      particles.push({ mesh, base: mesh.position.clone(), phase: i * 2.4 + sign });
    }
  }
  let time = 0;
  return delta => {
    time += delta;
    for (const particle of particles) {
      particle.mesh.position.x = particle.base.x + 0.12 * Math.sin(time * 1.3 + particle.phase);
      particle.mesh.position.z = particle.base.z + 0.12 * Math.cos(time * 1.7 + particle.phase);
    }
  };
}

export function buildExperiment(group: THREE.Group, id: ExperimentId, p: Parameters): Animate {
  switch (id) {
    case 'field': fieldScene(group, p); break;
    case 'flux': fluxScene(group, p); break;
    case 'energy': energyScene(group, p); break;
    case 'potential': potentialScene(group, p); break;
    case 'capacitance': plateScene(group, p, false); break;
    case 'dielectric': plateScene(group, p, true); break;
    case 'current': return currentScene(group, p);
    case 'semiconductor': return semiconductorScene(group, p);
    default: return buildTestScene(group, id, p);
  }
  return still;
}