// MLS-MPM over gaussian kernels.
//
// An independent implementation of the idea in PhysGaussian (arXiv 2311.12198):
// the gaussian kernels of a splat reconstruction *are* the continuum particles,
// so there is no mesh, no tetrahedralisation, and no round trip between a
// simulation representation and a render representation. Each kernel carries a
// deformation gradient F, and its covariance is transported by it:
//
//     Σ' = F Σ Fᵀ
//
// which is what makes a stretched blob render as a stretched blob rather than a
// blob that moved. That upstream project ships no licence, so nothing is taken
// from it — this follows the method as described in the paper.
//
// The transfer is MLS-MPM (Hu et al. 2018): particle to grid, solve on the grid,
// grid back to particles. The APIC affine velocity C doubles as the velocity
// gradient, which is why MLS-MPM needs no explicit gradient estimation.

// Stiffness has to be read against the scene, not copied from a table of real
// materials. The box is one unit tall at g=9.8, so the stress a falling object
// generates here is of order rho·v² ≈ 5; at E=4000 that is a fifth of a percent of
// strain — correct, invisible, and useless as a demo. These values put the
// interesting range of each material inside the stress this scene actually
// produces, keeping their ordering intact.
export const MATERIALS = {
    foam: {
        label: 'Foam', E: 90, nu: 0.15, rho: 0.3,
        model: 'plastic', yield: 0.03, hardening: 0, color: '#e2c8d8',
    },
    jelly: {
        label: 'Jelly', E: 240, nu: 0.32, rho: 1.0,
        model: 'neohookean', hardening: 0, color: '#7fd4a8',
    },
    plasticine: {
        label: 'Plasticine', E: 320, nu: 0.35, rho: 1.2,
        model: 'plastic', yield: 0.012, hardening: 0, color: '#d9a066',
    },
    snow: {
        label: 'Snow', E: 650, nu: 0.2, rho: 0.4,
        model: 'snow', critC: 2.5e-2, critS: 6.5e-3, hardening: 10, color: '#eef4f8',
    },
    sand: {
        label: 'Sand', E: 950, nu: 0.3, rho: 1.6,
        model: 'sand', friction: 38, hardening: 0, color: '#d8c48a',
    },
    metal: {
        label: 'Metal', E: 4200, nu: 0.3, rho: 4.0,
        model: 'plastic', yield: 0.004, hardening: 0, color: '#9aa3ad',
    },
};

// 32 rather than 40: the grid is cleared every substep, and 41³ cells cost ~275k
// float writes per step before any physics happens. Dropping to 33³ halves that
// for a difference in the result you cannot see at this scale.
const N = 32;                    // grid resolution per axis
const DX = 1 / N;
const INV_DX = N;

// ---------------------------------------------------------------- 3x3 helpers --
const m3 = () => new Float32Array(9);
const ident = () => Float32Array.from([1, 0, 0, 0, 1, 0, 0, 0, 1]);

function mul(a, b, out) {
    for (let i = 0; i < 3; i++) {
        for (let j = 0; j < 3; j++) {
            out[i * 3 + j] = a[i * 3] * b[j] + a[i * 3 + 1] * b[3 + j] + a[i * 3 + 2] * b[6 + j];
        }
    }
    return out;
}
function transpose(a, out) {
    out[0] = a[0]; out[1] = a[3]; out[2] = a[6];
    out[3] = a[1]; out[4] = a[4]; out[5] = a[7];
    out[6] = a[2]; out[7] = a[5]; out[8] = a[8];
    return out;
}
function det3(a) {
    return a[0] * (a[4] * a[8] - a[5] * a[7])
         - a[1] * (a[3] * a[8] - a[5] * a[6])
         + a[2] * (a[3] * a[7] - a[4] * a[6]);
}

/**
 * Polar decomposition F = R S by Newton iteration on R.
 * The rotation is what the elastic stress needs; iterating R -> (R + R^-T)/2
 * converges quickly and avoids a full SVD in the inner loop.
 */
const _polarTmp = m3(), _polarInv = m3();
function polarR(F, R) {
    R.set(F);
    const tmp = _polarTmp, inv = _polarInv;
    // Four iterations reach visual convergence; twelve cost three times as much
    // per particle per substep and change nothing you can see.
    for (let it = 0; it < 4; it++) {
        // inverse-transpose of R
        const d = det3(R);
        if (Math.abs(d) < 1e-12) break;
        const id = 1 / d;
        inv[0] = (R[4] * R[8] - R[5] * R[7]) * id;
        inv[3] = -(R[1] * R[8] - R[2] * R[7]) * id;
        inv[6] = (R[1] * R[5] - R[2] * R[4]) * id;
        inv[1] = -(R[3] * R[8] - R[5] * R[6]) * id;
        inv[4] = (R[0] * R[8] - R[2] * R[6]) * id;
        inv[7] = -(R[0] * R[5] - R[2] * R[3]) * id;
        inv[2] = (R[3] * R[7] - R[4] * R[6]) * id;
        inv[5] = -(R[0] * R[7] - R[1] * R[6]) * id;
        inv[8] = (R[0] * R[4] - R[1] * R[3]) * id;
        let diff = 0;
        for (let i = 0; i < 9; i++) {
            tmp[i] = 0.5 * (R[i] + inv[i]);
            diff += Math.abs(tmp[i] - R[i]);
        }
        R.set(tmp);
        if (diff < 1e-5) break;
    }
    return R;
}

/** Clamp singular values — the mechanism behind snow's cracking and sand's flow. */
const _clR = m3(), _clRt = m3(), _clS = m3();
function clampSingular(F, lo, hi) {
    // Jacobi SVD on the small symmetric F^T F is overkill here; clamping the
    // stretch through a few polar iterations keeps the volume in range and is
    // stable enough for a real-time demo.
    const R = _clR, Rt = _clRt, S = _clS;
    polarR(F, R);
    transpose(R, Rt);
    mul(Rt, F, S);
    for (let i = 0; i < 3; i++) {
        const d = S[i * 3 + i];
        S[i * 3 + i] = Math.min(hi, Math.max(lo, d));
    }
    mul(R, S, F);
    return F;
}

export class Simulation {
    /**
     * @param {Float32Array} positions xyz per particle, expected inside the unit box
     * @param {Float32Array} scales    per-particle gaussian sigma (xyz)
     */
    constructor(positions, scales, material = 'jelly') {
        this.n = positions.length / 3;
        this.x = Float32Array.from(positions);
        this.v = new Float32Array(this.n * 3);
        this.C = new Float32Array(this.n * 9);          // affine velocity (APIC)
        this.F = new Float32Array(this.n * 9);          // deformation gradient
        this.Jp = new Float32Array(this.n).fill(1);     // plastic volume ratio
        this.sigma0 = Float32Array.from(scales);
        this.sigma = Float32Array.from(scales);
        this.rot = new Float32Array(this.n * 9);        // per-particle frame for render
        for (let i = 0; i < this.n; i++) {
            this.F.set(ident(), i * 9);
            this.rot.set(ident(), i * 9);
        }
        this.setMaterial(material);

        const g = (N + 1) ** 3;
        this.gv = new Float32Array(g * 3);
        this.gm = new Float32Array(g);
        this.gravity = -9.8;
        this.time = 0;

        // Scratch reused across every particle and substep. Allocating these inside
        // the loop was costing more than the arithmetic: at 4.5k particles and 12
        // substeps a frame that is ~1.6M short-lived arrays for the GC to chase.
        this._wx = new Float32Array(3);
        this._wy = new Float32Array(3);
        this._wz = new Float32Array(3);
        this._R = m3(); this._Rt = m3(); this._FT = m3();
        this._tmpA = m3(); this._stress = m3(); this._affine = m3();
        this._newC = m3(); this._Fnew = m3(); this._upd = m3();
        this._Fp = m3();
    }

    setMaterial(name) {
        const m = MATERIALS[name] || MATERIALS.jelly;
        this.mat = m;
        this.matName = name;
        this.mu0 = m.E / (2 * (1 + m.nu));
        this.lambda0 = m.E * m.nu / ((1 + m.nu) * (1 - 2 * m.nu));
        this.pMass = m.rho * DX * DX * DX * 0.25;
        this.pVol = DX * DX * DX * 0.25;
    }

    reset(positions) {
        this.x.set(positions);
        this.v.fill(0);
        this.C.fill(0);
        this.Jp.fill(1);
        this.sigma.set(this.sigma0);
        for (let i = 0; i < this.n; i++) {
            this.F.set(ident(), i * 9);
            this.rot.set(ident(), i * 9);
        }
        this.time = 0;
    }

    step(dt) {
        const { x, v, C, F, Jp, gv, gm } = this;
        gv.fill(0); gm.fill(0);
        const stride = N + 1;
        const gi = (i, j, k) => (i * stride + j) * stride + k;

        const R = this._R, FT = this._FT, tmpA = this._tmpA;
        const stress = this._stress, affine = this._affine;
        const wx = this._wx, wy = this._wy, wz = this._wz;

        // ---- particle to grid
        for (let p = 0; p < this.n; p++) {
            const px = x[p * 3] * INV_DX, py = x[p * 3 + 1] * INV_DX, pz = x[p * 3 + 2] * INV_DX;
            const bi = Math.floor(px - 0.5), bj = Math.floor(py - 0.5), bk = Math.floor(pz - 0.5);
            if (bi < 0 || bj < 0 || bk < 0 || bi + 2 >= stride || bj + 2 >= stride || bk + 2 >= stride) continue;
            const fx = px - bi, fy = py - bj, fz = pz - bk;
            wx[0] = 0.5 * (1.5 - fx) ** 2; wx[1] = 0.75 - (fx - 1) ** 2; wx[2] = 0.5 * (fx - 0.5) ** 2;
            wy[0] = 0.5 * (1.5 - fy) ** 2; wy[1] = 0.75 - (fy - 1) ** 2; wy[2] = 0.5 * (fy - 0.5) ** 2;
            wz[0] = 0.5 * (1.5 - fz) ** 2; wz[1] = 0.75 - (fz - 1) ** 2; wz[2] = 0.5 * (fz - 0.5) ** 2;

            const Fp = this._Fp;
            for (let q = 0; q < 9; q++) Fp[q] = F[p * 9 + q];
            const J = det3(Fp);
            let mu = this.mu0, lambda = this.lambda0;
            if (this.mat.hardening) {
                const h = Math.exp(this.mat.hardening * (1 - Jp[p]));
                mu *= h; lambda *= h;
            }
            if (this.mat.model === 'sand') mu *= 0.35;    // grains shear far more easily

            // fixed corotated stress:  2mu (F - R) F^T + lambda (J-1) J I
            polarR(Fp, R);
            transpose(Fp, FT);
            for (let i = 0; i < 9; i++) tmpA[i] = Fp[i] - R[i];
            mul(tmpA, FT, stress);
            const vol = lambda * (J - 1) * J;
            for (let i = 0; i < 9; i++) stress[i] *= 2 * mu;
            stress[0] += vol; stress[4] += vol; stress[8] += vol;

            const k = -dt * this.pVol * 4 * INV_DX * INV_DX;
            for (let i = 0; i < 9; i++) affine[i] = stress[i] * k + this.pMass * C[p * 9 + i];

            for (let a = 0; a < 3; a++) {
                for (let b = 0; b < 3; b++) {
                    for (let c = 0; c < 3; c++) {
                        const w = wx[a] * wy[b] * wz[c];
                        const dpx = (a - fx) * DX, dpy = (b - fy) * DX, dpz = (c - fz) * DX;
                        const g = gi(bi + a, bj + b, bk + c);
                        gm[g] += w * this.pMass;
                        gv[g * 3] += w * (this.pMass * v[p * 3] + affine[0] * dpx + affine[1] * dpy + affine[2] * dpz);
                        gv[g * 3 + 1] += w * (this.pMass * v[p * 3 + 1] + affine[3] * dpx + affine[4] * dpy + affine[5] * dpz);
                        gv[g * 3 + 2] += w * (this.pMass * v[p * 3 + 2] + affine[6] * dpx + affine[7] * dpy + affine[8] * dpz);
                    }
                }
            }
        }

        // ---- grid update: gravity, boundaries, friction on the floor
        for (let i = 0; i <= N; i++) {
            for (let j = 0; j <= N; j++) {
                for (let k = 0; k <= N; k++) {
                    const g = gi(i, j, k);
                    const m = gm[g];
                    if (m <= 0) continue;
                    let vx = gv[g * 3] / m, vy = gv[g * 3 + 1] / m + dt * this.gravity, vz = gv[g * 3 + 2] / m;
                    // A thin grid boundary only; the real contact is applied per
                    // particle after the transfer.
                    const B = 2;
                    if (i < B && vx < 0) vx = 0;
                    if (i > N - B && vx > 0) vx = 0;
                    if (k < B && vz < 0) vz = 0;
                    if (k > N - B && vz > 0) vz = 0;
                    if (j > N - B && vy > 0) vy = 0;
                    if (j < B && vy < 0) vy = 0;
                    gv[g * 3] = vx; gv[g * 3 + 1] = vy; gv[g * 3 + 2] = vz;
                }
            }
        }

        // ---- grid to particle, and transport the covariance by F
        const newC = this._newC, Fnew = this._Fnew, upd = this._upd;
        for (let p = 0; p < this.n; p++) {
            const px = x[p * 3] * INV_DX, py = x[p * 3 + 1] * INV_DX, pz = x[p * 3 + 2] * INV_DX;
            const bi = Math.floor(px - 0.5), bj = Math.floor(py - 0.5), bk = Math.floor(pz - 0.5);
            if (bi < 0 || bj < 0 || bk < 0 || bi + 2 >= stride || bj + 2 >= stride || bk + 2 >= stride) continue;
            const fx = px - bi, fy = py - bj, fz = pz - bk;
            wx[0] = 0.5 * (1.5 - fx) ** 2; wx[1] = 0.75 - (fx - 1) ** 2; wx[2] = 0.5 * (fx - 0.5) ** 2;
            wy[0] = 0.5 * (1.5 - fy) ** 2; wy[1] = 0.75 - (fy - 1) ** 2; wy[2] = 0.5 * (fy - 0.5) ** 2;
            wz[0] = 0.5 * (1.5 - fz) ** 2; wz[1] = 0.75 - (fz - 1) ** 2; wz[2] = 0.5 * (fz - 0.5) ** 2;

            let nvx = 0, nvy = 0, nvz = 0;
            newC.fill(0);
            for (let a = 0; a < 3; a++) {
                for (let b = 0; b < 3; b++) {
                    for (let c = 0; c < 3; c++) {
                        const w = wx[a] * wy[b] * wz[c];
                        const g = gi(bi + a, bj + b, bk + c);
                        const gx = gv[g * 3], gy = gv[g * 3 + 1], gz = gv[g * 3 + 2];
                        nvx += w * gx; nvy += w * gy; nvz += w * gz;
                        const dpx = (a - fx), dpy = (b - fy), dpz = (c - fz);
                        const s = 4 * INV_DX * w;
                        newC[0] += s * gx * dpx; newC[1] += s * gx * dpy; newC[2] += s * gx * dpz;
                        newC[3] += s * gy * dpx; newC[4] += s * gy * dpy; newC[5] += s * gy * dpz;
                        newC[6] += s * gz * dpx; newC[7] += s * gz * dpy; newC[8] += s * gz * dpz;
                    }
                }
            }
            v[p * 3] = nvx; v[p * 3 + 1] = nvy; v[p * 3 + 2] = nvz;
            C.set(newC, p * 9);
            let nx = x[p * 3] + dt * nvx;
            let ny = x[p * 3 + 1] + dt * nvy;
            let nz = x[p * 3 + 2] + dt * nvz;

            // Contact at the particle, not only on the grid. Zeroing downward
            // velocity across a three-cell band lets a soft material sink into it
            // and compact to a wafer, because nothing stops a particle that is
            // already inside. A hard plane cannot be passed at any stiffness.
            const FLOOR = 0.02, WALL = 0.02;
            if (ny < FLOOR) {
                ny = FLOOR;
                if (nvy < 0) {
                    v[p * 3 + 1] = 0;
                    const fr = this.mat.model === 'sand' ? 0.6 : 0.2;
                    v[p * 3] *= (1 - fr); v[p * 3 + 2] *= (1 - fr);
                }
            }
            if (nx < WALL) { nx = WALL; if (nvx < 0) v[p * 3] = 0; }
            if (nx > 1 - WALL) { nx = 1 - WALL; if (nvx > 0) v[p * 3] = 0; }
            if (nz < WALL) { nz = WALL; if (nvz < 0) v[p * 3 + 2] = 0; }
            if (nz > 1 - WALL) { nz = 1 - WALL; if (nvz > 0) v[p * 3 + 2] = 0; }
            if (ny > 1 - WALL) { ny = 1 - WALL; if (nvy > 0) v[p * 3 + 1] = 0; }

            x[p * 3] = nx; x[p * 3 + 1] = ny; x[p * 3 + 2] = nz;

            // F <- (I + dt C) F
            const Fp = this._Fp;
            for (let q = 0; q < 9; q++) Fp[q] = F[p * 9 + q];
            upd.set(newC);
            for (let i = 0; i < 9; i++) upd[i] *= dt;
            upd[0] += 1; upd[4] += 1; upd[8] += 1;
            mul(upd, Fp, Fnew);

            switch (this.mat.model) {
                case 'snow':
                    clampSingular(Fnew, 1 - this.mat.critC, 1 + this.mat.critS);
                    Jp[p] = Math.max(0.6, Math.min(1.4, det3(Fnew)));
                    break;
                case 'sand':
                    clampSingular(Fnew, 0.92, 1.06);
                    break;
                case 'plastic': {
                    const y = this.mat.yield;
                    clampSingular(Fnew, 1 - y, 1 + y * 3);
                    break;
                }
                default:
                    break;                                  // jelly: purely elastic
            }
            // Σ' = F Σ Fᵀ. With Σ diagonal the render only needs the transformed
            // axes, so the frame is stored and the scales come out of its columns.
            for (let q = 0; q < 9; q++) {
                F[p * 9 + q] = Fnew[q];
                this.rot[p * 9 + q] = Fnew[q];
            }
        }
        this.time += dt;
    }
}