// tracker.js — лёгкий корреляционный WGSL-трекер боксов между кейфреймами.
//
// Схема: на кейфрейме из каждого бокса ресемплируется grayscale-шаблон
// 24x24 (zero-mean нормировка); на каждом кадре для бокса перебираются
// смещения в окне +-R пикселей (сетка (2R/STEP+1)^2), скор — кросс-корреляция
// нормированного шаблона с сырыми пикселями (инвариант к сдвигу яркости).
// Один dispatch на кадр: workgroup на бокс, 256 потоков делят смещения,
// argmax — в shared. Ридбек: nBoxes * (dx,dy,score).
//
// Стоимость: box * 1089 смещений * 576 пикс ~ 0.6 MFLOP/бокс — микросекунды.

const TPL = 24;          // сторона шаблона
const R = 16;            // радиус поиска, px кадра
const STEP = 1;          // шаг сетки смещений
const GRID = 2 * R / STEP + 1;              // 33
const NSHIFT = GRID * GRID;                 // 1089
const MAXB = 64;         // максимум треков

const WGSL = /* wgsl */ `
struct Params { n_boxes : u32, fw : u32, fh : u32, p3 : u32 }
@group(0) @binding(0) var<uniform> params : Params;
@group(0) @binding(1) var<storage, read> frame : array<f32>;     // [fh, fw] gray
@group(0) @binding(2) var<storage, read> tpl : array<f32>;       // [MAXB, 24*24] zero-mean
// boxinfo: [MAXB] vec4<f32> = (x0, y0, sx, sy) — угол окна поиска в px кадра
// и шаг сэмплирования шаблона (bw/24, bh/24)
@group(0) @binding(3) var<storage, read> boxinfo : array<vec4<f32>>;
@group(0) @binding(4) var<storage, read_write> out : array<vec4<f32>>; // (dx, dy, score, 0)

const TPL : u32 = ${TPL}u;
const GRID : i32 = ${GRID};
const STEP : i32 = ${STEP};
const R : i32 = ${R};
const NSHIFT : u32 = ${NSHIFT}u;

var<workgroup> best : array<f32, 256>;
var<workgroup> bestI : array<u32, 256>;

fn sample(x : f32, y : f32) -> f32 {
  // билинейно, с клампом к краям кадра
  let fx = clamp(x, 0.0, f32(params.fw - 1u));
  let fy = clamp(y, 0.0, f32(params.fh - 1u));
  let x0 = u32(fx); let y0 = u32(fy);
  let x1 = min(x0 + 1u, params.fw - 1u);
  let y1 = min(y0 + 1u, params.fh - 1u);
  let ax = fx - f32(x0); let ay = fy - f32(y0);
  let v00 = frame[y0 * params.fw + x0];
  let v01 = frame[y0 * params.fw + x1];
  let v10 = frame[y1 * params.fw + x0];
  let v11 = frame[y1 * params.fw + x1];
  return mix(mix(v00, v01, ax), mix(v10, v11, ax), ay);
}

@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wg_id : vec3<u32>,
        @builtin(local_invocation_index) lid : u32) {
  let b = wg_id.x;
  let info = boxinfo[b];
  var myBest = -1.0e30;
  var myIdx = 0u;
  for (var s = lid; s < NSHIFT; s = s + 256u) {
    let dy = f32(i32(s / u32(GRID)) * STEP - R);
    let dx = f32(i32(s % u32(GRID)) * STEP - R);
    var score = 0.0;
    for (var i = 0u; i < TPL; i = i + 1u) {
      let y = info.y + dy + (f32(i) + 0.5) * info.w;
      for (var j = 0u; j < TPL; j = j + 1u) {
        let x = info.x + dx + (f32(j) + 0.5) * info.z;
        score = score + tpl[b * (TPL * TPL) + i * TPL + j] * sample(x, y);
      }
    }
    if (score > myBest) { myBest = score; myIdx = s; }
  }
  best[lid] = myBest; bestI[lid] = myIdx;
  workgroupBarrier();
  var stride = 128u;
  loop {
    if (stride == 0u) { break; }
    if (lid < stride) {
      if (best[lid + stride] > best[lid]) {
        best[lid] = best[lid + stride]; bestI[lid] = bestI[lid + stride];
      }
    }
    workgroupBarrier();
    stride = stride >> 1u;
  }
  if (lid == 0u) {
    let s = bestI[0];
    let dy = f32(i32(s / u32(GRID)) * STEP - R);
    let dx = f32(i32(s % u32(GRID)) * STEP - R);
    out[b] = vec4<f32>(dx, dy, best[0], 0.0);
  }
}
`;

export class Tracker {
  constructor(device) {
    this.device = device;
    this.boxes = [];       // [{box:{x1..y2 норм.}, det}] — текущее состояние
    this.fw = 0; this.fh = 0;
  }

  async init() {
    const d = this.device;
    this.pipe = await d.createComputePipelineAsync({
      layout: 'auto',
      compute: { module: d.createShaderModule({ code: WGSL }), entryPoint: 'main' },
    });
    this.u = d.createBuffer({ size: 16, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
    this.tplBuf = d.createBuffer({ size: MAXB * TPL * TPL * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
    this.infoBuf = d.createBuffer({ size: MAXB * 16, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
    this.outBuf = d.createBuffer({ size: MAXB * 16, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC });
    this.frameBuf = null;
    return this;
  }

  _ensureFrame(fw, fh) {
    if (this.frameBuf && this.fw === fw && this.fh === fh) return;
    this.frameBuf?.destroy();
    this.fw = fw; this.fh = fh;
    this.frameBuf = this.device.createBuffer({
      size: fw * fh * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
    this.bind = this.device.createBindGroup({
      layout: this.pipe.getBindGroupLayout(0),
      entries: [this.u, this.frameBuf, this.tplBuf, this.infoBuf, this.outBuf]
        .map((buffer, i) => ({ binding: i, resource: { buffer } })),
    });
  }

  uploadFrame(gray, fw, fh) {
    this._ensureFrame(fw, fh);
    this.device.queue.writeBuffer(this.frameBuf, 0, gray);
    this.device.queue.writeBuffer(this.u, 0, new Uint32Array([this.boxes.length, fw, fh, 0]));
  }

  // Кейфрейм: детекции + grayscale кадра -> шаблоны и опорные боксы.
  // ВАЖНО: шаблоны режутся по dets (позиции В КАДРЕ gray); стартовое состояние
  // треков может быть сдвинуто (stateBoxes) — пан, накопленный за время детекции.
  setKeyframe(dets, gray, fw, fh, stateBoxes = null) {
    const src = dets.slice(0, MAXB);
    this.boxes = src.map((d, i) => ({ det: d,
      box: { ...((stateBoxes && stateBoxes[i]) || d.box) } }));
    this.uploadFrame(gray, fw, fh);
    const tpl = new Float32Array(MAXB * TPL * TPL);
    for (let b = 0; b < src.length; b++) {
      const bx = src[b].box;
      const x0 = bx.x1 * fw, y0 = bx.y1 * fh;
      const sx = (bx.x2 - bx.x1) * fw / TPL, sy = (bx.y2 - bx.y1) * fh / TPL;
      const t = new Float32Array(TPL * TPL);
      let mean = 0;
      for (let i = 0; i < TPL; i++)
        for (let j = 0; j < TPL; j++) {
          const y = Math.min(fh - 1, Math.max(0, Math.round(y0 + (i + 0.5) * sy)));
          const x = Math.min(fw - 1, Math.max(0, Math.round(x0 + (j + 0.5) * sx)));
          const v = gray[y * fw + x];
          t[i * TPL + j] = v; mean += v;
        }
      mean /= TPL * TPL;
      for (let k = 0; k < TPL * TPL; k++) tpl[b * TPL * TPL + k] = t[k] - mean;
    }
    this.device.queue.writeBuffer(this.tplBuf, 0, tpl);
  }

  // Трек кадра: обновляет боксы, возвращает детекции с новыми координатами.
  async track(gray, fw, fh) {
    if (!this.boxes.length) return [];
    this.uploadFrame(gray, fw, fh);
    const info = new Float32Array(MAXB * 4);
    for (let b = 0; b < this.boxes.length; b++) {
      const bx = this.boxes[b].box;
      info[b * 4] = bx.x1 * fw;                       // x угла (без смещения)
      info[b * 4 + 1] = bx.y1 * fh;                   // y
      info[b * 4 + 2] = (bx.x2 - bx.x1) * fw / TPL;   // sx
      info[b * 4 + 3] = (bx.y2 - bx.y1) * fh / TPL;   // sy
    }
    this.device.queue.writeBuffer(this.infoBuf, 0, info);
    const enc = this.device.createCommandEncoder();
    const pass = enc.beginComputePass();
    pass.setPipeline(this.pipe);
    pass.setBindGroup(0, this.bind);
    pass.dispatchWorkgroups(this.boxes.length);
    pass.end();
    const st = this.device.createBuffer({
      size: this.boxes.length * 16, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
    enc.copyBufferToBuffer(this.outBuf, 0, st, 0, this.boxes.length * 16);
    this.device.queue.submit([enc.finish()]);
    await st.mapAsync(GPUMapMode.READ);
    const res = new Float32Array(st.getMappedRange().slice(0));
    st.unmap(); st.destroy();
    const dxs = [], dys = [];
    const out = this.boxes.map((t, b) => {
      const dx = res[b * 4] / fw, dy = res[b * 4 + 1] / fh;
      dxs.push(res[b * 4]); dys.push(res[b * 4 + 1]);
      t.box.x1 += dx; t.box.x2 += dx;
      t.box.y1 += dy; t.box.y2 += dy;
      return { ...t.det, box: { ...t.box } };
    });
    const med = (a) => a.slice().sort((x, y) => x - y)[a.length >> 1];
    this.lastShift = { dx: med(dxs) / fw, dy: med(dys) / fh }; // медианный пан кадра
    return out;
  }
}

export function toGray(imageData, fw, fh) {
  const g = new Float32Array(fw * fh);
  const d = imageData.data;
  for (let i = 0; i < fw * fh; i++) {
    g[i] = (0.299 * d[i * 4] + 0.587 * d[i * 4 + 1] + 0.114 * d[i * 4 + 2]) / 255;
  }
  return g;
}