// TRITIUM-PROVENANCE sha256=6d7ce1815a5ea4eacf5df829ac23026deec566513c6bfb4d5c0fbfd4974218ee lang=rust
// TRITIUM-HEIGHT high=1.00 mid=0.00 low=0.00
fn fmean(vals: &mut [f64], n: i64) -> f64 {

    let mut x = 0.0;
    let mut i = 0;
    while i < n {
        x = x + vals[i as usize];
        i += 1;
    }
    x / ((n) as f64)
}

fn fsqrt(x: f64) -> f64 {
    let mut i: f64 = 0;
    let mut i_2 = 0.0;
    let mut x_2 = 0.0;

    let mut fcmp4 = ((x) <= (0.0)) as i64;
    if (fcmp4) != 0 {
        return 0.0;
    } else {
        i_2 = x;
        i = 0;
        while i < 8 {
            x_2 = 0.5 * (i_2 + (x / i_2));
            i_2 = x_2;
            i += 1;
        }
        return i_2;
    }
}

fn layer_norm(vals: &mut [f64], n: i64, out: &mut [f64], eps: f64) -> f64 {
    let mut x_2 = 0.0;

    let mut mu = fmean(vals, n);
    let mut x = 0.0;
    let mut i = 0;
    while i < n {
        x_2 = vals[i as usize] - mu;
        x = x + (x_2 * x_2);
        i += 1;
    }
    x = (x / ((n) as f64));
    let mut x_3 = 1.0 / (fsqrt(x + eps));
    i = 0;
    while i < n {
        out[i as usize] = ((vals[i as usize] - mu) * x_3);
        i += 1;
    }
    mu
}

fn ln_rows(x: &mut [f64], n_tokens: i64, d: i64, out: &mut [f64], tmp: &mut [f64], nrm: &mut [f64], eps: f64) -> f64 {
    let mut i: f64 = 0;

    let mut t = 0;
    while t < n_tokens {
        i = 0;
        while i < d {
            tmp[i as usize] = x[((((t * d) + i)) as usize)];
            i += 1;
        }
        layer_norm(tmp, d, nrm, eps);
        i = 0;
        while i < d {
            out[((((t * d) + i)) as usize)] = nrm[i as usize];
            i += 1;
        }
        t += 1;
    }
    let mut x_2 = ((n_tokens) as f64) * 1.0;
    x_2
}

fn gemv_f(a: &mut [f64], x: &mut [f64], n: i64, m: i64, out: &mut [f64]) -> i64 {
    let mut j = 0;
    let mut x_2 = 0.0;

    let mut i = 0;
    while i < n {
        x_2 = 0.0;
        j = 0;
        while j < m {
            x_2 = x_2 + (a[((((i * m) + j)) as usize)] * x[j as usize]);
            j += 1;
        }
        out[i as usize] = x_2;
        i += 1;
    }
    n
}

fn qkv_seq(x: &mut [f64], n_tokens: i64, d: i64, w: &mut [f64], out: &mut [f64], tmp: &mut [f64], row: &mut [f64]) -> i64 {
    let mut i: f64 = 0;

    let mut t = 0;
    while t < n_tokens {
        i = 0;
        while i < d {
            tmp[i as usize] = x[((((t * d) + i)) as usize)];
            i += 1;
        }
        gemv_f(w, tmp, d, d, row);
        i = 0;
        while i < d {
            out[((((t * d) + i)) as usize)] = row[i as usize];
            i += 1;
        }
        t += 1;
    }
    n_tokens
}

fn pack_heads(src: &mut [f64], n_tokens: i64, heads: i64, dh: i64, dest: &mut [f64]) -> i64 {
    let mut t = 0;
    let mut u = 0;

    let mut h = 0;
    while h < heads {
        t = 0;
        while t < n_tokens {
            u = 0;
            while u < dh {
                dest[((((((h * n_tokens) * dh) + (t * dh)) + u)) as usize)] = src[((((((t * heads) * dh) + (h * dh)) + u)) as usize)];
                u += 1;
            }
            t += 1;
        }
        h += 1;
    }
    (heads * n_tokens) * dh
}

fn fexp(x: f64) -> f64 {
    let mut j = 0;
    let mut ifexp_2 = 0.0;
    let mut x_2 = 0.0;

    let mut ln2 = 0.6931471805599453;
    let mut inv_ln2 = 1.4426950408889634;
    let mut fcmp15 = ((x) >= (0.0)) as i64;
    if (fcmp15) != 0 {
        ifexp_2 = 0.5;
    } else {
        ifexp_2 = -0.5;
    }
    let mut ifexp = ifexp_2;
    let mut xi = (((((x * inv_ln2) + ifexp) as i64)));
    if xi > 60 {
        xi = 60;
    }
    if xi < -60 {
        xi = -60;
    }
    let mut x_5 = 1.0;
    let mut x_4 = 1.0;
    let mut i = 1;
    while i <= 14 {
        x_4 = (x_4 * (x - (((xi) as f64) * ln2))) / ((i) as f64);
        x_5 = x_5 + x_4;
        i += 1;
    }
    let mut x_3 = 1.0;
    let mut i_2 = xi;
    if i_2 < 0 {
        i_2 = 0 - i_2;
        j = 0;
        while j < i_2 {
            x_3 = x_3 * 2.0;
            j += 1;
        }
        return x_5 / x_3;
    } else {
        j = 0;
        while j < i_2 {
            x_2 = x_3 * 2.0;
            x_3 = x_2;
            j += 1;
        }
        return x_5 * x_3;
    }
}

fn mha_fwd(q: &mut [f64], k: &mut [f64], v: &mut [f64], scores: &mut [f64], out: &mut [f64], h: i64, l: i64, d: i64, scale: f64) -> i64 {
    let mut i: f64 = 0;
    let mut j = 0;
    let mut t = 0;
    let mut x_2 = 0.0;
    let mut i_2 = 0.0;
    let mut fcmp96 = 0;
    let mut e = 0.0;
    let mut x = 0.0;

    let mut heap_dist = 0;
    while heap_dist < h {
        i = 0;
        while i < l {
            j = 0;
            while j < l {
                x_2 = 0.0;
                t = 0;
                while t < d {
                    x_2 = x_2 + (q[((((((heap_dist * l) * d) + (i * d)) + t)) as usize)] * k[((((((heap_dist * l) * d) + (j * d)) + t)) as usize)]);
                    t += 1;
                }
                scores[((((i * l) + j)) as usize)] = x_2 * scale;
                j += 1;
            }
            i += 1;
        }
        i = 0;
        while i < l {
            i_2 = scores[((i * l) as usize)];
            j = 1;
            while j < l {
                fcmp96 = ((scores[((((i * l) + j)) as usize)]) > (i_2)) as i64;
                if (fcmp96) != 0 {
                    i_2 = scores[((((i * l) + j)) as usize)];
                }
                j += 1;
            }
            x = 0.0;
            j = 0;
            while j < l {
                e = fexp((scores[((((i * l) + j)) as usize)] - i_2));
                scores[((((i * l) + j)) as usize)] = e;
                x = x + e;
                j += 1;
            }
            j = 0;
            while j < l {
                scores[((((i * l) + j)) as usize)] = (scores[((((i * l) + j)) as usize)] / x);
                j += 1;
            }
            i += 1;
        }
        i = 0;
        while i < l {
            t = 0;
            while t < d {
                x_2 = 0.0;
                j = 0;
                while j < l {
                    x_2 = (x_2 + (scores[((((i * l) + j)) as usize)] * v[((((((heap_dist * l) * d) + (j * d)) + t)) as usize)]));
                    j += 1;
                }
                out[((((((heap_dist * l) * d) + (i * d)) + t)) as usize)] = x_2;
                t += 1;
            }
            i += 1;
        }
        heap_dist += 1;
    }
    (h * l) * d
}

fn unpack_heads(src: &mut [f64], n_tokens: i64, heads: i64, dh: i64, dest: &mut [f64]) -> i64 {
    let mut t = 0;
    let mut u = 0;

    let mut h = 0;
    while h < heads {
        t = 0;
        while t < n_tokens {
            u = 0;
            while u < dh {
                dest[((((((t * heads) * dh) + (h * dh)) + u)) as usize)] = src[((((((h * n_tokens) * dh) + (t * dh)) + u)) as usize)];
                u += 1;
            }
            t += 1;
        }
        h += 1;
    }
    (n_tokens * heads) * dh
}

fn residual_add(a: &mut [f64], b: &mut [f64], n: i64, out: &mut [f64]) -> i64 {

    let mut i = 0;
    while i < n {
        out[i as usize] = a[i as usize] + b[i as usize];
        i += 1;
    }
    n
}

fn ftanh(x: f64) -> f64 {

    let mut i = x;
    let mut fcmp5 = ((i) > (20.0)) as i64;
    if (fcmp5) != 0 {
        i = 20.0;
    }
    let mut fcmp7 = ((i) < (-20.0)) as i64;
    if (fcmp7) != 0 {
        i = -20.0;
    }
    let mut e = fexp(i + i);
    let mut x_2 = e - 1.0;
    let mut x_3 = e + 1.0;
    x_2 / x_3
}

fn gelu_tanh(x: f64) -> f64 {

    let mut x_5 = 0.5 * x;
    let mut x_4 = 0.044715 * ((x * x) * x);
    let mut x_2 = 0.79788456 * (x + x_4);
    let mut x_3 = 1.0 + (ftanh(x_2));
    x_5 * x_3
}

fn mlp_gelu_vec(x: &mut [f64], w1: &mut [f64], w2: &mut [f64], d: i64, hidden: i64, hidden_out: &mut [f64], out: &mut [f64]) -> i64 {

    gemv_f(w1, x, hidden, d, hidden_out);
    let mut i = 0;
    while i < hidden {
        hidden_out[i as usize] = (gelu_tanh(hidden_out[i as usize]));
        i += 1;
    }
    gemv_f(w2, hidden_out, d, hidden, out);
    d
}

fn transformer_block(x: &mut [f64], n_tokens: i64, d: i64, heads: i64, dh: i64, hidden: i64, scale: f64, wq: &mut [f64], wk: &mut [f64], wv: &mut [f64], wo: &mut [f64], w1: &mut [f64], w2: &mut [f64], q: &mut [f64], k: &mut [f64], v: &mut [f64], qh: &mut [f64], kh: &mut [f64], vh: &mut [f64], scores: &mut [f64], attention: &mut [f64], flat: &mut [f64], normed: &mut [f64], normed2: &mut [f64], tmp: &mut [f64], nrm: &mut [f64], row: &mut [f64], hidden_out: &mut [f64], mlp_out: &mut [f64], projected: &mut [f64], residual: &mut [f64], out: &mut [f64]) -> i64 {
    let mut i: f64 = 0;

    ln_rows(x, n_tokens, d, normed, tmp, nrm, 0.0);
    qkv_seq(normed, n_tokens, d, wq, q, tmp, row);
    qkv_seq(normed, n_tokens, d, wk, k, tmp, row);
    qkv_seq(normed, n_tokens, d, wv, v, tmp, row);
    pack_heads(q, n_tokens, heads, dh, qh);
    pack_heads(k, n_tokens, heads, dh, kh);
    pack_heads(v, n_tokens, heads, dh, vh);
    mha_fwd(qh, kh, vh, scores, attention, heads, n_tokens, dh, scale);
    unpack_heads(attention, n_tokens, heads, dh, flat);
    let mut t = 0;
    while t < n_tokens {
        i = 0;
        while i < d {
            tmp[i as usize] = flat[((((t * d) + i)) as usize)];
            i += 1;
        }
        gemv_f(wo, tmp, d, d, row);
        i = 0;
        while i < d {
            projected[((((t * d) + i)) as usize)] = row[i as usize];
            i += 1;
        }
        t += 1;
    }
    residual_add(x, projected, n_tokens * d, residual);
    ln_rows(residual, n_tokens, d, normed2, tmp, nrm, 0.0);
    t = 0;
    while t < n_tokens {
        i = 0;
        while i < d {
            tmp[i as usize] = normed2[((((t * d) + i)) as usize)];
            i += 1;
        }
        mlp_gelu_vec(tmp, w1, w2, d, hidden, hidden_out, row);
        i = 0;
        while i < d {
            mlp_out[((((t * d) + i)) as usize)] = row[i as usize];
            i += 1;
        }
        t += 1;
    }
    residual_add(residual, mlp_out, n_tokens * d, out);
    n_tokens
}