const positionBuffer = instancedArray(COUNT, 'vec3');
const velocityBuffer = instancedArray(COUNT, 'vec3');
const prevVelocityBuffer = instancedArray(COUNT, 'vec3');
const matrixBuffer = instancedArray(COUNT, 'mat4');
const aoBuffer = instancedArray(COUNT, 'float');

// Build lookAt matrix from direction (reusable in compute)
const buildLookAtMatrix = Fn(([direction, position]) => {
  const up = vec3(0, 1, 0);
  const dir = normalize(direction);

  const right = normalize(cross(up, dir));
  const correctedUp = normalize(cross(dir, right));

  // Build full transformation matrix with rotation + translation
  return mat4(
    vec4(right, 0),
    vec4(correctedUp, 0),
    vec4(dir, 0),
    vec4(position, 1)
  );
});

// computeInit: spawn particles at center with curl noise distribution
const computeInit = Fn(() => {
  const i = float(instanceIndex);

  // Spherical distribution from center
  const phi = i.mul(2.399963);  // golden angle
  const theta = i.div(COUNT).mul(Math.PI);
  const r = i.div(COUNT).sqrt().mul(3.0);

  const pos = vec3(
    r.mul(theta.sin()).mul(phi.cos()),
    r.mul(theta.cos()),
    r.mul(theta.sin()).mul(phi.sin())
  );

  const initVel = vec3(0, 1, 0);

  positionBuffer.element(instanceIndex).assign(pos);
  velocityBuffer.element(instanceIndex).assign(initVel);
  prevVelocityBuffer.element(instanceIndex).assign(initVel);
  matrixBuffer.element(instanceIndex).assign(buildLookAtMatrix(initVel, pos));
  aoBuffer.element(instanceIndex).assign(0.0);
});

renderer.compute(computeInit().compute(COUNT));

// Uniforms
const uFrequency = uniform(FREQUENCY);
const uAmplitude = uniform(AMPLITUDE);
const uSpeed = uniform(SPEED);
const uSteps = uniform(STEPS);
const uDt = uniform(DT);

// computeUpdate: advect with curl noise, compute mat4
const computeUpdateFn = Fn(() => {
  const pos = positionBuffer.element(instanceIndex).toVar();
  const vel = velocityBuffer.element(instanceIndex);
  const prevVel = prevVelocityBuffer.element(instanceIndex);

  // Store previous velocity before computing new one
  prevVel.assign(vel);

  // Oscillating scale for grow/shrink effect
  const scale = sin(time.mul(uSpeed)).mul(0.5).add(1.0);
  const dt = float(uDt).mul(scale);

  // Advect through curl noise field
  Loop(int(uSteps), () => {
    const v = curlNoise(pos, uFrequency, time.mul(0.05), uAmplitude, false);
    pos.addAssign(v.mul(dt));
  });

  // Store current velocity
  vel.assign(pos.sub(positionBuffer.element(instanceIndex)));

  // Smooth direction from mix of prev and current velocity
  const smoothDir = normalize(mix(prevVel, vel, 0.5).add(vec3(0.0001)));

  // Compute and store mat4
  matrixBuffer.element(instanceIndex).assign(buildLookAtMatrix(smoothDir, pos));

  // Fake AO based on multiple factors:
  // 1. Distance from center (darker at center where particles cluster)
  const distFromCenter = length(pos);
  const centerAo = float(1.0).sub(clamp(distFromCenter.div(5.0), 0.0, 1.0));

  // 2. Vertical position (darker at bottom)
  const heightAo = clamp(pos.z.div(4.8).add(0.5), 0.0, 1.0);

  // 3. Velocity magnitude (slower = more clustered = darker)
  const speed = length(vel);
  const speedAo = float(1.0).sub(clamp(speed.mul(50.0), 0.0, 1.0));

  // Combine AO factors
  const ao = clamp(centerAo.oneMinus().mul(1).add(heightAo.mul(0.3)).add(speedAo.mul(0.3)), 0.0, 1.0);
  aoBuffer.element(instanceIndex).assign(ao);

  positionBuffer.element(instanceIndex).assign(pos);
});

computeUpdate = computeUpdateFn().compute(COUNT);

// Material with direction-based coloring
const material = new THREE.MeshPhysicalNodeMaterial({
  metalness: 0.8,
  roughness: 0.3
});

const vel = velocityBuffer.element(instanceIndex);
const prevVel = prevVelocityBuffer.element(instanceIndex);

// Color based on velocity direction
const smoothDir = normalize(mix(prevVel, vel, 0.5).add(vec3(0.0001)));
const colX = abs(smoothDir.x);
const colY = abs(smoothDir.y);
const colZ = abs(smoothDir.z);

const cA = color(0x00ffff);  // cyan
const cB = color(0xff00ff);  // magenta
const cC = color(0xffff00);  // yellow

material.colorNode = mix(mix(cA, cB, colX), cC, colY.add(colZ).mul(0.5));

// Apply fake AO from buffer
material.aoNode = aoBuffer.element(instanceIndex);

// Mesh with small box geometry
const geometry = new THREE.BoxGeometry(0.08, 0.28, 0.08);
const mesh = new THREE.InstancedMesh(geometry, material, COUNT);

// Use mat4 buffer directly for instance matrices
mesh.instanceMatrix = matrixBuffer.value;
scene.add(mesh);