import * as THREE from "three";
import {
  positionLocal,
  normalLocal,
  normalize,
  modelWorldMatrix,
  cameraProjectionMatrix,
  cameraViewMatrix,
  mix,
  attributeArray,
  clamp,
  time,
  mx_noise_float,
  Fn,
  uint,
  float,
  cross,
  If,
  Continue,
  distance,
  length,
  attribute,
  max,
  exp,
  mat3,
  vec3,
  select,
  Loop,
  instanceIndex,
  uniform
} from "three/tsl";
import { OrbitControls } from "three/addons/controls/OrbitControls.js";
import { GLTFLoader } from "three/addons/loaders/GLTFLoader.js";
import { DRACOLoader } from "three/addons/loaders/DRACOLoader.js";
import { GUI } from "three/addons/libs/lil-gui.module.min.js";
import Stats from 'three/addons/libs/stats.module.js';

class FlockGeometry extends THREE.BufferGeometry {
  constructor(geo) {
    super();

    const geometry = geo.toNonIndexed();
    const srcPosAttr = geometry.getAttribute( "position" );
    const srcNormAttr = geometry.getAttribute( "normal" );
    let count = srcPosAttr.count;
    const total = count * BOIDS;
    
    const posAttr = new THREE.BufferAttribute(new Float32Array(total * 3), 3); 
    const normAttr = new THREE.BufferAttribute(new Float32Array(total * 3), 3); 
    const instanceIDAttr = new THREE.BufferAttribute(new Uint32Array(total), 1);

    this.setAttribute("instanceID", instanceIDAttr);
    this.setAttribute("position", posAttr);
    this.setAttribute("normal", normAttr);


    for (let b = 0; b < BOIDS; b++) {
      let offset = b * count * 3;
      for (let i = 0; i < count * 3; i++) {
        posAttr.array[offset + i] = srcPosAttr.array[i];
        normAttr.array[offset + i] = srcNormAttr.array[i];
      }
  
      offset = b * count;
      for (let i = 0; i < count; i++) {
        instanceIDAttr.array[offset + i] = b
      }
    }
  }
}

let container;

let camera,
  scene,
  renderer,
  options,
  material,
  assetPath,
  clock,
  boid,
  flock,
  deltaTime,
  stats,
  computeVelocity,
  computePosition,
  computeTest;

const BOIDS = 2000;
const BOUNDS = 20,
  BOUNDS_HALF = BOUNDS / 2;

init();

function init() {
  container = document.createElement("div");
  document.body.appendChild(container);
  
  stats = new Stats();
  stats.domElement.style.right = "0px";
	container.appendChild( stats.dom );

  camera = new THREE.PerspectiveCamera(
    40,
    window.innerWidth / window.innerHeight,
    1,
    100
  );
  camera.position.set(0.0, 1, 12);

  //

  scene = new THREE.Scene();
  scene.background = new THREE.Color(0x444488);

  //

  renderer = new THREE.WebGPURenderer({ antialias: true });
  renderer.setPixelRatio(window.devicePixelRatio);
  renderer.setSize(window.innerWidth, window.innerHeight);
  renderer.setAnimationLoop(render);
  container.appendChild(renderer.domElement);

  //

  //content

  const ambient = new THREE.HemisphereLight(0xaaaaaa, 0x333333);
  const light = new THREE.DirectionalLight(0xffffff, 3);
  light.position.set(3, 3, 1);
  scene.add(ambient);
  scene.add(light);

  clock = new THREE.Clock();
  
  const controls = new OrbitControls(camera, renderer.domElement);

  assetPath = "https://assets.codepen.io/2666677/";

  loadGLB("boid");

  window.addEventListener("resize", onWindowResize);
}

function loadGLB(name) {
  const loader = new GLTFLoader().setPath(assetPath);
  const dracoLoader = new DRACOLoader();
  dracoLoader.setDecoderPath(
    "https://cdn.jsdelivr.net/npm/three@v0.170.0/examples/jsm/libs/draco/gltf/"
  );
  loader.setDRACOLoader(dracoLoader);

  loader.load(`${name}.glb`, (gltf) => {
    boid = gltf.scene.children[0];
    const scale = 0.2;
    boid.geometry.scale( scale, scale, scale );

    tsl();
    //scene.add(boid);
  });
}

function initStorage() {
  const positionArray = new Float32Array(BOIDS * 3);
  const directionArray = new Float32Array(BOIDS * 3);
  const noiseArray = new Float32Array(BOIDS);

  const q = new THREE.Quaternion();
  const v = new THREE.Euler();

  for (let i = 0; i < BOIDS; i++) {
    const offset = i * 3;

    for( let j=0; j<3; j++){
      positionArray[offset + j] = Math.random() * BOUNDS - BOUNDS_HALF;
    }

    q.random();
    v.setFromQuaternion(q);

    directionArray[offset + 0] = v.x;
    directionArray[offset + 1] = v.y;
    directionArray[offset + 2] = v.z;
    
    noiseArray[i] = Math.random() * 1000.0;
  }

  const positionStorage = attributeArray(positionArray, "vec3").label(
    "positionStorage"
  );
  const directionStorage = attributeArray(directionArray, "vec3").label(
    "directionStorage"
  );
  const noiseStorage = attributeArray(noiseArray, "float").label(
    "noiseStorage"
  );

  // The Pixel Buffer Object (PBO) is required to get the GPU computed data in the WebGL2 fallback.
  positionStorage.setPBO(true);
  directionStorage.setPBO(true);
  noiseStorage.setPBO(true);

  return [positionStorage, directionStorage, noiseStorage];
}

function tsl() {
  const [positionStorage, directionStorage, noiseStorage] = initStorage();

  deltaTime = uniform(float());
  const boidSpeed = uniform(3);
  const flockPosition = uniform(vec3());
  const neighbourDistance = uniform(5);
  const rotationSpeed = uniform(1);

  const flockVertexTSL = Fn(() => {
    const instanceID = attribute("instanceID");
    const normal = normalLocal.toVar();
    const dir = normalize(directionStorage.element(instanceID)).toVar();
    
    //Create matrix
    //float4x4 create_matrix(float3 pos, float3 dir, float3 up) {
    const zaxis = dir.negate().normalize().toVar();
    const xaxis = cross(vec3(0, 1, 0), zaxis).normalize().toVar();
    const yaxis = cross(zaxis, xaxis).toVar();
    const mat = mat3( 
      xaxis.x,
      yaxis.x,
      zaxis.x,
      xaxis.y,
      yaxis.y,
      zaxis.y,
      xaxis.z,
      yaxis.z,
      zaxis.z
    ).toVar();
    
    const finalVert = modelWorldMatrix.mul(mat.mul(positionLocal)).add(positionStorage.element(instanceID));

    return cameraProjectionMatrix.mul(cameraViewMatrix).mul(finalVert);
  });

  computeVelocity = Fn(() => {
    const boid_pos = positionStorage.element(instanceIndex).toVar();
    const boid_dir = directionStorage.element(instanceIndex).toVar();
    
    const separation = vec3(0).toVar();
    const alignment = vec3(0).toVar();
    const cohesion = vec3(flockPosition).toVar();

    const nearbyCount = uint(1).toVar(); // Add self that is ignored in loop
    
    Loop(
      { start: uint(0), end: uint(BOIDS), type: "uint", condition: "<" },
      ({ i }) => {
        If(i == instanceIndex, () => {
          Continue();
        });

        const tempBoid_pos = positionStorage.element(i).toVar();
        const tempBoid_dir = directionStorage.element(i).toVar();

        const offset = boid_pos.sub(tempBoid_pos).toVar();
        const dist = length(offset).toVar();
        
        If( dist.lessThan(neighbourDistance), () => {
          If( dist.lessThan( 0.0001 ), () => {
							Continue();
						} ); 
          
          //separation += offset * (1.0/dist - 1.0/neighbourDistance);
          const s = offset.mul(float(1.0).div(dist).sub(float(1.0).div(neighbourDistance))).toVar();
          separation.addAssign( s );
          alignment.addAssign(tempBoid_dir);
          cohesion.addAssign(tempBoid_pos);

          nearbyCount.addAssign(1);
        }); //If
      }); //Loop
    
    const avg = float(1.0).div(nearbyCount).toVar();
    alignment.mulAssign(avg);
    cohesion.mulAssign(avg);
    cohesion.assign(cohesion.normalize().sub(boid_pos));

    const direction = alignment.add(separation).add(cohesion).toVar();

    const ip = exp(rotationSpeed.mul(-1).mul(deltaTime));
    boid_dir.assign(mix(direction, boid_dir.normalize(), ip));
    directionStorage.element(instanceIndex).assign(boid_dir);
  })().compute(BOIDS);

  computePosition = Fn( () => {
    const boid_pos = positionStorage.element(instanceIndex).toVar();
    const boid_dir = directionStorage.element(instanceIndex).toVar();
    const noise_offset = noiseStorage.element(instanceIndex).toVar();
    const noise = mx_noise_float( boid_pos.mul( time.div(100.0).add(noise_offset))).add(1).div(2.0).toVar();
	const velocity = boidSpeed.mul(float(1.0).add(noise)).toVar();// * boidSpeedVariation);
    
    boid_pos.addAssign( boid_dir.mul( velocity ).mul(deltaTime));
    positionStorage.element(instanceIndex).assign(boid_pos);
  })().compute(BOIDS);
  
  computeTest = Fn( () => {
    const position = positionStorage.element( instanceIndex );
    position.addAssign( vec3( 0, deltaTime.mul(100), 0) );
  })().compute(BOIDS);
  
  const geometry = new FlockGeometry(boid.geometry);
  const material = new THREE.MeshStandardNodeMaterial();

  flock = new THREE.Mesh(geometry, material);
  scene.add(flock);

  material.vertexNode = flockVertexTSL();
  /*options = {
    delta: 0
  };

  const gui = new GUI();
  gui.add(options, "delta", 0, 1).onChange((value) => {
    deltaUniform.value = value;
  });*/
}

function onWindowResize() {
  camera.aspect = window.innerWidth / window.innerHeight;
  camera.updateProjectionMatrix();

  renderer.setSize(window.innerWidth, window.innerHeight);
}

//

function render() {
  if (deltaTime && clock) deltaTime.value = clock.getDelta();
  if (computeVelocity) renderer.compute(computeVelocity);
  if (computePosition) renderer.compute(computePosition);
  stats.update();
  renderer.render(scene, camera);
}