p5.disableFriendlyErrors = true;

let cam, segmenter, seg = null, maskGfx, sh;
const W = 640, H = 480;

const VERT = `
precision mediump float;
attribute vec3 aPosition;
attribute vec2 aTexCoord;
varying vec2 vUv;
void main(){
  vUv = aTexCoord;
  vec4 pos = vec4(aPosition, 1.0);
  pos.xy = pos.xy * 2.0 - 1.0;
  gl_Position = pos;
}
`;

const FRAG = `
precision mediump float;

uniform vec2  uResolution;
uniform float uTime;
uniform sampler2D uVideo;
uniform sampler2D uMask;
uniform float uHasMask;

varying vec2 vUv;

vec4 tex2D(sampler2D s, vec2 uv){
  return texture2D(s, vec2(uv.x, 1.0 - uv.y));
}

vec3 bg(vec2 uv){
  vec2 p = uv - 0.5;
  float r = length(p);
  vec3 a = 0.6 + 0.4 * cos(6.28318 * (r + vec3(0.0, 0.15, 0.33) + 0.03 * uTime));
  return a;
}

float maskPresent(vec2 uv){
  vec4 c = tex2D(uMask, uv);
  return step(0.001, c.a + c.r + c.g + c.b);
}

float edgeStrength(vec2 uv){
  vec2 px = 1.0 / uResolution;
  float m = maskPresent(uv);
  float neigh = (
    maskPresent(uv + vec2( px.x, 0.0)) +
    maskPresent(uv + vec2(-px.x, 0.0)) +
    maskPresent(uv + vec2(0.0,  px.y)) +
    maskPresent(uv + vec2(0.0, -px.y))
  ) * 0.25;
  return max(0.0, m - neigh);
}

/*
vec3 partsTint(float id){
  return 0.6 + 0.4 * cos(6.28318 * (id * 0.07 + vec3(0.0, 0.15, 0.33)));
}
*/



vec3 partsTint(float id){
  return 0.3 + 0.7 * cos(6.28318 * (id * 0.07 + vec3(0.0, 0.15, 0.33)));
}


void main(){
  vec2 uv = vUv;

  vec4 vid = tex2D(uVideo, uv);
  vec4 msk = tex2D(uMask,  uv);

  float present = (uHasMask > 0.5 ? 1.0 : 0.0) * maskPresent(uv);
  float id = msk.r * 255.0;
  vec3 tint = partsTint(id);

// vec3 tint = vec3(0.1, 1.0, 0.0);

  vec3 base = mix(bg(uv), vid.rgb * 0.7, 0.7);
  vec3 body = vid.rgb * tint; // * mix(vec3(1.0), tint, 0.55);

  // waves
  float waves = sin(uv.x * 100.0 * sin(uTime) * 5.0) * sin(uv.y * 100.0 + uTime * 3.0);
  body += waves * 0.01;

  float edge = smoothstep(0.0, 0.8, edgeStrength(uv)) * 0.8;

  vec3 col = mix(base, body, present);
  col += edge * vec3(1.1, 0.9, 1.2);

  gl_FragColor = vec4(col, 1.0);
}
`;

async function setup(){
  createCanvas(W, H, WEBGL);
  pixelDensity(1);
  noStroke();

  cam = createCapture({ video:{ width:W, height:H, facingMode:'user' }, audio:false });
  cam.size(W, H);
  cam.elt.setAttribute('playsinline','');
  cam.hide();
  await new Promise(res => (cam.elt.readyState >= 2 ? res() : (cam.elt.onloadeddata = res)));

  window.video = cam;

  if (window.tf && tf.setBackend) {
    if (tf.getBackend && tf.getBackend() !== 'webgl') {
      await tf.setBackend('webgl');
    }
    await tf.ready();
  }

  sh = createShader(VERT, FRAG);
  maskGfx = createGraphics(W, H);
  maskGfx.pixelDensity(1);

  segmenter = await ml5.bodySegmentation('BodyPix', { maskType: 'parts' });

  if (segmenter.ready && typeof segmenter.ready.then === 'function') {
    await segmenter.ready;
  } else {
    await new Promise((res) => {
      let n = 0;
      const t = setInterval(() => {
        if (segmenter.segmenter) { clearInterval(t); res(); }
        else if ((n += 1) > 400) { clearInterval(t); res(); }
      }, 25);
    });
  }

  segmenter.detectStart(cam, r => { seg = r; });
}

function draw(){
  if (seg && seg.mask){
    maskGfx.clear();
    maskGfx.image(seg.mask, 0, 0, maskGfx.width, maskGfx.height);
  } else {
    maskGfx.clear();
  }

  shader(sh);
  sh.setUniform('uResolution', [width, height]);
  sh.setUniform('uTime', millis()/1000.0);
  sh.setUniform('uVideo', cam);
  sh.setUniform('uMask',  maskGfx);
  sh.setUniform('uHasMask', (seg && seg.mask) ? 1.0 : 0.0);

  rect(-width/2, -height/2, width, height);
}