function zeros(n: number) {
  return Array.from({ length: n }, () => 0);
}

function zeros2(r: number, c: number) {
  return Array.from({ length: r }, () => zeros(c));
}

function randn(rng: () => number) {
  const u = Math.max(rng(), 1e-9);
  const v = Math.max(rng(), 1e-9);
  return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v);
}

function sigmoid(x: number) {
  if (x > 20) return 1;
  if (x < -20) return 0;
  return 1 / (1 + Math.exp(-x));
}

function clip(x: number, m = 5) {
  return Math.max(-m, Math.min(m, x));
}

export class LSTM {
  V: number;
  H: number;
  Wx: number[][];
  Wh: number[][];
  b: number[];
  Why: number[][];
  by: number[];
  lr: number;

  constructor(V: number, H: number, rng: () => number, lr = 0.08) {
    this.V = V;
    this.H = H;
    this.lr = lr;
    const sV = Math.sqrt(2 / (V + H));
    const sH = Math.sqrt(2 / (H + H));
    this.Wx = zeros2(4 * H, V).map((row) => row.map(() => randn(rng) * sV));
    this.Wh = zeros2(4 * H, H).map((row) => row.map(() => randn(rng) * sH));
    this.b = zeros(4 * H);
    this.b.splice(H, H, ...Array.from({ length: H }, () => 1));
    this.Why = zeros2(V, H).map((row) => row.map(() => randn(rng) * Math.sqrt(2 / H)));
    this.by = zeros(V);
  }

  private gates(xIdx: number, h: number[], c: number[]) {
    const { H, Wx, Wh, b } = this;
    const pre = zeros(4 * H);
    for (let i = 0; i < 4 * H; i++) {
      let s = b[i] + Wx[i][xIdx];
      const wh = Wh[i];
      for (let j = 0; j < H; j++) s += wh[j] * h[j];
      pre[i] = s;
    }
    const i = zeros(H);
    const f = zeros(H);
    const o = zeros(H);
    const g = zeros(H);
    for (let j = 0; j < H; j++) {
      i[j] = sigmoid(pre[j]);
      f[j] = sigmoid(pre[H + j]);
      o[j] = sigmoid(pre[2 * H + j]);
      g[j] = Math.tanh(pre[3 * H + j]);
    }
    const c2 = zeros(H);
    const h2 = zeros(H);
    for (let j = 0; j < H; j++) {
      c2[j] = f[j] * c[j] + i[j] * g[j];
      h2[j] = o[j] * Math.tanh(c2[j]);
    }
    return { i, f, o, g, c2, h2, pre };
  }

  softmax(h: number[]) {
    const logits = zeros(this.V);
    let max = -Infinity;
    for (let v = 0; v < this.V; v++) {
      let s = this.by[v];
      const row = this.Why[v];
      for (let j = 0; j < this.H; j++) s += row[j] * h[j];
      logits[v] = s;
      if (s > max) max = s;
    }
    const p = zeros(this.V);
    let sum = 0;
    for (let v = 0; v < this.V; v++) {
      const e = Math.exp(logits[v] - max);
      p[v] = e;
      sum += e;
    }
    for (let v = 0; v < this.V; v++) p[v] /= sum;
    return { logits, p };
  }

  sequenceNll(ids: number[]) {
    let h = zeros(this.H);
    let c = zeros(this.H);
    let nll = 0;
    let steps = 0;
    for (let t = 0; t < ids.length - 1; t++) {
      const step = this.gates(ids[t], h, c);
      h = step.h2;
      c = step.c2;
      const { p } = this.softmax(h);
      nll += -Math.log(Math.max(p[ids[t + 1]], 1e-12));
      steps += 1;
    }
    return steps ? nll / steps : 0;
  }

  hiddenTrace(ids: number[]) {
    let h = zeros(this.H);
    let c = zeros(this.H);
    const trace: number[][] = [];
    for (let t = 0; t < ids.length; t++) {
      const step = this.gates(ids[t], h, c);
      h = step.h2;
      c = step.c2;
      trace.push([...h]);
    }
    return trace;
  }

  topk(h: number[], k: number) {
    const { p } = this.softmax(h);
    return [...p.keys()]
      .sort((a, b) => p[b] - p[a])
      .slice(0, k)
      .map((id) => ({ id, p: p[id] }));
  }

  trainSequence(ids: number[]) {
    const T = ids.length - 1;
    if (T <= 0) return 0;
    const H = this.H;
    const xs: number[] = [];
    const hs: number[][] = [zeros(H)];
    const cs: number[][] = [zeros(H)];
    const is: number[][] = [];
    const fs: number[][] = [];
    const os: number[][] = [];
    const gs: number[][] = [];
    const ps: number[][] = [];
    let nll = 0;

    for (let t = 0; t < T; t++) {
      xs.push(ids[t]);
      const step = this.gates(ids[t], hs[t], cs[t]);
      is.push(step.i);
      fs.push(step.f);
      os.push(step.o);
      gs.push(step.g);
      hs.push(step.h2);
      cs.push(step.c2);
      const sm = this.softmax(step.h2);
      ps.push(sm.p);
      nll += -Math.log(Math.max(sm.p[ids[t + 1]], 1e-12));
    }

    const dWx = zeros2(4 * H, this.V);
    const dWh = zeros2(4 * H, H);
    const db = zeros(4 * H);
    const dWhy = zeros2(this.V, H);
    const dby = zeros(this.V);
    let dhNext = zeros(H);
    let dcNext = zeros(H);

    for (let t = T - 1; t >= 0; t--) {
      const target = ids[t + 1];
      const p = ps[t];
      const h = hs[t + 1];
      const dy = p.map((v, i) => v - (i === target ? 1 : 0));
      for (let v = 0; v < this.V; v++) {
        dby[v] += dy[v];
        const row = dWhy[v];
        for (let j = 0; j < H; j++) row[j] += dy[v] * h[j];
      }
      const dh = zeros(H);
      for (let j = 0; j < H; j++) {
        let s = dhNext[j];
        for (let v = 0; v < this.V; v++) s += this.Why[v][j] * dy[v];
        dh[j] = s;
      }

      const i = is[t];
      const f = fs[t];
      const o = os[t];
      const g = gs[t];
      const c = cs[t + 1];
      const cPrev = cs[t];
      const hPrev = hs[t];
      const tanhC = c.map(Math.tanh);
      const do_ = zeros(H);
      const dc = zeros(H);
      const di = zeros(H);
      const df = zeros(H);
      const dg = zeros(H);
      for (let j = 0; j < H; j++) {
        do_[j] = dh[j] * tanhC[j];
        dc[j] = dh[j] * o[j] * (1 - tanhC[j] * tanhC[j]) + dcNext[j];
        di[j] = dc[j] * g[j];
        df[j] = dc[j] * cPrev[j];
        dg[j] = dc[j] * i[j];
        dcNext[j] = dc[j] * f[j];
      }
      const dpre = zeros(4 * H);
      for (let j = 0; j < H; j++) {
        dpre[j] = di[j] * i[j] * (1 - i[j]);
        dpre[H + j] = df[j] * f[j] * (1 - f[j]);
        dpre[2 * H + j] = do_[j] * o[j] * (1 - o[j]);
        dpre[3 * H + j] = dg[j] * (1 - g[j] * g[j]);
      }
      const xIdx = xs[t];
      dhNext = zeros(H);
      for (let u = 0; u < 4 * H; u++) {
        const gp = dpre[u];
        db[u] += gp;
        dWx[u][xIdx] += gp;
        const wh = this.Wh[u];
        const dwh = dWh[u];
        for (let j = 0; j < H; j++) {
          dwh[j] += gp * hPrev[j];
          dhNext[j] += wh[j] * gp;
        }
      }
    }

    const scale = this.lr / T;
    const apply2 = (w: number[][], d: number[][]) => {
      for (let r = 0; r < w.length; r++) {
        const wr = w[r];
        const dr = d[r];
        for (let c = 0; c < wr.length; c++) wr[c] -= scale * clip(dr[c]);
      }
    };
    apply2(this.Wx, dWx);
    apply2(this.Wh, dWh);
    apply2(this.Why, dWhy);
    for (let i = 0; i < this.b.length; i++) this.b[i] -= scale * clip(db[i]);
    for (let i = 0; i < this.by.length; i++) this.by[i] -= scale * clip(dby[i]);
    return nll / T;
  }
}