class TinyLM {
  constructor(config, weights) {
    this.config = config;
    this.weights = weights;
  }

  norm(x, name) {
    const {mean, variance} = tf.moments(x, -1, true);
    return x.sub(mean).mul(variance.add(this.config.epsilon).rsqrt())
      .mul(this.weights[`${name}.gamma`]).add(this.weights[`${name}.beta`]);
  }

  dense(x, name) {
    return tf.matMul(x, this.weights[`${name}.kernel`]).add(this.weights[`${name}.bias`]);
  }

  step(bytes, cache = []) {
    const {width, heads, layers, context} = this.config;
    const offset = cache.length ? cache[0].key.shape[1] : 0;
    const length = bytes.length;
    if (!length || offset + length > context) throw new Error("Context limit reached");
    return tf.tidy(() => {
      const ids = tf.tensor1d(bytes, "int32");
      let x = tf.gather(this.weights.tokens, ids)
        .add(tf.gather(this.weights.positions, tf.range(offset, offset + length, 1, "int32")));
      const next = [];
      for (let i = 0; i < layers; i++) {
        const name = `b${i}`;
        const normalized = this.norm(x, `${name}.norm1`);
        const project = part => this.dense(normalized, `${name}.${part}`)
          .reshape([length, heads, width / heads]).transpose([1, 0, 2]);
        const query = project("q");
        let key = project("k"), value = project("v");
        if (cache.length) {
          key = tf.concat([cache[i].key, key], 1);
          value = tf.concat([cache[i].value, value], 1);
        }
        next.push({key, value});
        let scores = tf.matMul(query, key, false, true).mul(1 / Math.sqrt(width / heads));
        if (length > 1) {
          const future = tf.range(0, offset + length, 1, "int32").expandDims(0)
            .greater(tf.range(offset, offset + length, 1, "int32").expandDims(1));
          scores = scores.add(future.cast("float32").mul(-1e9));
        }
        const attention = tf.matMul(tf.softmax(scores, -1), value)
          .transpose([1, 0, 2]).reshape([length, width]);
        x = x.add(this.dense(attention, `${name}.o`));
        const expanded = this.dense(this.norm(x, `${name}.norm2`), `${name}.expand`);
        const gelu = expanded.mul(0.5).mul(tf.erf(expanded.div(Math.SQRT2)).add(1));
        x = x.add(this.dense(gelu, `${name}.contract`));
      }
      const last = this.norm(x.slice([length - 1, 0], [1, width]), "norm");
      return {logits: this.dense(last, "head").reshape([256]), cache: next};
    });
  }
}

self.TinyLM = TinyLM;