// Plasmodio — a Physarum polycephalum transport-network simulation on the GPU (WebGL2).
//
// Model: Jones (2010), "Characteristics of pattern formation and evolution in
// approximations of Physarum transport networks", Artificial Life 16(2).
// Each agent senses the chemoattractant trail at three points ahead of it,
// turns toward the strongest, moves one step and deposits trail. The trail
// diffuses and decays. Cities emit trail as food. Sea is impassable except
// along ferry corridors. That is the whole organism.
//
// Textures
//   agents  RGBA32F  (AW x AH)   x, y in grid px  |  heading (rad)  |  spare
//   trail   R16F     (W x H)     chemoattractant concentration
//   mask    RGBA8    (W x H)     R = land coverage, G = ferry corridor, B = unused

const QUAD_VS = `#version 300 es
precision highp float;
out vec2 vUv;
void main(){
  vec2 p = vec2(float((gl_VertexID << 1) & 2), float(gl_VertexID & 2));
  vUv = p;
  gl_Position = vec4(p * 2.0 - 1.0, 0.0, 1.0);
}`;

const UPDATE_FS = `#version 300 es
precision highp float;
uniform highp sampler2D uAgents;
uniform highp sampler2D uTrail;
uniform highp sampler2D uMask;
uniform vec2 uGrid;
uniform float uSA, uSD, uRA, uSS, uSeed, uStep, uFerry, uRand;
out vec4 o;

float hash12(vec2 p){
  vec3 p3 = fract(vec3(p.xyx) * 0.1031);
  p3 += dot(p3, p3.yzx + 33.33);
  return fract((p3.x + p3.y) * p3.z);
}
// .r = land coverage, .g = ferry corridor; (0,0) outside the grid
vec2 maskAt(vec2 q){
  if (q.x < 0.0 || q.y < 0.0 || q.x >= uGrid.x || q.y >= uGrid.y) return vec2(0.0);
  return texelFetch(uMask, ivec2(q), 0).rg;
}
float sense(vec2 p, float ang){
  vec2 q = p + vec2(cos(ang), sin(ang)) * uSD;
  vec2 m = maskAt(q);
  if (max(m.r, m.g) < 0.5) return -1.0;
  return texelFetch(uTrail, ivec2(q), 0).r;
}
void main(){
  vec4 a = texelFetch(uAgents, ivec2(gl_FragCoord.xy), 0);
  // a.w is the step at which this agent joins the colony
  if (a.w > uStep) { o = a; return; }
  vec2 p = a.xy;
  float h = a.z;
  float F = sense(p, h), L = sense(p, h + uSA), R = sense(p, h - uSA);
  float r = hash12(gl_FragCoord.xy + uSeed);
  if (F > L && F > R) { }
  else if (F < L && F < R) { h += (r < 0.5) ? uRA : -uRA; }
  else if (L > R) { h += uRA; }
  else if (R > L) { h -= uRA; }
  // a little wanderlust: rarely, an agent forgets the trail and picks a new direction
  if (hash12(gl_FragCoord.xy * 0.77 + uSeed + 9.0) < uRand) h = hash12(gl_FragCoord.yx * 1.91 + uSeed) * 6.2831853;
  vec2 np = p + vec2(cos(h), sin(h)) * uSS;
  vec2 mHere = maskAt(p), mNext = maskAt(np);
  bool blocked = max(mNext.r, mNext.g) < 0.5;
  // boarding a ferry (stepping from land onto a corridor) is rare: most of the colony stays ashore
  if (!blocked && mHere.r >= 0.5 && mNext.r < 0.5 && hash12(gl_FragCoord.xy * 1.37 + uSeed + 5.0) > uFerry) blocked = true;
  if (blocked) {
    h = hash12(gl_FragCoord.yx + uSeed * 1.7 + 11.0) * 6.2831853;
    np = p;
  }
  o = vec4(np, h, a.w);
}`;

const DEPOSIT_VS = `#version 300 es
precision highp float;
uniform highp sampler2D uAgents;
uniform int uAW;
uniform vec2 uGrid;
uniform float uStep;
void main(){
  ivec2 ij = ivec2(gl_VertexID % uAW, gl_VertexID / uAW);
  vec4 a = texelFetch(uAgents, ij, 0);
  gl_PointSize = 1.0;
  if (a.w > uStep) { gl_Position = vec4(-2.0, -2.0, 0.0, 1.0); return; }
  gl_Position = vec4((a.xy / uGrid) * 2.0 - 1.0, 0.0, 1.0);
}`;

const DEPOSIT_FS = `#version 300 es
precision highp float;
uniform float uDep;
out vec4 o;
void main(){ o = vec4(uDep, 0.0, 0.0, 0.0); }`;

const FOOD_VS = `#version 300 es
precision highp float;
in vec3 aCity;
uniform vec2 uGrid;
uniform float uRadius;
out float vW;
void main(){
  gl_Position = vec4((aCity.xy / uGrid) * 2.0 - 1.0, 0.0, 1.0);
  vW = aCity.z;
  // bigger towns are bigger flakes of oats
  gl_PointSize = uRadius * (0.6 + 0.8 * vW) * 2.0;
}`;

const FOOD_FS = `#version 300 es
precision highp float;
uniform float uFood;
in float vW;
out vec4 o;
void main(){
  vec2 d2 = gl_PointCoord - 0.5;
  float d = dot(d2, d2) * 4.0;
  o = vec4(uFood * (0.5 + 0.5 * vW) * max(0.0, 1.0 - d), 0.0, 0.0, 0.0);
}`;

const DIFFUSE_FS = `#version 300 es
precision highp float;
uniform highp sampler2D uTrail;
uniform highp sampler2D uMask;
uniform vec2 uGrid;
uniform float uDecay, uMax;
out vec4 o;
void main(){
  ivec2 c = ivec2(gl_FragCoord.xy);
  ivec2 hi = ivec2(uGrid) - 1;
  float s = 0.0;
  for (int j = -1; j <= 1; j++)
    for (int i = -1; i <= 1; i++)
      s += texelFetch(uTrail, clamp(c + ivec2(i, j), ivec2(0), hi), 0).r;
  vec2 m = texelFetch(uMask, c, 0).rg;
  float pass = step(0.5, max(m.r, m.g));
  // the trail saturates, like the 8-bit map in Jones (2010): a vein cannot outshine a town
  o = vec4(min(s / 9.0 * uDecay, uMax) * pass, 0.0, 0.0, 0.0);
}`;

const RENDER_FS = `#version 300 es
precision highp float;
uniform highp sampler2D uTrail;
uniform highp sampler2D uMask;
uniform vec2 uGrid, uCanvas, uOffset;
uniform float uScale, uExposure;
uniform vec3 uSea, uLand, uCoast, uLo, uHi;
out vec4 o;
float T(ivec2 q){ return texelFetch(uTrail, clamp(q, ivec2(0), ivec2(uGrid) - 1), 0).r; }
void main(){
  vec2 sp = vec2(gl_FragCoord.x, uCanvas.y - gl_FragCoord.y);
  vec2 g = sp / uScale + uOffset;
  if (g.x < 0.0 || g.y < 0.0 || g.x >= uGrid.x || g.y >= uGrid.y) { o = vec4(uSea, 1.0); return; }
  vec2 gc = g - 0.5;
  ivec2 c = ivec2(floor(gc));
  vec2 f = fract(gc);
  float t = mix(mix(T(c), T(c + ivec2(1, 0)), f.x), mix(T(c + ivec2(0, 1)), T(c + ivec2(1, 1)), f.x), f.y);
  vec3 m = texelFetch(uMask, ivec2(g), 0).rgb;
  // soft coastline: where land coverage is fractional, lighten a touch
  float edge = m.r * (1.0 - m.r) * 4.0;
  vec3 base = mix(uSea, uLand, m.r) + uCoast * edge + uCoast * m.g * (1.0 - m.r) * 0.18;
  float e = 1.0 - exp(-t * uExposure);
  vec3 trail = mix(uLo, uHi, e * e) * e;
  o = vec4(base + trail, 1.0);
}`;

function compile(gl, type, src) {
  const sh = gl.createShader(type);
  gl.shaderSource(sh, src);
  gl.compileShader(sh);
  if (!gl.getShaderParameter(sh, gl.COMPILE_STATUS)) {
    throw new Error('Shader error: ' + gl.getShaderInfoLog(sh) + '\n' + src);
  }
  return sh;
}
function program(gl, vs, fs) {
  const p = gl.createProgram();
  gl.attachShader(p, compile(gl, gl.VERTEX_SHADER, vs));
  gl.attachShader(p, compile(gl, gl.FRAGMENT_SHADER, fs));
  gl.linkProgram(p);
  if (!gl.getProgramParameter(p, gl.LINK_STATUS)) throw new Error('Link error: ' + gl.getProgramInfoLog(p));
  const u = {};
  const n = gl.getProgramParameter(p, gl.ACTIVE_UNIFORMS);
  for (let i = 0; i < n; i++) { const info = gl.getActiveUniform(p, i); u[info.name] = gl.getUniformLocation(p, info.name); }
  return { p, u };
}
function tex2d(gl, w, h, internal, format, type, data) {
  const t = gl.createTexture();
  gl.bindTexture(gl.TEXTURE_2D, t);
  gl.texImage2D(gl.TEXTURE_2D, 0, internal, w, h, 0, format, type, data || null);
  gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
  gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
  gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
  gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
  return t;
}
function fbo(gl, tex) {
  const f = gl.createFramebuffer();
  gl.bindFramebuffer(gl.FRAMEBUFFER, f);
  gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, tex, 0);
  const st = gl.checkFramebufferStatus(gl.FRAMEBUFFER);
  if (st !== gl.FRAMEBUFFER_COMPLETE) throw new Error('Framebuffer incomplete: 0x' + st.toString(16));
  gl.bindFramebuffer(gl.FRAMEBUFFER, null);
  return f;
}
function hex(c) { return [((c >> 16) & 255) / 255, ((c >> 8) & 255) / 255, (c & 255) / 255]; }

export const AGENT_PRESETS = {
  '65k':  [256, 256],
  '262k': [512, 512],
  '1M':   [1024, 1024],
  '2M':   [2048, 1024],
};

export class Sim {
  constructor(canvas, mask, opts = {}) {
    const gl = canvas.getContext('webgl2', { antialias: false, alpha: false, depth: false, stencil: false, preserveDrawingBuffer: true, powerPreference: 'high-performance' });
    if (!gl) throw new Error('This browser has no WebGL2. Try Chrome, Firefox, Edge or Safari 15+.');
    if (!gl.getExtension('EXT_color_buffer_float')) throw new Error('This GPU cannot render to float textures (EXT_color_buffer_float).');
    this.gl = gl;
    this.canvas = canvas;
    this.W = opts.W; this.H = opts.H;
    this.mask = mask;              // { data: Uint8ClampedArray (RGBA, W*H), image }
    this.p = {
      sensorAngle: 0.5, sensorDist: 9, turnAngle: 0.45, stepSize: 1.0,
      deposit: 0.02,     // per agent per step at 262k agents; scaled so total trail mass is constant
      decay: 0.94, food: 1.0, cityRadius: 4, exposure: 3.4,
      trailMax: 1.0,     // saturation level of the trail; towns emit enough to pin it there
      ferryRate: 0.03,   // probability that an agent at the shore boards a corridor
      randomTurn: 0.003, // probability per step of forgetting the trail and turning at random
    };
    Object.assign(this.p, opts.params || {});
    this.colors = {
      sea: hex(0x0a0e1a), land: hex(0x1a1f2c), coast: hex(0x2c3446),
      lo: hex(0xd9772a), hi: hex(0xfff3c4),
    };
    this.step = 0;
    this.agentCount = 0;
    this.cityCount = 0;
    this._build();
  }

  _build() {
    const gl = this.gl;
    this.prog = {
      update: program(gl, QUAD_VS, UPDATE_FS),
      deposit: program(gl, DEPOSIT_VS, DEPOSIT_FS),
      food: program(gl, FOOD_VS, FOOD_FS),
      diffuse: program(gl, QUAD_VS, DIFFUSE_FS),
      render: program(gl, QUAD_VS, RENDER_FS),
    };
    // mask texture straight from the image: row 0 = north
    this.maskTex = gl.createTexture();
    gl.bindTexture(gl.TEXTURE_2D, this.maskTex);
    gl.pixelStorei(gl.UNPACK_FLIP_Y_WEBGL, false);
    gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA8, gl.RGBA, gl.UNSIGNED_BYTE, this.mask.image);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
    // trail ping-pong
    this.trail = [0, 1].map(() => {
      const t = tex2d(gl, this.W, this.H, gl.R16F, gl.RED, gl.HALF_FLOAT, null);
      return { t, f: fbo(gl, t) };
    });
    this.cur = 0;
    this.cityBuf = gl.createBuffer();
    this.cityVao = gl.createVertexArray();
    gl.bindVertexArray(this.cityVao);
    gl.bindBuffer(gl.ARRAY_BUFFER, this.cityBuf);
    const loc = gl.getAttribLocation(this.prog.food.p, 'aCity');
    gl.enableVertexAttribArray(loc);
    gl.vertexAttribPointer(loc, 3, gl.FLOAT, false, 0, 0);
    gl.bindVertexArray(null);
    this.emptyVao = gl.createVertexArray();
  }

  // cities: array of {x, y, w} in grid px, w in [0,1]
  setCities(cities) {
    const gl = this.gl;
    const arr = new Float32Array(cities.length * 3);
    cities.forEach((c, i) => { arr[i * 3] = c.x; arr[i * 3 + 1] = c.y; arr[i * 3 + 2] = c.w; });
    gl.bindBuffer(gl.ARRAY_BUFFER, this.cityBuf);
    gl.bufferData(gl.ARRAY_BUFFER, arr, gl.STATIC_DRAW);
    this.cityCount = cities.length;
  }

  // Seeds the population. mode: 'everywhere' | {x, y, spread, release}
  // 'release' spreads the agents' entry into the colony over that many steps, so it grows.
  seed(preset, mode) {
    const gl = this.gl;
    const [AW, AH] = AGENT_PRESETS[preset] || AGENT_PRESETS['1M'];
    const n = AW * AH;
    const data = new Float32Array(n * 4);
    const md = this.mask.data, W = this.W, H = this.H;
    const land = (x, y) => {
      const i = ((y | 0) * W + (x | 0)) * 4;
      return Math.max(md[i], md[i + 1]) >= 128;
    };
    // land cell list for uniform spawning
    if (!this._landCells) {
      const cells = [];
      for (let y = 0; y < H; y++) for (let x = 0; x < W; x++) if (md[(y * W + x) * 4] >= 128) cells.push(y * W + x);
      this._landCells = Uint32Array.from(cells);
    }
    const lc = this._landCells;
    for (let i = 0; i < n; i++) {
      let x, y;
      if (mode === 'everywhere') {
        const c = lc[(Math.random() * lc.length) | 0];
        x = (c % W) + Math.random(); y = Math.floor(c / W) + Math.random();
      } else {
        let tries = 0;
        do {
          const r = mode.spread * Math.sqrt(-2 * Math.log(1 - Math.random())), a = Math.random() * 6.2831853;
          x = mode.x + r * Math.cos(a); y = mode.y + r * Math.sin(a);
        } while ((x < 0 || y < 0 || x >= W || y >= H || !land(x, y)) && ++tries < 50);
        if (tries >= 50) { x = mode.x; y = mode.y; }
      }
      data[i * 4] = x; data[i * 4 + 1] = y; data[i * 4 + 2] = Math.random() * 6.2831853;
      data[i * 4 + 3] = mode.release ? (i / n) * mode.release : 0;
    }
    this.release = (mode && mode.release) || 0;
    if (this.agents) this.agents.forEach(a => { gl.deleteTexture(a.t); gl.deleteFramebuffer(a.f); });
    this.agents = [0, 1].map((k) => {
      const t = tex2d(gl, AW, AH, gl.RGBA32F, gl.RGBA, gl.FLOAT, k === 0 ? data : null);
      return { t, f: fbo(gl, t) };
    });
    this.AW = AW; this.AH = AH; this.agentCount = n;
    this.clearTrail();
    this.step = 0;
  }

  clearTrail() {
    const gl = this.gl;
    for (const tr of this.trail) {
      gl.bindFramebuffer(gl.FRAMEBUFFER, tr.f);
      gl.clearColor(0, 0, 0, 0);
      gl.clear(gl.COLOR_BUFFER_BIT);
    }
    gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    this.step = 0;
  }

  tick() {
    const gl = this.gl, p = this.p;
    if (!this.agents) return;
    const src = this.agents[0], dst = this.agents[1];
    const tr = this.trail[this.cur], nx = this.trail[this.cur ^ 1];
    gl.disable(gl.BLEND);
    gl.bindVertexArray(this.emptyVao);

    // 1. sense + move
    let pr = this.prog.update;
    gl.useProgram(pr.p);
    gl.bindFramebuffer(gl.FRAMEBUFFER, dst.f);
    gl.viewport(0, 0, this.AW, this.AH);
    gl.activeTexture(gl.TEXTURE0); gl.bindTexture(gl.TEXTURE_2D, src.t); gl.uniform1i(pr.u.uAgents, 0);
    gl.activeTexture(gl.TEXTURE1); gl.bindTexture(gl.TEXTURE_2D, tr.t); gl.uniform1i(pr.u.uTrail, 1);
    gl.activeTexture(gl.TEXTURE2); gl.bindTexture(gl.TEXTURE_2D, this.maskTex); gl.uniform1i(pr.u.uMask, 2);
    gl.uniform2f(pr.u.uGrid, this.W, this.H);
    gl.uniform1f(pr.u.uSA, p.sensorAngle);
    gl.uniform1f(pr.u.uSD, p.sensorDist);
    gl.uniform1f(pr.u.uRA, p.turnAngle);
    gl.uniform1f(pr.u.uSS, p.stepSize);
    gl.uniform1f(pr.u.uSeed, (this.step % 1000) * 0.6180339 + 0.137);
    gl.uniform1f(pr.u.uStep, this.step);
    gl.uniform1f(pr.u.uFerry, p.ferryRate);
    gl.uniform1f(pr.u.uRand, p.randomTurn);
    gl.drawArrays(gl.TRIANGLES, 0, 3);
    this.agents = [dst, src];

    // 2. deposit (additive)
    gl.bindFramebuffer(gl.FRAMEBUFFER, tr.f);
    gl.viewport(0, 0, this.W, this.H);
    gl.enable(gl.BLEND);
    gl.blendFunc(gl.ONE, gl.ONE);
    pr = this.prog.deposit;
    gl.useProgram(pr.p);
    gl.activeTexture(gl.TEXTURE0); gl.bindTexture(gl.TEXTURE_2D, dst.t); gl.uniform1i(pr.u.uAgents, 0);
    gl.uniform1i(pr.u.uAW, this.AW);
    gl.uniform2f(pr.u.uGrid, this.W, this.H);
    gl.uniform1f(pr.u.uStep, this.step);
    gl.uniform1f(pr.u.uDep, p.deposit * (262144 / this.agentCount));
    gl.drawArrays(gl.POINTS, 0, this.agentCount);

    // 3. food from cities
    if (this.cityCount) {
      pr = this.prog.food;
      gl.useProgram(pr.p);
      gl.bindVertexArray(this.cityVao);
      gl.uniform2f(pr.u.uGrid, this.W, this.H);
      gl.uniform1f(pr.u.uRadius, p.cityRadius);
      gl.uniform1f(pr.u.uFood, p.food);
      gl.drawArrays(gl.POINTS, 0, this.cityCount);
      gl.bindVertexArray(this.emptyVao);
    }
    gl.disable(gl.BLEND);

    // 4. diffuse + decay
    pr = this.prog.diffuse;
    gl.useProgram(pr.p);
    gl.bindFramebuffer(gl.FRAMEBUFFER, nx.f);
    gl.activeTexture(gl.TEXTURE0); gl.bindTexture(gl.TEXTURE_2D, tr.t); gl.uniform1i(pr.u.uTrail, 0);
    gl.activeTexture(gl.TEXTURE1); gl.bindTexture(gl.TEXTURE_2D, this.maskTex); gl.uniform1i(pr.u.uMask, 1);
    gl.uniform2f(pr.u.uGrid, this.W, this.H);
    gl.uniform1f(pr.u.uDecay, p.decay);
    gl.uniform1f(pr.u.uMax, p.trailMax);
    gl.drawArrays(gl.TRIANGLES, 0, 3);
    this.cur ^= 1;
    this.step++;
  }

  // view: { scale: canvas px per grid px, ox, oy: grid coords at the canvas top-left }
  render(view) {
    const gl = this.gl, c = this.colors, pr = this.prog.render;
    gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    gl.viewport(0, 0, this.canvas.width, this.canvas.height);
    gl.disable(gl.BLEND);
    gl.bindVertexArray(this.emptyVao);
    gl.useProgram(pr.p);
    gl.activeTexture(gl.TEXTURE0); gl.bindTexture(gl.TEXTURE_2D, this.trail[this.cur].t); gl.uniform1i(pr.u.uTrail, 0);
    gl.activeTexture(gl.TEXTURE1); gl.bindTexture(gl.TEXTURE_2D, this.maskTex); gl.uniform1i(pr.u.uMask, 1);
    gl.uniform2f(pr.u.uGrid, this.W, this.H);
    gl.uniform2f(pr.u.uCanvas, this.canvas.width, this.canvas.height);
    gl.uniform2f(pr.u.uOffset, view.ox, view.oy);
    gl.uniform1f(pr.u.uScale, view.scale);
    gl.uniform1f(pr.u.uExposure, this.p.exposure);
    gl.uniform3fv(pr.u.uSea, c.sea);
    gl.uniform3fv(pr.u.uLand, c.land);
    gl.uniform3fv(pr.u.uCoast, c.coast);
    gl.uniform3fv(pr.u.uLo, c.lo);
    gl.uniform3fv(pr.u.uHi, c.hi);
    gl.drawArrays(gl.TRIANGLES, 0, 3);
  }
}