import { aE as E, ab as T, ah as O, N as V, aW as B, K as F, aM as K, aX as W } from "./index-DSGwv2Yx.js";
import { r as $ } from "./Reshape-BZC-ebeR.js";
import { a as A, b as k, d as C, c as N, e as R } from "./axis_util-QWWgLjut.js";
import { t as U, m as _ } from "./shared-zTaJ5siv.js";
import { c as j } from "./backend_util-qwSFfxYx.js";
import { f as y } from "./gpgpu_math-3bCb5ooU.js";
import { g as G, b as L } from "./kernel_funcs_utils-DGqzNlHT.js";
class w {
  constructor(s, e) {
    this.variableNames = ["x"];
    const { windowSize: t, batchSize: n, inSize: l, outSize: r } = s;
    this.outputShape = [n, r];
    const i = Math.floor(t / 4) * 4, c = t % 4;
    let o = "sumValue += dot(values, ones);";
    if (e != null) {
      const p = 1 / e;
      o = `sumValue += dot(values * ${E(p) ? p.toPrecision(2) : p}, ones);`;
    }
    let u = "";
    l % t > 0 && (u = `
        if (inIdx < 0 || inIdx >= ${l}) {
          return 0.0;
        }
      `), this.userCode = `
      const vec4 ones = vec4(1.0, 1.0, 1.0, 1.0);

      float getValue(int batch, int inIdx) {
        ${u}
        return getX(batch, inIdx);
      }

      void main() {
        ivec2 coords = getOutputCoords();
        int batch = coords[0];
        int outIdx = coords[1];
        int inOffset = outIdx * ${t};

        float sumValue = 0.0;

        for (int i = 0; i < ${i}; i += 4) {
          int inIdx = inOffset + i;
          vec4 values = vec4(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            getValue(batch, inIdx + 2),
            getValue(batch, inIdx + 3)
          );

          ${o}
        }

        int inIdx = inOffset + ${i};
        if (${c === 1}) {
          vec4 values = vec4(getValue(batch, inIdx), 0.0, 0.0, 0.0);

          ${o}
        } else if (${c === 2}) {
          vec4 values = vec4(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1), 0.0, 0.0);

          ${o}
        } else if (${c === 3}) {
          vec4 values = vec4(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            getValue(batch, inIdx + 2), 0.0);

          ${o}
        }
        setOutput(sumValue);
      }
    `;
  }
}
class X {
  constructor(s, e) {
    this.variableNames = ["x"];
    const { windowSize: t, batchSize: n, inSize: l, outSize: r } = s;
    this.outputShape = [n, r];
    let i = "0.0", c = "";
    e === "prod" ? i = "1.0" : e === "min" ? (i = "1.0 / 1e-20", c = "min") : e === "max" && (i = "-1.0 / 1e-20", c = "max");
    let o = `${e}(${e}(${e}(minMaxValue[0], minMaxValue[1]), minMaxValue[2]), minMaxValue[3])`;
    e === "sum" ? o = "sumValue" : e === "prod" ? o = "prodValue" : e === "all" ? o = "allValue" : e === "any" && (o = "anyValue");
    const u = Math.floor(t / 4) * 4, p = t % 4;
    let h = `
      if (${e === "sum"}) {
        sumValue += dot(values, ones);
      } else if (${e === "prod"}) {
        vec2 tmp = vec2(values[0], values[1]) * vec2(values[2], values[3]);
        prodValue *= tmp[0] * tmp[1];
      } else {
        minMaxValue = ${c}(values, minMaxValue);
        if (${e === "min"} || ${e === "max"}) {
          minMaxValue = ${c}(values, minMaxValue);
          bvec4 isNaN = isnan(values);
          if (isNaN.r || isNaN.g || isNaN.b || isNaN.a) {
            minMaxValue = vec4(NAN);
          }
        }
      }
    `, d = "vec4";
    e === "all" ? (i = "1.0", h = `
        bool reducedAllValue = all(values);
        float floatedReducedAllValue = float(reducedAllValue);
        allValue = float(allValue >= 1.0 && floatedReducedAllValue >= 1.0);
      `, d = "bvec4") : e === "any" && (i = "0.0", h = `
        bool reducedAnyValue = any(values);
        float floatedReducedAnyValue = float(reducedAnyValue);
        anyValue = float(anyValue >= 1.0 || floatedReducedAnyValue >= 1.0);
      `, d = "bvec4");
    let f = "";
    l % t > 0 && (f = `
        if (inIdx < 0 || inIdx >= ${l}) {
          return initializationValue;
        }
      `), this.userCode = `
      const float initializationValue = ${i};
      const vec4 ones = vec4(1.0, 1.0, 1.0, 1.0);

      float getValue(int batch, int inIdx) {
        ${f}
        return getX(batch, inIdx);
      }

      void main() {
        ivec2 coords = getOutputCoords();
        int batch = coords[0];
        int outIdx = coords[1];
        int inOffset = outIdx * ${t};

        vec4 minMaxValue = vec4(${i});
        float prodValue = 1.0;
        float sumValue = 0.0;
        float allValue = 1.0;
        float anyValue = 0.0;

        for (int i = 0; i < ${u}; i += 4) {
          int inIdx = inOffset + i;
          ${d} values = ${d}(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            getValue(batch, inIdx + 2),
            getValue(batch, inIdx + 3)
          );

          ${h}
        }

        int inIdx = inOffset + ${u};
        if (${p === 1}) {
          ${d} values = ${d}(
            getValue(batch, inIdx),
            initializationValue,
            initializationValue,
            initializationValue
          );

          ${h}
        } else if (${p === 2}) {
          ${d} values = ${d}(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            initializationValue,
            initializationValue
          );

          ${h}
        } else if (${p === 3}) {
          ${d} values = ${d}(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            getValue(batch, inIdx + 2),
            initializationValue
          );

          ${h}
        }
        setOutput(${o});
      }
    `;
  }
}
function q(a) {
  const s = [];
  for (; s.length === 0 || s[s.length - 1].outSize !== 1; ) {
    const e = s.length ? s[s.length - 1].outSize : a[1], t = j(e);
    s.push({
      inSize: e,
      windowSize: t,
      outSize: Math.ceil(e / t)
    });
  }
  return s;
}
function M(a, s, e, t) {
  const n = q(a.shape);
  let l = a;
  for (let r = 0; r < n.length; r++) {
    const { inSize: i, windowSize: c, outSize: o } = n[r];
    let u, p;
    e === "mean" ? u = r === 0 ? new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, i) : new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }) : u = new X({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, e), p = l, l = t.runWebGLProgram(u, [l], s), p.dataId !== a.dataId && t.disposeIntermediateTensorInfo(p);
  }
  return l;
}
class Y {
  constructor(s, e) {
    this.variableNames = ["A"];
    const t = new Array(s.length);
    for (let r = 0; r < t.length; r++)
      t[r] = s[e[r]];
    this.outputShape = t, this.rank = t.length;
    const n = y(this.rank), l = H(e);
    this.userCode = `
    void main() {
      ${n} resRC = getOutputCoords();
      setOutput(getA(${l}));
    }
    `;
  }
}
function H(a) {
  const s = a.length;
  if (s > 6)
    throw Error(`Transpose for rank ${s} is not yet supported`);
  const e = ["resRC.x", "resRC.y", "resRC.z", "resRC.w", "resRC.u", "resRC.v"], t = new Array(s);
  for (let n = 0; n < a.length; n++)
    t[a[n]] = e[n];
  return t.join();
}
class J {
  constructor(s, e) {
    this.variableNames = ["A"], this.packedInputs = !0, this.packedOutput = !0;
    const t = new Array(s.length);
    for (let u = 0; u < t.length; u++)
      t[u] = s[e[u]];
    if (this.outputShape = t, this.rank = t.length, this.rank > 6)
      throw Error(`Packed transpose for rank ${this.rank} is not yet supported.`);
    const n = y(this.rank), l = G("rc", this.rank), r = new Array(this.rank);
    for (let u = 0; u < e.length; u++)
      r[e[u]] = l[u];
    const i = `vec2(${r.slice(-2).join()})`, c = `++${l[this.rank - 1]} < ${t[this.rank - 1]}`, o = `getChannel(getA(${r.join()}), ${i})`;
    this.userCode = `
    void main() {
      ${n} rc = getOutputCoords();
      vec4 result = vec4(0.);
      result[0] = ${o};
      if(${c}) {
        result[1] = ${o};
      }
      --${l[this.rank - 1]};
      if(++${l[this.rank - 2]} < ${t[this.rank - 2]}) {
        result[2] = ${o};
        if(${c}) {
          result[3] = ${o};
        }
      }
      setOutput(result);
    }
    `;
  }
}
function P(a, s, e) {
  const t = T().getBool("WEBGL_PACK_ARRAY_OPERATIONS") ? new J(a.shape, s) : new Y(a.shape, s);
  return e.runWebGLProgram(t, [a], a.dtype);
}
function Q(a, s, e, t) {
  const n = s, l = a.shape.length, r = O(n, a.shape);
  let i = r;
  const c = A(i, l), o = c != null;
  let u = a;
  o && (u = P(a, c, t), i = k(i.length, l)), C("sum", i, l);
  const [p, h] = N(u.shape, i);
  let d = p;
  e && (d = R(p, r));
  const f = V(h), g = V(a.shape) / f, x = $({ inputs: { x: u }, attrs: { shape: [g, f] }, backend: t }), S = B(a.dtype), I = M(x, S, "sum", t), m = $({ inputs: { x: I }, attrs: { shape: d }, backend: t });
  return t.disposeIntermediateTensorInfo(x), t.disposeIntermediateTensorInfo(I), o && t.disposeIntermediateTensorInfo(u), m;
}
function Z(a) {
  const { inputs: s, backend: e, attrs: t } = a, { x: n } = s, { axis: l, keepDims: r } = t;
  return Q(n, l, r, e);
}
const pe = {
  kernelName: F,
  backendName: "webgl",
  kernelFunc: Z
};
function ee(a, s, e, t) {
  const n = V(s), r = V(a.shape) / n, i = $({ inputs: { x: a }, attrs: { shape: [r, n] }, backend: t }), c = M(i, a.dtype, "max", t), o = $({ inputs: { x: c }, attrs: { shape: e }, backend: t });
  return t.disposeIntermediateTensorInfo(i), t.disposeIntermediateTensorInfo(c), o;
}
function te(a) {
  const { inputs: s, backend: e, attrs: t } = a, { x: n } = s, { reductionIndices: l, keepDims: r } = t, i = n.shape.length, c = O(l, n.shape);
  let o = c;
  const u = A(o, i), p = u != null, h = e.shouldExecuteOnCPU([n]);
  let d = n;
  if (p) {
    if (h) {
      const I = e.texData.get(d.dataId).values, m = new Array(i);
      for (let v = 0; v < m.length; v++)
        m[v] = n.shape[u[v]];
      const z = U(I, n.shape, n.dtype, u, m);
      d = e.makeTensorInfo(m, n.dtype);
      const D = e.texData.get(d.dataId);
      D.values = z;
    } else
      d = P(n, u, e);
    o = k(o.length, i);
  }
  C("max", o, i);
  const [f, b] = N(d.shape, o);
  let g = f;
  r && (g = R(f, c));
  let x;
  if (h) {
    const I = e.texData.get(d.dataId).values, m = _(I, V(b), g, n.dtype);
    x = e.makeTensorInfo(g, n.dtype);
    const z = e.texData.get(x.dataId);
    z.values = m;
  } else
    x = ee(d, b, g, e);
  return p && e.disposeIntermediateTensorInfo(d), x;
}
const he = {
  kernelName: K,
  backendName: "webgl",
  kernelFunc: te
};
const ae = `
if (a == b) {
  return 1.0;
};
return a / b;`, se = `
  // vec4 one = vec4(equal(a, b));
  // return one + (vec4(1.0) - one) * a / b;
  vec4 result = a / b;
  if(a.x == b.x) {
    result.x = 1.;
  }
  if(a.y == b.y) {
    result.y = 1.;
  }
  if(a.z == b.z) {
    result.z = 1.;
  }
  if(a.w == b.w) {
    result.w = 1.;
  }

  return result;
`, ne = L({ opSnippet: ae, packedOpSnippet: se, checkOutOfBounds: !0 }), fe = {
  kernelName: W,
  backendName: "webgl",
  kernelFunc: ne
};
export {
  M as a,
  he as b,
  fe as c,
  pe as d,
  te as m,
  ne as r,
  Z as s,
  P as t
};
