from stdint import *
import vipersimd
import numpy
from stdio import printf
import stdlib
from stdlib import calloc
import t, c
import memhub
import testcheck


# 全局内存池
arena: t.CUInt8T | t.CPtr = None
pool: memhub.MemPool | t.CPtr = None

# 外部函数: clock() 用于性能计时
def clock() -> t.CLong | t.CExtern: pass

CLOCK_PER_SEC: t.CDefine = 1000


# ============================================================
# 1. AVX2 double4 正确性测试
# ============================================================

def test_avx2_add_sub_mul_div():
    testcheck.section("AVX2 add/sub/mul/div 4d")
    a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    a.data[0] = t.CDouble(1.0)
    a.data[1] = t.CDouble(2.0)
    a.data[2] = t.CDouble(3.0)
    a.data[3] = t.CDouble(4.0)

    b: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    b.data[0] = t.CDouble(5.0)
    b.data[1] = t.CDouble(6.0)
    b.data[2] = t.CDouble(7.0)
    b.data[3] = t.CDouble(8.0)

    out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)

    # add
    vipersimd.simd_add4d(a.data, b.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(6.0), "add[0]==6", "add[0]==6")
    testcheck.check(out.data[1] == t.CDouble(8.0), "add[1]==8", "add[1]==8")
    testcheck.check(out.data[2] == t.CDouble(10.0), "add[2]==10", "add[2]==10")
    testcheck.check(out.data[3] == t.CDouble(12.0), "add[3]==12", "add[3]==12")

    # sub
    vipersimd.simd_sub4d(a.data, b.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(-4.0), "sub[0]==-4", "sub[0]==-4")
    testcheck.check(out.data[3] == t.CDouble(-4.0), "sub[3]==-4", "sub[3]==-4")

    # mul
    vipersimd.simd_mul4d(a.data, b.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(5.0), "mul[0]==5", "mul[0]==5")
    testcheck.check(out.data[1] == t.CDouble(12.0), "mul[1]==12", "mul[1]==12")
    testcheck.check(out.data[2] == t.CDouble(21.0), "mul[2]==21", "mul[2]==21")
    testcheck.check(out.data[3] == t.CDouble(32.0), "mul[3]==32", "mul[3]==32")

    # div
    vipersimd.simd_div4d(b.data, a.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(5.0), "div[0]==5", "div[0]==5")
    testcheck.check(out.data[1] == t.CDouble(3.0), "div[1]==3", "div[1]==3")

    a.delete()
    b.delete()
    out.delete()


def test_avx2_neg_abs_sqrt():
    testcheck.section("AVX2 neg/abs/sqrt 4d")
    a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    a.data[0] = t.CDouble(1.0)
    a.data[1] = t.CDouble(-4.0)
    a.data[2] = t.CDouble(9.0)
    a.data[3] = t.CDouble(-16.0)

    out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)

    # neg
    vipersimd.simd_neg4d(a.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(-1.0), "neg[0]==-1", "neg[0]==-1")
    testcheck.check(out.data[1] == t.CDouble(4.0), "neg[1]==4", "neg[1]==4")
    testcheck.check(out.data[2] == t.CDouble(-9.0), "neg[2]==-9", "neg[2]==-9")
    testcheck.check(out.data[3] == t.CDouble(16.0), "neg[3]==16", "neg[3]==16")

    # abs
    vipersimd.simd_abs4d(a.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(1.0), "abs[0]==1", "abs[0]==1")
    testcheck.check(out.data[1] == t.CDouble(4.0), "abs[1]==4", "abs[1]==4")
    testcheck.check(out.data[2] == t.CDouble(9.0), "abs[2]==9", "abs[2]==9")
    testcheck.check(out.data[3] == t.CDouble(16.0), "abs[3]==16", "abs[3]==16")

    # sqrt
    vipersimd.simd_sqrt4d(out.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(1.0), "sqrt[0]==1", "sqrt[0]==1")
    testcheck.check(out.data[1] == t.CDouble(2.0), "sqrt[1]==2", "sqrt[1]==2")
    testcheck.check(out.data[2] == t.CDouble(3.0), "sqrt[2]==3", "sqrt[2]==3")
    testcheck.check(out.data[3] == t.CDouble(4.0), "sqrt[3]==4", "sqrt[3]==4")

    a.delete()
    out.delete()


def test_avx2_min_max():
    testcheck.section("AVX2 min/max 4d")
    a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    a.data[0] = t.CDouble(1.0)
    a.data[1] = t.CDouble(8.0)
    a.data[2] = t.CDouble(3.0)
    a.data[3] = t.CDouble(6.0)

    b: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    b.data[0] = t.CDouble(5.0)
    b.data[1] = t.CDouble(2.0)
    b.data[2] = t.CDouble(7.0)
    b.data[3] = t.CDouble(4.0)

    out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)

    vipersimd.simd_min4d(a.data, b.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(1.0), "min[0]==1", "min[0]==1")
    testcheck.check(out.data[1] == t.CDouble(2.0), "min[1]==2", "min[1]==2")
    testcheck.check(out.data[2] == t.CDouble(3.0), "min[2]==3", "min[2]==3")
    testcheck.check(out.data[3] == t.CDouble(4.0), "min[3]==4", "min[3]==4")

    vipersimd.simd_max4d(a.data, b.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(5.0), "max[0]==5", "max[0]==5")
    testcheck.check(out.data[1] == t.CDouble(8.0), "max[1]==8", "max[1]==8")
    testcheck.check(out.data[2] == t.CDouble(7.0), "max[2]==7", "max[2]==7")
    testcheck.check(out.data[3] == t.CDouble(6.0), "max[3]==6", "max[3]==6")

    a.delete()
    b.delete()
    out.delete()


def test_avx2_scalar_ops():
    testcheck.section("AVX2 scalar broadcast ops")
    a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    a.data[0] = t.CDouble(1.0)
    a.data[1] = t.CDouble(2.0)
    a.data[2] = t.CDouble(3.0)
    a.data[3] = t.CDouble(4.0)

    s: numpy.ndarray | t.CPtr = numpy.zeros(pool, 1)
    s.data[0] = t.CDouble(10.0)

    out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)

    vipersimd.simd_mul_scalar4d(a.data, s.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(10.0), "mul_scalar[0]==10", "mul_scalar[0]==10")
    testcheck.check(out.data[1] == t.CDouble(20.0), "mul_scalar[1]==20", "mul_scalar[1]==20")
    testcheck.check(out.data[3] == t.CDouble(40.0), "mul_scalar[3]==40", "mul_scalar[3]==40")

    vipersimd.simd_add_scalar4d(a.data, s.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(11.0), "add_scalar[0]==11", "add_scalar[0]==11")
    testcheck.check(out.data[3] == t.CDouble(14.0), "add_scalar[3]==14", "add_scalar[3]==14")

    a.delete()
    s.delete()
    out.delete()


def test_avx2_hsum_dot():
    testcheck.section("AVX2 hsum/dot 4d")
    a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    a.data[0] = t.CDouble(1.0)
    a.data[1] = t.CDouble(2.0)
    a.data[2] = t.CDouble(3.0)
    a.data[3] = t.CDouble(4.0)

    b: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    b.data[0] = t.CDouble(5.0)
    b.data[1] = t.CDouble(6.0)
    b.data[2] = t.CDouble(7.0)
    b.data[3] = t.CDouble(8.0)

    out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 1)

    vipersimd.simd_hsum4d(a.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(10.0), "hsum==10", "hsum==10")

    vipersimd.simd_dot4d(a.data, b.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(70.0), "dot==70", "dot==70")

    a.delete()
    b.delete()
    out.delete()


def test_avx2_fma():
    testcheck.section("AVX2 FMA 4d")
    a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    a.data[0] = t.CDouble(1.0)
    a.data[1] = t.CDouble(2.0)
    a.data[2] = t.CDouble(3.0)
    a.data[3] = t.CDouble(4.0)

    b: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    b.data[0] = t.CDouble(5.0)
    b.data[1] = t.CDouble(6.0)
    b.data[2] = t.CDouble(7.0)
    b.data[3] = t.CDouble(8.0)

    c_val: numpy.ndarray | t.CPtr = numpy.ones(pool, 4)

    out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)

    vipersimd.simd_fma4d(a.data, b.data, c_val.data, out.data)
    testcheck.check(out.data[0] == t.CDouble(6.0), "fma[0]==6", "fma[0]==6")
    testcheck.check(out.data[1] == t.CDouble(13.0), "fma[1]==13", "fma[1]==13")
    testcheck.check(out.data[2] == t.CDouble(22.0), "fma[2]==22", "fma[2]==22")
    testcheck.check(out.data[3] == t.CDouble(33.0), "fma[3]==33", "fma[3]==33")

    a.delete()
    b.delete()
    c_val.delete()
    out.delete()


# ============================================================
# 2. 批量数组运算正确性测试
# ============================================================

def test_batch_array_ops():
    testcheck.section("Batch array operations")
    n: t.CSizeT = 16
    a: numpy.ndarray | t.CPtr = numpy.arange(pool, t.CDouble(1.0), t.CDouble(17.0), t.CDouble(1.0))
    b: numpy.ndarray | t.CPtr = numpy.full(pool, 16, t.CDouble(2.0))
    out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 16)

    # add array
    vipersimd.simd_add_array(a.data, b.data, out.data, n)
    testcheck.check(out.data[0] == t.CDouble(3.0), "add_arr[0]==3", "add_arr[0]==3")
    testcheck.check(out.data[15] == t.CDouble(18.0), "add_arr[15]==18", "add_arr[15]==18")

    # mul array
    vipersimd.simd_mul_array(a.data, b.data, out.data, n)
    testcheck.check(out.data[0] == t.CDouble(2.0), "mul_arr[0]==2", "mul_arr[0]==2")
    testcheck.check(out.data[3] == t.CDouble(8.0), "mul_arr[3]==8", "mul_arr[3]==8")

    # sqrt array
    c: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    c.data[0] = t.CDouble(4.0)
    c.data[1] = t.CDouble(9.0)
    c.data[2] = t.CDouble(16.0)
    c.data[3] = t.CDouble(25.0)
    out4: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    vipersimd.simd_sqrt_array(c.data, out4.data, t.CSizeT(4))
    testcheck.check(out4.data[0] == t.CDouble(2.0), "sqrt_arr[0]==2", "sqrt_arr[0]==2")
    testcheck.check(out4.data[3] == t.CDouble(5.0), "sqrt_arr[3]==5", "sqrt_arr[3]==5")

    # abs array
    d: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    d.data[0] = t.CDouble(-1.0)
    d.data[1] = t.CDouble(-2.0)
    d.data[2] = t.CDouble(3.0)
    d.data[3] = t.CDouble(-4.0)
    vipersimd.simd_abs_array(d.data, out4.data, t.CSizeT(4))
    testcheck.check(out4.data[0] == t.CDouble(1.0), "abs_arr[0]==1", "abs_arr[0]==1")
    testcheck.check(out4.data[1] == t.CDouble(2.0), "abs_arr[1]==2", "abs_arr[1]==2")
    testcheck.check(out4.data[3] == t.CDouble(4.0), "abs_arr[3]==4", "abs_arr[3]==4")

    # dot array
    e: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
    e.data[0] = t.CDouble(1.0)
    e.data[1] = t.CDouble(2.0)
    e.data[2] = t.CDouble(3.0)
    e.data[3] = t.CDouble(4.0)
    f: numpy.ndarray | t.CPtr = numpy.ones(pool, 4)
    out1: numpy.ndarray | t.CPtr = numpy.zeros(pool, 1)
    vipersimd.simd_dot_array(e.data, f.data, t.CSizeT(4), out1.data)
    testcheck.check(out1.data[0] == t.CDouble(10.0), "dot_arr==10", "dot_arr==10")

    a.delete()
    b.delete()
    out.delete()
    c.delete()
    out4.delete()
    d.delete()
    e.delete()
    f.delete()
    out1.delete()


# ============================================================
# 3. LLVMIR 精确标量运算测试
# ============================================================

def test_llvmir_scalar():
    testcheck.section("LLVMIR precise scalar ops")

    # sqrt
    r: t.CDouble = vipersimd.simd_sqrt(t.CDouble(4.0))
    testcheck.check(r == t.CDouble(2.0), "sqrt(4)==2", "sqrt(4)==2")

    # fabs
    r = vipersimd.simd_fabs(t.CDouble(-3.5))
    testcheck.check(r == t.CDouble(3.5), "fabs(-3.5)==3.5", "fabs(-3.5)==3.5")

    # floor
    r = vipersimd.simd_floor(t.CDouble(3.7))
    testcheck.check(r == t.CDouble(3.0), "floor(3.7)==3", "floor(3.7)==3")

    # ceil
    r = vipersimd.simd_ceil(t.CDouble(3.2))
    testcheck.check(r == t.CDouble(4.0), "ceil(3.2)==4", "ceil(3.2)==4")

    # round
    r = vipersimd.simd_round(t.CDouble(3.5))
    testcheck.check(r == t.CDouble(4.0), "round(3.5)==4", "round(3.5)==4")

    # trunc
    r = vipersimd.simd_trunc(t.CDouble(3.9))
    testcheck.check(r == t.CDouble(3.0), "trunc(3.9)==3", "trunc(3.9)==3")

    # fma: 2*3+1 = 7
    r = vipersimd.simd_fma(t.CDouble(2.0), t.CDouble(3.0), t.CDouble(1.0))
    testcheck.check(r == t.CDouble(7.0), "fma(2,3,1)==7", "fma(2,3,1)==7")

    # copysign
    r = vipersimd.simd_copysign(t.CDouble(3.0), t.CDouble(-1.0))
    testcheck.check(r == t.CDouble(-3.0), "copysign(3,-1)==-3", "copysign(3,-1)==-3")

    # neg
    r = vipersimd.simd_neg(t.CDouble(1.5))
    testcheck.check(r == t.CDouble(-1.5), "neg(1.5)==-1.5", "neg(1.5)==-1.5")

    # minnum
    r = vipersimd.simd_minnum(t.CDouble(2.0), t.CDouble(3.0))
    testcheck.check(r == t.CDouble(2.0), "minnum(2,3)==2", "minnum(2,3)==2")

    # maxnum
    r = vipersimd.simd_maxnum(t.CDouble(2.0), t.CDouble(3.0))
    testcheck.check(r == t.CDouble(3.0), "maxnum(2,3)==3", "maxnum(2,3)==3")

    # exp
    r = vipersimd.simd_exp(t.CDouble(1.0))
    testcheck.check(r > t.CDouble(2.71) and r < t.CDouble(2.72), "exp(1)~=2.718", "exp(1)~=2.718")

    # log
    r = vipersimd.simd_log(t.CDouble(2.718281828459045))
    testcheck.check(r > t.CDouble(0.999) and r < t.CDouble(1.001), "log(e)~=1", "log(e)~=1")

    # sin
    r = vipersimd.simd_sin(t.CDouble(0.0))
    testcheck.check(r > t.CDouble(-0.001) and r < t.CDouble(0.001), "sin(0)~=0", "sin(0)~=0")

    # cos
    r = vipersimd.simd_cos(t.CDouble(0.0))
    testcheck.check(r > t.CDouble(0.999) and r < t.CDouble(1.001), "cos(0)~=1", "cos(0)~=1")

    # pow
    r = vipersimd.simd_pow(t.CDouble(2.0), t.CDouble(10.0))
    testcheck.check(r == t.CDouble(1024.0), "pow(2,10)==1024", "pow(2,10)==1024")


# ============================================================
# 4. 性能基准测试: SIMD vs 标量
# ============================================================

N_ELEM: t.CDefine = 40000
N_ITER: t.CDefine = 50000

def scalar_add_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
    i: t.CSizeT = 0
    da: t.CDouble | t.CPtr = a
    db: t.CDouble | t.CPtr = b
    dout: t.CDouble | t.CPtr = out
    while i < n:
        dout[i] = da[i] + db[i]
        i += 1

def scalar_mul_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
    i: t.CSizeT = 0
    da: t.CDouble | t.CPtr = a
    db: t.CDouble | t.CPtr = b
    dout: t.CDouble | t.CPtr = out
    while i < n:
        dout[i] = da[i] * db[i]
        i += 1

def scalar_sqrt_array(a: t.CPtr, out: t.CPtr, n: t.CSizeT):
    i: t.CSizeT = 0
    da: t.CDouble | t.CPtr = a
    dout: t.CDouble | t.CPtr = out
    while i < n:
        dout[i] = vipersimd.simd_sqrt(da[i])
        i += 1

def benchmark_simd_vs_scalar():
    testcheck.section("Benchmark: SIMD vs Scalar")

    a: numpy.ndarray | t.CPtr = numpy.full(pool, N_ELEM, t.CDouble(1.5))
    b: numpy.ndarray | t.CPtr = numpy.full(pool, N_ELEM, t.CDouble(2.5))
    out: numpy.ndarray | t.CPtr = numpy.zeros(pool, N_ELEM)

    # --- SIMD add benchmark ---
    t0: t.CLong = clock()
    iter: t.CInt = 0
    while iter < N_ITER:
        vipersimd.simd_add_array(a.data, b.data, out.data, t.CSizeT(N_ELEM))
        iter += 1
    t1: t.CLong = clock()
    simd_add_ms: t.CLong = t1 - t0

    # --- Scalar add benchmark ---
    t0 = clock()
    iter = 0
    while iter < N_ITER:
        scalar_add_array(a.data, b.data, out.data, t.CSizeT(N_ELEM))
        iter += 1
    t1 = clock()
    scalar_add_ms: t.CLong = t1 - t0

    printf("  ADD:  SIMD=%ldms  Scalar=%ldms", simd_add_ms, scalar_add_ms)
    if scalar_add_ms > 0:
        printf("  speedup=%.1fx\n", t.CDouble(scalar_add_ms) / t.CDouble(simd_add_ms))
    else:
        printf("  speedup=N/A\n")

    # --- SIMD mul benchmark ---
    t0 = clock()
    iter = 0
    while iter < N_ITER:
        vipersimd.simd_mul_array(a.data, b.data, out.data, t.CSizeT(N_ELEM))
        iter += 1
    t1 = clock()
    simd_mul_ms: t.CLong = t1 - t0

    # --- Scalar mul benchmark ---
    t0 = clock()
    iter = 0
    while iter < N_ITER:
        scalar_mul_array(a.data, b.data, out.data, t.CSizeT(N_ELEM))
        iter += 1
    t1 = clock()
    scalar_mul_ms: t.CLong = t1 - t0

    printf("  MUL:  SIMD=%ldms  Scalar=%ldms", simd_mul_ms, scalar_mul_ms)
    if scalar_mul_ms > 0:
        printf("  speedup=%.1fx\n", t.CDouble(scalar_mul_ms) / t.CDouble(simd_mul_ms))
    else:
        printf("  speedup=N/A\n")

    # --- SIMD sqrt benchmark ---
    t0 = clock()
    iter = 0
    while iter < N_ITER:
        vipersimd.simd_sqrt_array(a.data, out.data, t.CSizeT(N_ELEM))
        iter += 1
    t1 = clock()
    simd_sqrt_ms: t.CLong = t1 - t0

    # --- Scalar sqrt benchmark ---
    t0 = clock()
    iter = 0
    while iter < N_ITER:
        scalar_sqrt_array(a.data, out.data, t.CSizeT(N_ELEM))
        iter += 1
    t1 = clock()
    scalar_sqrt_ms: t.CLong = t1 - t0

    printf("  SQRT: SIMD=%ldms  Scalar=%ldms", simd_sqrt_ms, scalar_sqrt_ms)
    if scalar_sqrt_ms > 0:
        printf("  speedup=%.1fx\n", t.CDouble(scalar_sqrt_ms) / t.CDouble(simd_sqrt_ms))
    else:
        printf("  speedup=N/A\n")

    # 验证结果正确性
    testcheck.check(out.data[0] > t.CDouble(0.0), "bench result[0] valid", "bench result[0] valid")

    a.delete()
    b.delete()
    out.delete()


# ============================================================
# Main
# ============================================================

def main() -> t.CInt:
    global arena, pool

    arena = calloc(4096 * 1024, 1)  # 4MB arena
    pool = memhub.MemPool(arena, 4096 * 1024)  # bump allocator

    testcheck.begin("ViperSIMD Test Suite")

    test_avx2_add_sub_mul_div()
    test_avx2_neg_abs_sqrt()
    test_avx2_min_max()
    test_avx2_scalar_ops()
    test_avx2_hsum_dot()
    test_avx2_fma()
    test_batch_array_ops()
    test_llvmir_scalar()
    benchmark_simd_vs_scalar()

    return testcheck.end()
