import * as THREE from 'three';
import { OrbitControls } from './vendor/OrbitControls.js';
import { Simulation, MATERIALS } from './mpm.js';
import { makeAsset, SHAPE_LIST } from './assets.js';

const canvas = document.getElementById('view');
const renderer = new THREE.WebGLRenderer({ canvas, antialias: true, preserveDrawingBuffer: true });
renderer.setPixelRatio(Math.min(devicePixelRatio, 2));
renderer.outputColorSpace = THREE.SRGBColorSpace;

const scene = new THREE.Scene();
scene.background = new THREE.Color('#10151b');

const camera = new THREE.PerspectiveCamera(45, 1, 0.01, 100);
camera.position.set(1.5, 1.05, 1.6);

const controls = new OrbitControls(camera, canvas);
controls.enableDamping = true;
controls.target.set(0.5, 0.35, 0.5);

scene.add(new THREE.HemisphereLight('#cfe0f0', '#20262e', 1.3));
const key = new THREE.DirectionalLight('#fff4e2', 1.6);
key.position.set(2, 3, 1.5);
scene.add(key);
const rim = new THREE.DirectionalLight('#6ea8d8', 0.7);
rim.position.set(-2, 1.2, -1.5);
scene.add(rim);

// the unit box the solver works in
const floor = new THREE.Mesh(
    new THREE.PlaneGeometry(1, 1).rotateX(-Math.PI / 2).translate(0.5, 0, 0.5),
    new THREE.MeshStandardMaterial({ color: '#2b333d', roughness: 0.95 })
);
scene.add(floor);
const cage = new THREE.LineSegments(
    new THREE.EdgesGeometry(new THREE.BoxGeometry(1, 1, 1).translate(0.5, 0.5, 0.5)),
    new THREE.LineBasicMaterial({ color: '#33414f' })
);
scene.add(cage);

// ------------------------------------------------------------------ kernels ---
// One instanced ellipsoid per gaussian. A unit sphere scaled by the transformed
// axes is exactly what Σ' = F Σ Fᵀ describes, so the deformation is visible in
// the shape of each kernel and not only in where it moved.
let sim = null, asset = null, mesh = null;
const dummy = new THREE.Object3D();
const mat3 = new THREE.Matrix4();

function buildMesh(n, colors) {
    if (mesh) { scene.remove(mesh); mesh.geometry.dispose(); mesh.material.dispose(); }
    mesh = new THREE.InstancedMesh(
        new THREE.SphereGeometry(1, 6, 4),
        new THREE.MeshStandardMaterial({ roughness: 0.55, metalness: 0.05, vertexColors: false }),
        n
    );
    mesh.instanceColor = new THREE.InstancedBufferAttribute(Float32Array.from(colors), 3);
    mesh.frustumCulled = false;
    scene.add(mesh);
}

function syncMesh() {
    const { x, sigma, rot, n } = sim;
    for (let i = 0; i < n; i++) {
        // columns of F scaled by sigma give the transformed principal axes
        const f = rot.subarray(i * 9, i * 9 + 9);
        const sx = sigma[i * 3], sy = sigma[i * 3 + 1], sz = sigma[i * 3 + 2];
        mat3.set(
            f[0] * sx, f[1] * sy, f[2] * sz, x[i * 3],
            f[3] * sx, f[4] * sy, f[5] * sz, x[i * 3 + 1],
            f[6] * sx, f[7] * sy, f[8] * sz, x[i * 3 + 2],
            0, 0, 0, 1
        );
        mesh.setMatrixAt(i, mat3);
    }
    mesh.instanceMatrix.needsUpdate = true;
}

// --------------------------------------------------------------------- ui -----
const shapeSel = document.getElementById('shape');
const matSel = document.getElementById('material');
const countEl = document.getElementById('count');
const statsEl = document.getElementById('stats');

shapeSel.innerHTML = SHAPE_LIST.map(s => `<option value="${s.id}">${s.label}</option>`).join('');
matSel.innerHTML = Object.entries(MATERIALS)
    .map(([id, m]) => `<option value="${id}">${m.label}</option>`).join('');

let running = false;
// Explicit MPM needs many small steps, but each one is a full P2G/G2P sweep. Eight
// keeps a soft material stable while leaving the frame budget usable.
let substeps = 8;

function rebuild() {
    const n = parseInt(countEl.value, 10) || 4000;
    asset = makeAsset(shapeSel.value, n, 7);
    sim = new Simulation(asset.positions, asset.sigma, matSel.value);
    buildMesh(asset.count, asset.color);
    syncMesh();
    report(0);
}

function report(ms) {
    const m = MATERIALS[sim.matName];
    statsEl.textContent =
        `${asset.count.toLocaleString()} kernels · ${m.label} · E=${m.E.toExponential(1)} · ` +
        `t=${sim.time.toFixed(2)}s${ms ? ` · ${ms.toFixed(1)} ms/frame` : ''}`;
}

document.getElementById('go').onclick = () => {
    running = !running;
    document.getElementById('go').textContent = running ? 'Pause' : 'Drop';
    document.getElementById('go').classList.toggle('on', running);
};
document.getElementById('reset').onclick = () => {
    running = false;
    document.getElementById('go').textContent = 'Drop';
    document.getElementById('go').classList.remove('on');
    rebuild();
};
shapeSel.onchange = () => { running = false; document.getElementById('go').textContent = 'Drop'; rebuild(); };
matSel.onchange = () => {
    sim.setMaterial(matSel.value);
    report(0);
};
countEl.onchange = () => { running = false; rebuild(); };

// ------------------------------------------------------------------- loop -----
function resize() {
    const w = canvas.clientWidth, h = canvas.clientHeight;
    if (canvas.width !== w || canvas.height !== h) {
        renderer.setSize(w, h, false);
        camera.aspect = w / h;
        camera.updateProjectionMatrix();
    }
}

let frameMs = 0;
function tick() {
    if (!sim || !running) return;
    const t0 = performance.now();
    // Many small steps: MPM is explicit, so stability is set by the step size and
    // the stiffest material has to survive it.
    const dt = 2.4e-4;
    for (let s = 0; s < substeps; s++) sim.step(dt);
    syncMesh();
    frameMs = frameMs * 0.85 + (performance.now() - t0) * 0.15;
    report(frameMs);
}

function loop() {
    resize();
    controls.update();
    tick();
    renderer.render(scene, camera);
    requestAnimationFrame(loop);
}

rebuild();
requestAnimationFrame(loop);

window.gp = {
    renderer, scene, camera, controls, sim: () => sim,
    step: (steps = 240) => { const d = 2.4e-4; for (let i = 0; i < steps; i++) sim.step(d); syncMesh(); },
    setShape: (s) => { shapeSel.value = s; rebuild(); },
    setMaterial: (m) => { matSel.value = m; sim.setMaterial(m); },
    frame: () => { resize(); renderer.render(scene, camera); },
};