// In-browser port of brain.py + the step() of server.py: leaky integrate-and-fire network
// over the FlyWire v783 connectome (Shiu et al. 2024 constants), 1 ms steps.
const V0 = -52, VTH = -45, T_MBR = 20, TAU = 5, DT = 1, REFRAC = 2;
const DECAY = Math.exp(-DT / TAU), K_V = DT / T_MBR;

let meta, N, indptr, idx, w, c;           // connectome (shared, read-only)
let v, g, rfc, ring, tick = 0;            // state of this simulation
let groupBits, groupNames, isDN, dnPos, recent, drive = [], spkFlag, forcedBuf;
// mushroom-body learning (real KC -> MBON synapses, depressed when a dopamine neuron is co-active)
let mb = null, kcOrd, mbonOrd, danOrd, plOff, plK, plPre, plPost, mult, kcTrace, danTrace, mbCnt;
const ETA = 0.04;
let RECOVER = 0.00025, clock = 1;               // memory fades on the fly-day clock (set by the page)

async function gunzip(buf) {
  if (buf[0] === 0x1f && buf[1] === 0x8b) {
    const s = new Blob([buf]).stream().pipeThrough(new DecompressionStream('gzip'));
    return await new Response(s).arrayBuffer();
  }
  return buf.buffer;
}

async function load(url, rawBytes) {
  const res = await fetch(url);
  const len = +res.headers.get('Content-Length') || 0;
  const reader = res.body.getReader();
  const parts = []; let got = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    parts.push(value); got += value.length;
    if (len) postMessage({ type: 'progress', file: url, frac: got / len });
  }
  const buf = new Uint8Array(got); let o = 0;
  for (const p of parts) { buf.set(p, o); o += p.length; }
  if (rawBytes) return buf.buffer;                         // a part of a larger gzip stream
  return await gunzip(buf);
}

async function init(ver) {
  meta = await (await fetch('meta.json?v=' + ver)).json();
  N = meta.n; c = meta.wsyn * meta.gain;
  postMessage({ type: 'status', text: 'downloading connectome (~37 MB)…' });
  indptr = new Uint32Array(await load('indptr.dat'));
  const parts = [];                                           // delta.dat is split in 4 parts of one gzip stream
  for (let i = 0; i < 4; i++) parts.push(new Uint8Array(await load(`delta${i}.dat`, true)));
  let tot = 0; for (const p of parts) tot += p.length;
  const joined = new Uint8Array(tot); { let o = 0; for (const p of parts) { joined.set(p, o); o += p.length; } }
  idx = new Uint32Array(await gunzip(joined));
  w = new Int16Array(await load('weight.dat'));
  postMessage({ type: 'status', text: 'decoding…' });
  for (let r = 0; r < N; r++) {                              // undo the per-row delta coding
    let acc = 0;
    for (let k = indptr[r], e = indptr[r + 1]; k < e; k++) { acc = k === indptr[r] ? idx[k] : acc + idx[k]; idx[k] = acc; }
  }
  groupNames = Object.keys(meta.groups);
  groupBits = new Uint8Array(N);
  groupNames.forEach((gn, b) => { for (const i of meta.groups[gn]) groupBits[i] |= 1 << b; });
  isDN = new Uint8Array(N); dnPos = new Int32Array(N).fill(-1);
  meta.dn.forEach((i, p) => { isDN[i] = 1; dnPos[i] = p; });
  recent = new Float32Array(meta.dn.length);
  spkFlag = new Uint8Array(N);
  if (meta.mb) setupLearning();
  reset();
  postMessage({ type: 'ready', info: { neurons: N, synapses: meta.synapses, connections: meta.connections,
                                       pos: meta.pos, dn: meta.dn } });
}

function setupLearning() {
  mb = meta.mb;
  kcOrd = new Int32Array(N).fill(-1); mbonOrd = new Int32Array(N).fill(-1); danOrd = new Int32Array(N).fill(-1);
  mb.kc.forEach((i, o) => { kcOrd[i] = o; });
  mb.mbon.forEach((i, o) => { mbonOrd[i] = o; });
  mb.pam.forEach((i, o) => { danOrd[i] = o; });
  mb.ppl.forEach((i, o) => { danOrd[i] = mb.pam.length + o; });
  const K = [], pre = [], post = []; plOff = new Int32Array(mb.kc.length + 1);
  mb.kc.forEach((i, o) => {
    plOff[o] = K.length;
    for (let k = indptr[i], e = indptr[i + 1]; k < e; k++) {
      const m = mbonOrd[idx[k]];
      if (m >= 0) { K.push(k); pre.push(i); post.push(m); }
    }
  });
  plOff[mb.kc.length] = K.length;
  plK = Int32Array.from(K); plPre = Int32Array.from(pre); plPost = Int32Array.from(post);
  mult = new Float32Array(K.length).fill(1);
  kcTrace = new Float32Array(mb.kc.length);
  danTrace = new Float32Array(mb.pam.length + mb.ppl.length);
  mbCnt = new Float32Array(mb.mbon.length);
}

function reset() {
  v = new Float32Array(N).fill(V0); g = new Float32Array(N); rfc = new Uint8Array(N);
  ring = [new Float32Array(N), new Float32Array(N)]; tick = 0;
  if (recent) recent.fill(0);
  if (mb) { mult.fill(1); kcTrace.fill(0); danTrace.fill(0); mbCnt.fill(0); }
}

// one 1 ms step; returns the list of neurons that spiked
function step() {
  const cur = ring[tick++ & 1];
  // forced (sensory) spikes, decided on the refractory state before this step
  let nf = 0;
  for (let j = 0; j < drive.length; j++) {
    const d = drive[j], ids = d.ids, p = d.p;
    for (let k = 0; k < ids.length; k++) {
      const i = ids[k];
      if (rfc[i] <= 0 && Math.random() < p && !spkFlag[i]) { spkFlag[i] = 1; forcedBuf[nf++] = i; }
    }
  }
  const spikes = [];
  for (let i = 0; i < N; i++) {
    const gi = g[i];
    const vi = v[i] + (V0 - v[i] + gi) * K_V;
    g[i] = gi * DECAY + cur[i];
    cur[i] = 0;
    v[i] = vi;
    const r = rfc[i];
    if (spkFlag[i] || (vi > VTH && r <= 0)) spikes.push(i);
    if (r > 0) rfc[i] = r - 1;
  }
  for (let k = 0; k < nf; k++) spkFlag[forcedBuf[k]] = 0;
  for (let s = 0; s < spikes.length; s++) {
    const i = spikes[s];
    v[i] = V0; g[i] = 0; rfc[i] = REFRAC;
    for (let k = indptr[i], e = indptr[i + 1]; k < e; k++) cur[idx[k]] += w[k] * c;
    if (mb) {
      const o = kcOrd[i];
      if (o >= 0) {
        kcTrace[o] += 1;
        for (let p = plOff[o], pe = plOff[o + 1]; p < pe; p++) cur[idx[plK[p]]] += w[plK[p]] * c * (mult[p] - 1);
      } else {
        const d = danOrd[i]; if (d >= 0) danTrace[d] += 1;
        const m = mbonOrd[i]; if (m >= 0) mbCnt[m] += 1;
      }
    }
  }
  return spikes;
}

function runStep(stim) {
  const t0 = performance.now();
  drive = []; let nd = 0;
  for (const [name, list] of Object.entries(meta.channels)) {
    const s = Math.min(1, Math.max(0, +stim[name] || 0));
    if (s > 0.001 && list.length) { drive.push({ ids: list, p: s * meta.max_rate }); nd += list.length; }
  }
  if (!forcedBuf || forcedBuf.length < nd) forcedBuf = new Int32Array(Math.max(nd, 1024));
  const G = groupNames.length, cnt = new Array(G).fill(0);
  const steps = meta.steps_per_call;
  let total = 0; const shown = [];
  for (let s = 0; s < steps; s++) {
    const sp = step();
    total += sp.length;
    for (let k = 0; k < sp.length; k++) {
      const i = sp[k], b = groupBits[i];
      if (b) for (let q = 0; q < G; q++) if (b & (1 << q)) cnt[q]++;
      if (isDN[i]) recent[dnPos[i]]++;
    }
    if (sp.length <= 120) for (let k = 0; k < sp.length; k++) shown.push(sp[k]);
    else for (let k = 0; k < 120; k++) shown.push(sp[(Math.random() * sp.length) | 0]);
  }
  for (let p = 0; p < recent.length; p++) recent[p] *= 0.9;
  let mbA = 0, mbV = 0, depressed = 0;
  if (mb) {
    const nP = mb.pam.length, nM = mb.mbon.length;
    // dopamine drive onto each MBON = share of its own dopamine input that is active
    const mod = new Float32Array(nM);
    for (let m = 0; m < nM; m++) {
      let num = 0, tot = 0;
      for (let d = 0; d < nP; d++) { const wgt = mb.Mp[d][m]; tot += wgt; num += Math.min(1, danTrace[d] / 3) * wgt; }
      for (let d = 0; d < mb.ppl.length; d++) { const wgt = mb.Ml[d][m]; tot += wgt; num += Math.min(1, danTrace[nP + d] / 3) * wgt; }
      mod[m] = Math.min(1, Math.max(0, num / (tot + 1e-9) - 0.6) / 0.4);
    }
    for (let p = 0; p < mult.length; p++) {
      const kt = Math.min(1, kcTrace[kcOrd[plPre[p]]] / 2), dm = mod[plPost[p]];
      let v = mult[p];
      if (kt > 0 && dm > 0) v *= 1 - ETA * kt * dm;
      v += (1 - v) * RECOVER * clock;
      mult[p] = v < 0.1 ? 0.1 : v;
      if (v < 0.9) depressed++;
    }
    for (let m = 0; m < nM; m++) { if (mb.avoid[m]) mbV += mbCnt[m]; else mbA += mbCnt[m]; }
    kcTrace.forEach((_, i) => { kcTrace[i] *= 0.5; });
    for (let d = 0; d < danTrace.length; d++) danTrace[d] *= 0.6;
    mbCnt.fill(0);
  }
  while (shown.length > 1500) shown.splice((Math.random() * shown.length) | 0, 1);
  const out = {};
  groupNames.forEach((gn, q) => { out[gn] = Math.min(1, cnt[q] / steps / meta.ref[gn]); });
  const order = Array.from(recent.keys()).sort((a, b) => recent[b] - recent[a]).slice(0, 5);
  out.spikes_per_ms = total / steps;
  out.spiking = shown;
  out.top_dn = order.filter(p => recent[p] > 0.5).map(p => [meta.dn_names[p], Math.round(recent[p] * 10) / 10]);
  out.mb_approach = mbA; out.mb_avoid = mbV;
  out.mb_base_approach = mb ? mb.base_approach * steps / 500 : 0;
  out.mb_base_avoid = mb ? mb.base_avoid * steps / 500 : 0;
  out.depressed = mb ? depressed / Math.max(1, mult.length) : 0;
  out.compute_ms = Math.round(performance.now() - t0);
  return out;
}

onmessage = (e) => {
  const m = e.data;
  if (m.type === 'init') init(m.v).catch(err => postMessage({ type: 'error', text: String(err) }));
  else if (m.type === 'clock') clock = Math.max(0.1, +m.speed || 1);
  else if (m.type === 'reset') reset();
  else if (m.type === 'getMemory') {
    const q = new Uint8Array(mult ? mult.length : 0);
    for (let p = 0; p < q.length; p++) q[p] = Math.round(mult[p] * 255);
    postMessage({ type: 'memory', data: q }, [q.buffer]);
  } else if (m.type === 'setMemory' && mult && m.data && m.data.length === mult.length) {
    for (let p = 0; p < mult.length; p++) mult[p] = Math.max(0.1, m.data[p] / 255);
  } else if (m.type === 'forget' && mult) mult.fill(1);
  else if (m.type === 'step') postMessage({ type: 'out', id: m.id, out: runStep(m.stim || {}) });
};