# Elementwise expressions and reductions over slices, as the compiled
# engines fuse them (lang/viv/src/fuse.rs): every result is compared to the
# element loop that writes the same arithmetic in the same order, to the bit.

N: Int = 300

#| y = a x + b z - c, elementwise
def axpbz(a: Num, x: Arr, b: Num, z: Arr, c: Num) -> Arr {
    return a * x + b * z - c
}

#| The same, one element at a time.
def axpbz_loop(a: Num, x: Arr, b: Num, z: Arr, c: Num) -> Arr {
    n = int(len(x))
    y = zeros(n)
    for i to n {
        y[i] = a * x[i] + b * z[i] - c
    }
    return y
}

#| Pearson correlation of two runs of a flat array (SISSO's ffcorr).
def corr(d: Arr, o1: Int, o2: Int, n: Int) -> Num {
    x1 = d[o1:o1 + n]
    x2 = d[o2:o2 + n]
    m1 = sum(x1, ordered=true) / n
    m2 = sum(x2, ordered=true) / n
    s1 = sqrt(sum((x1 - m1) ^ 2, ordered=true) / n)
    s2 = sqrt(sum((x2 - m2) ^ 2, ordered=true) / n)
    return dot(x1 - m1, x2 - m2, ordered=true) / n / s1 / s2
}

#| The same, with loops.
def corr_loop(d: Arr, o1: Int, o2: Int, n: Int) -> Num {
    m1 = 0.0
    m2 = 0.0
    for i to n {
        m1 += d[o1 + i]
        m2 += d[o2 + i]
    }
    m1 = m1 / n
    m2 = m2 / n
    v1 = 0.0
    v2 = 0.0
    c = 0.0
    for i to n {
        v1 += pow(d[o1 + i] - m1, 2.0)
        v2 += pow(d[o2 + i] - m2, 2.0)
        c += (d[o1 + i] - m1) * (d[o2 + i] - m2)
    }
    return c / n / sqrt(v1 / n) / sqrt(v2 / n)
}

#| A view, then a write to its source: the reads before the write see the
#| old elements, the reads after see the new ones.
def view_then_write(d: Arr) -> Num {
    x = d[2:6]
    s = sum(x, ordered=true)
    d[3] = 100.0
    return s + sum(x, ordered=true) + sum(d[2:6], ordered=true)
}

#| The same, with loops: the copy holds the old elements.
def vtw_loop(d: Arr) -> Num {
    old = zeros(4)
    for i to 4 {
        old[i] = d[2 + i]
    }
    s = 0.0
    for i to 4 {
        s += old[i]
    }
    d[3] = 100.0
    t = 0.0
    for i to 4 {
        t += old[i]
    }
    u = 0.0
    for i to 4 {
        u += d[2 + i]
    }
    return s + t + u
}

#| A slice assigned and the source written in between: the copy is kept.
def copy_kept(d: Arr) -> Arr {
    x = d[0:3]
    d[0] = -1.0
    return 2.0 * x + 1.0
}

#| Loss for the gradient: a fused tree inside a reduction.
def loss(w: Arr, x: Arr, y: Arr) -> Num {
    r = 3.0 * w * x - y + 0.5
    return sum(r * r) + dot(w[1:4], x[0:3])
}

d = [sin(0.37 * i) + 0.01 * i for i to N]
x = d[0:100]
z = d[100:200]
y = axpbz(1.5, x, -0.25, z, 0.125)
show sum(y), dot(y, y)
r = d[::-3]
show len(r), sum(r), sum(d[::-3], ordered=true), prod(d[1:9] + 1.0), mean(d[0:N:7]), norm(d[5:50:5])
show max(d[10:20] * 2.0 - 1.0), min(abs(d[10:20]) + 0.5)
e = d[5:5]
show sum(e), sum(e, ordered=true), dot(e, e), len(2.0 * e + 1.0)
M = reshape(d[0:12], 3, 4)
show sum(M[1, :]), sum(M[:, 2] * 2.0), dot(M[0, 1:3], M[2, 2:4])
show sum(M[0:2] + 1.0), sum(M[:, 1:3] * M[:, 0:2])
show corr(d, 10, 150, 100)
show view_then_write(copy(d)), copy_kept(copy(d))
w = [0.1, -0.2, 0.3, 0.4, 0.5]
g = grad(loss, w, d[0:5], d[5:10])
show g
print(f"{sum(y):.17g} {corr(d, 10, 150, 100):.17g} {sum(d[::-3]):.17g} {mean(d[0:N:7]):.17g} {norm(d[5:50:5]):.17g} {g[2]:.17g}")
# broadcasting is left to the unfused operations: a row against a matrix
show sum(M * [1.0, 2.0, 3.0, 4.0] + 1.0)
test "fused results equal the element loops to the bit" {
    assert_true(sum(axpbz(1.5, x, -0.25, z, 0.125) - axpbz_loop(1.5, x, -0.25, z, 0.125)) == 0.0)
    assert_true(all(axpbz(1.5, x, -0.25, z, 0.125) == axpbz_loop(1.5, x, -0.25, z, 0.125)))
    assert_true(corr(d, 10, 150, 100) == corr_loop(d, 10, 150, 100))
    assert_true(sum(d[::-3], ordered=true) == sum(r, ordered=true))
    assert_true(view_then_write(copy(d)) == vtw_loop(copy(d)))
    assert_true(copy_kept(copy(d))[0] == 2.0 * d[0] + 1.0)
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors