#version 300 es
precision highp float;
precision highp int;
uniform sampler2D u_in, u_w, u_res;
uniform int u_has_res;
uniform int u_mode, u_in_len, u_out_len, u_row_offset, u_elem_base;
uniform float u_f0, u_f1;
uniform int u_i0, u_i1, u_i2, u_i3, u_i4, u_i5;
out vec4 fragColor;

precise float dot4(precise vec4 a, precise vec4 b) {
  precise float s = a.x*b.x;
  s = s + a.y*b.y; s = s + a.z*b.z; s = s + a.w*b.w;
  return s;
}
ivec2 widx(int i) { int t = i >> 2; return ivec2(t & 4095, t >> 12); }
float getIn(int i) { return texelFetch(u_in, widx(i), 0)[i&3]; }
float getW(int i) { return texelFetch(u_w, widx(i), 0)[i&3]; }
float getRes(int i) { return texelFetch(u_res, widx(i), 0)[i&3]; }

void main() {
  if (u_mode != 35 && gl_FragCoord.y >= 1.0) { fragColor = vec4(0); return; }
  int px = int(gl_FragCoord.x);
  float y[4] = float[4](0., 0., 0., 0.);

  // mode 3: SILU (unchanged)
  if (u_mode == 3) {
    for (int l = 0; l < 4; l++) {
      int i = px*4 + l + u_elem_base;
      if (i < u_in_len) { float x = getIn(i); y[l] = x / (1.0 + exp(-x)); }
    }
  }
  // mode 7: RMSNORM — fast path via u_f1 (caller-precomputed r).
  // Slow in-shader reduction retained as fallback when u_f1 == 0.
  // WARNING: slow path is O(N) per pixel × N/4 pixels = O(N²) — catastrophic
  // for large n_embd. Callers SHOULD pass r via u_f1 whenever possible.
  else if (u_mode == 7) {
    float r;
    if (u_f1 > 0.0) {
      r = u_f1;                                  // fast path: caller-computed
    } else {
      // slow fallback: in-shader reduction (O(N²) total)
      float ss = 0.0;
      for (int i = 0; i < 8192; i++) {
        if (i >= u_in_len) break;
        float v = getIn(i); ss += v*v;
      }
      r = inversesqrt(ss / float(u_in_len) + u_f0);
    }
    for (int l = 0; l < 4; l++) {
      int i = px*4 + l + u_elem_base;
      if (i < u_in_len) y[l] = getIn(i) * r * getW(i);
    }
  }
  // mode 8: ROPE (runtime) — pow/cos/sin run per element.
  // Prefer mode 9 when possible: mode 8 kept for compat / A-B benching.
  else if (u_mode == 8) {
    float base = (u_f1 > 0.0) ? u_f1 : 10000.0;
    int hd = u_i0;
    float pos = u_f0;
    for (int l = 0; l < 4; l++) {
      int i = px*4 + l + u_elem_base;
      if (i >= u_in_len) continue;
      int pair = (i % hd) & ~1;
      int b = i - (i % hd) + pair;
      float th = pos / pow(base, float(pair) / float(hd));
      float c = cos(th), s = sin(th);
      float x0 = getIn(b), x1 = getIn(b + 1);
      if ((i & 1) == 0) y[l] = x0*c - x1*s;
      else y[l] = x0*s + x1*c;
    }
  }
  // mode 9: ROPE_PRECOMP — cos/sin pairs from u_w.
  // Layout: w[2k] = cos(theta_k), w[2k+1] = sin(theta_k), k = 0..hd/2-1,
  // where theta_k = pos / base^(2k/hd). Caller computes this in Python
  // (or C) and passes the hd-length table as the weight buffer.
  // Skips pow/cos/sin — only the four arithmetic ops per pair remain.
  // Numerics match REFOps.rope bit-for-bit given the same f32 table.
  else if (u_mode == 9) {
    int hd = u_i0;
    for (int l = 0; l < 4; l++) {
      int i = px*4 + l + u_elem_base;
      if (i >= u_in_len) continue;
      int in_head = i % hd;
      int pair_idx = in_head >> 1;
      float c = getW(2*pair_idx);
      float s = getW(2*pair_idx + 1);
      int b = i - in_head + (in_head & ~1);
      float x0 = getIn(b), x1 = getIn(b + 1);
      if ((in_head & 1) == 0) y[l] = x0*c - x1*s;
      else                    y[l] = x0*s + x1*c;
    }
  }
  // mode 15: MUL
  else if (u_mode == 15) {
    for (int l = 0; l < 4; l++) {
      int i = px*4 + l + u_elem_base;
      if (i < u_in_len) y[l] = getIn(i) * getW(i);
    }
  }
  // mode 16: ADD
  else if (u_mode == 16) {
    for (int l = 0; l < 4; l++) {
      int i = px*4 + l + u_elem_base;
      if (i < u_in_len) y[l] = getIn(i) + getW(i);
    }
  }
  // mode 23: SWIGLU
  else if (u_mode == 23) {
    int mid = u_in_len / 2;
    for (int l = 0; l < 4; l++) {
      int i = px*4 + l;
      if (i < u_out_len) {
        float g = getIn(i);
        float u = getIn(i + mid);
        float sg = g / (1.0 + exp(-g));
        y[l] = sg * u;
      }
    }
  }
  // mode 34: LINEAR f32 (4 rows per pixel, x4 hoisted out of jj loop)
  else if (u_mode == 34) {
    int base_row = (px + u_row_offset) * 4;
    precise float acc[4];
    acc[0] = (base_row+0 < u_out_len) ? getW(u_out_len*u_in_len + base_row+0) : 0.0;
    acc[1] = (base_row+1 < u_out_len) ? getW(u_out_len*u_in_len + base_row+1) : 0.0;
    acc[2] = (base_row+2 < u_out_len) ? getW(u_out_len*u_in_len + base_row+2) : 0.0;
    acc[3] = (base_row+3 < u_out_len) ? getW(u_out_len*u_in_len + base_row+3) : 0.0;
    for (int c = 0; c < 8192; c++) {
      if ((c << 2) >= u_in_len) break;
      precise vec4 x4 = precise vec4(getIn(c*4), getIn(c*4+1),
                                      getIn(c*4+2), getIn(c*4+3));
      for (int jj = 0; jj < 4; jj++) {
        int row = base_row + jj;
        if (row >= u_out_len) break;
        int t = row * (u_in_len >> 2) + c;
        precise vec4 w4 = texelFetch(u_w, ivec2(t & 4095, t >> 12), 0);
        acc[jj] = acc[jj] + dot4(w4, x4);
      }
    }
    for (int jj = 0; jj < 4; jj++) {
      int row = base_row + jj;
      if (row >= u_out_len) { y[jj] = -1e30; continue; }
      if (u_has_res == 1) acc[jj] += getRes(row);
      y[jj] = acc[jj];
    }
  }
  // mode 35: SDPA_DENSE
  else if (u_mode == 35) {
    int Tq = u_i0, Tkv = u_i1, hd = u_i2;
    int start_pos = u_i3, nh = u_i4, nkv = u_i5;
    int d_model = nh * hd, kv_dim = nkv * hd;
    int max_len = (u_f1 > 0.0) ? int(u_f1) : Tkv;
    int vb = max_len * kv_dim;
    int base_idx = (px + u_row_offset) * 4;
    for (int jj = 0; jj < 4; jj++) {
      int idx = base_idx + jj;
      if (idx < Tq * d_model) {
        int ti = idx / d_model; int r = idx - ti * d_model;
        int hh = r / hd; int ki = r - hh * hd;
        int hk = hh / (nh / nkv);
        float sc = 1.0 / sqrt(float(hd));
        int max_j = start_pos + ti + 1;
        if (max_j > max_len) max_j = max_len;
        precise float mx = -1e30, z = 0., acc = 0.;
        for (int j = 0; j < 8192; j++) {
          if (j >= max_j) break;
          precise float s = 0.;
          int qb = ti * d_model + hh * hd;
          int kb = j * kv_dim + hk * hd;
          int k = 0;
          for (; k + 4 <= hd; k += 4) {
            precise vec4 q4 = precise vec4(getIn(qb+k), getIn(qb+k+1),
                                            getIn(qb+k+2), getIn(qb+k+3));
            precise vec4 k4 = precise vec4(getW(kb+k), getW(kb+k+1),
                                            getW(kb+k+2), getW(kb+k+3));
            s = s + dot4(q4, k4);
          }
          for (; k < hd; k++) {
            s = s + getIn(qb + k) * getW(kb + k);
          }
          s = s * sc;
          precise float newmx = max(mx, s);
          precise float rs = exp(mx - newmx);
          precise float p = exp(s - newmx);
          acc = acc * rs + p * getW(vb + j*kv_dim + hk*hd + ki);
          z = z * rs + p;
          mx = newmx;
        }
        y[jj] = acc / z;
      }
    }
  }
  fragColor = vec4(y[0], y[1], y[2], y[3]);
}