<script src="https://threejs.org/build/three.min.js"></script>

<script src="https://threejs.org/examples/js/shaders/CopyShader.js"></script>

<script src="https://threejs.org/examples/js/postprocessing/EffectComposer.js"></script>
<script src="https://threejs.org/examples/js/postprocessing/RenderPass.js"></script>
<script src="https://threejs.org/examples/js/postprocessing/ShaderPass.js"></script>
<script src="https://threejs.org/examples/js/shaders/FXAAShader.js"></script>

<script src="https://threejs.org/examples/js/controls/OrbitControls.js"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/dat-gui/0.7.5/dat.gui.min.js"></script>

html, body {
  height: 100%;
  margin: 0;
  overflow: hidden;
}
canvas {
  width: 100%;
  height: 100%;
  display; block;
}

var materialShader;

var scene = new THREE.Scene();
var camera = new THREE.PerspectiveCamera(60, window.innerWidth / window.innerHeight, 1, 5000);
camera.position.setScalar(125);
renderer = new THREE.WebGLRenderer();
renderer.autoClear = false;
renderer.setPixelRatio( window.devicePixelRatio );
renderer.setClearColor(0x404040);
renderer.setSize( window.innerWidth, window.innerHeight );
document.body.appendChild( renderer.domElement );

/////////////////////////////////////////////////////////

var renderPass = new THREE.RenderPass(scene, camera);

//

fxaaPass = new THREE.ShaderPass(THREE.FXAAShader);
fxaaPass.renderToScreen = true;

var pixelRatio = renderer.getPixelRatio();

fxaaPass.material.uniforms["resolution"].value.x = 1 / (window.innerWidth * pixelRatio);
fxaaPass.material.uniforms["resolution"].value.y = 1 / (window.innerHeight * pixelRatio);

var composer = new THREE.EffectComposer(renderer);
composer.addPass(renderPass);
composer.addPass(fxaaPass);

/////////////////////////////////////////////////////////

var controls = new THREE.OrbitControls(camera, renderer.domElement);

var MAX_POINTS = 10000;

var colors = [
  new THREE.Color(1, 0, 1),
  new THREE.Color(1, 1, 0),
  new THREE.Color(0, 1, 1)
];

var pts = [];
var clrs = [];
var shpIdx = [];
var shpScl = [];

for (let i = 0; i < MAX_POINTS; i++) {
  pts.push(
    new THREE.Vector3(
      THREE.Math.randFloatSpread(250),
      THREE.Math.randFloatSpread(250),
      THREE.Math.randFloatSpread(250)
    )
  );
  let clrIdx = THREE.Math.randInt(0, 2);
  clrs.push(colors[clrIdx].r, colors[clrIdx].g, colors[clrIdx].b);
  shpIdx.push(
    clrIdx // correlation between shape and colour :)
  );
  shpScl.push(
    1 // all have the same scale
  );
}

var geometry = new THREE.BufferGeometry().setFromPoints(pts);
geometry.addAttribute(
  "color",
  new THREE.BufferAttribute(new Float32Array(clrs), 3)
);
geometry.addAttribute(
  "shapeIdx",
  new THREE.BufferAttribute(new Float32Array(shpIdx), 1)
);
geometry.addAttribute(
  "shapeScale",
  new THREE.BufferAttribute(new Float32Array(shpScl), 1)
);

var material = new THREE.PointsMaterial({
  size: 5,
  vertexColors: THREE.VertexColors
});

material.onBeforeCompile = shader => {
  shader.uniforms.inner = { value: 0.25 };
  shader.uniforms.outer = { value: 0.5 };
  shader.vertexShader =
    `
    attribute float shapeIdx;
    attribute float shapeScale;
    
    varying float vShapeIdx;
  ` + shader.vertexShader;
  shader.vertexShader = shader.vertexShader.replace(
    "gl_PointSize = size;",
    `gl_PointSize = size * shapeScale;
    vShapeIdx = shapeIdx;
    `
  );
  shader.fragmentShader =
    `
    uniform float inner;
    uniform float outer;
    varying float vShapeIdx;
  ` + shader.fragmentShader;
  shader.fragmentShader = shader.fragmentShader.replace(
    "void main() {",
    `
    float circle(float inner, float outer, vec2 uv){
      float dist = length(uv);
      return smoothstep(inner - 0.01, inner, dist) - smoothstep(outer - 0.01, outer,  dist);
    }
    
    float square(float inner, float outer, vec2 uv){
      vec2 absUv = abs(uv);
      float maxUv = abs(max(absUv.x, absUv.y));
      return smoothstep(inner - 0.01, inner, maxUv) - smoothstep(outer - 0.01, outer, maxUv);
    }
    
    // honestly stolen from https://www.shadertoy.com/view/lsBfRc
    float triangle(float inner, float outer, vec2 uv){
        uv *= 2.0;
        float a = atan(uv.x,-uv.y)+PI;
        float r = 3.1415926 * 2.0 / 3.0;

        float d = cos( floor( .5 + a/r ) * r - a ) * length( uv );

        return smoothstep(inner - 0.01, inner, d) - smoothstep(outer - 0.01, outer,d);
    }
    
    void main() {`
  );
  shader.fragmentShader = shader.fragmentShader.replace(
    `gl_FragColor = vec4( outgoingLight, diffuseColor.a );`,
    `vec2 uv = gl_PointCoord;
    uv -= 0.5;

    float color = 0.0;
    //float shape = 0.0;

    float shapeIdx = floor(vShapeIdx + 0.5);

    if (shapeIdx == 0.0){
      //shape = circle(0.0, 0.5, uv);
      color = circle(inner, outer, uv);
    } else if (shapeIdx == 1.0) {
      //shape = 1.0;
      color = square(inner, outer, uv);
    } else if (shapeIdx == 2.0){
      //shape = triangle(0.0, 0.5, uv);
      color = triangle(inner, outer, uv);
    }

    if (color < 0.5) discard;

    vec3 baseColor = vec3(0);
    vec3 col = mix(baseColor, vColor.rgb, color);
    gl_FragColor = vec4( col, diffuseColor.a );
    `
  );

  materialShader = shader;
  var gui = new dat.GUI();
  gui.add(materialShader.uniforms.inner, "value", 0.0, 0.4).name("inner");
};

var points = new THREE.Points(geometry, material);

scene.add(points);

window.addEventListener( 'resize', onWindowResize, false );

animate();

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

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

  var pixelRatio = renderer.getPixelRatio();

  fxaaPass.material.uniforms["resolution"].value.x =
    1 / (window.innerWidth * pixelRatio);
  fxaaPass.material.uniforms["resolution"].value.y =
    1 / (window.innerHeight * pixelRatio);
}

function animate() {
  requestAnimationFrame(animate);
  renderer.render(scene, camera);
}