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


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

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

CLOCK_PER_SEC: t.CDefine = 1000

test_passed: t.CInt = 0
test_failed: t.CInt = 0

def check(name: str, condition: t.CInt, detail: str):
    global test_passed, test_failed
    if condition:
        test_passed += 1
        printf("  [PASS] %s\n", name)
    else:
        test_failed += 1
        printf("  [FAIL] %s  -- %s\n", name, detail if detail else "")

def section(title: str):
    printf("\n=== %s ===\n", title)


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

def test_avx2_add_sub_mul_div():
    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)
    check("add[0]==6", out.data[0] == t.CDouble(6.0), "")
    check("add[1]==8", out.data[1] == t.CDouble(8.0), "")
    check("add[2]==10", out.data[2] == t.CDouble(10.0), "")
    check("add[3]==12", out.data[3] == t.CDouble(12.0), "")

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

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

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

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


def test_avx2_neg_abs_sqrt():
    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)
    check("neg[0]==-1", out.data[0] == t.CDouble(-1.0), "")
    check("neg[1]==4", out.data[1] == t.CDouble(4.0), "")
    check("neg[2]==-9", out.data[2] == t.CDouble(-9.0), "")
    check("neg[3]==16", out.data[3] == t.CDouble(16.0), "")

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

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

    a.delete()
    out.delete()


def test_avx2_min_max():
    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)
    check("min[0]==1", out.data[0] == t.CDouble(1.0), "")
    check("min[1]==2", out.data[1] == t.CDouble(2.0), "")
    check("min[2]==3", out.data[2] == t.CDouble(3.0), "")
    check("min[3]==4", out.data[3] == t.CDouble(4.0), "")

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

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


def test_avx2_scalar_ops():
    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)
    check("mul_scalar[0]==10", out.data[0] == t.CDouble(10.0), "")
    check("mul_scalar[1]==20", out.data[1] == t.CDouble(20.0), "")
    check("mul_scalar[3]==40", out.data[3] == t.CDouble(40.0), "")

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

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


def test_avx2_hsum_dot():
    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)
    check("hsum==10", out.data[0] == t.CDouble(10.0), "")

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

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


def test_avx2_fma():
    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)
    check("fma[0]==6", out.data[0] == t.CDouble(6.0), "")
    check("fma[1]==13", out.data[1] == t.CDouble(13.0), "")
    check("fma[2]==22", out.data[2] == t.CDouble(22.0), "")
    check("fma[3]==33", out.data[3] == t.CDouble(33.0), "")

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


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

def test_batch_array_ops():
    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)
    check("add_arr[0]==3", out.data[0] == t.CDouble(3.0), "")
    check("add_arr[15]==18", out.data[15] == t.CDouble(18.0), "")

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

    # 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))
    check("sqrt_arr[0]==2", out4.data[0] == t.CDouble(2.0), "")
    check("sqrt_arr[3]==5", out4.data[3] == t.CDouble(5.0), "")

    # 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))
    check("abs_arr[0]==1", out4.data[0] == t.CDouble(1.0), "")
    check("abs_arr[1]==2", out4.data[1] == t.CDouble(2.0), "")
    check("abs_arr[3]==4", out4.data[3] == t.CDouble(4.0), "")

    # 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)
    check("dot_arr==10", out1.data[0] == t.CDouble(10.0), "")

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


# ============================================================
# 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():
    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")

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

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


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

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

    arena = calloc(4096 * 1024, 1)  # 4MB arena
    pool = mpool.MPool(arena, 4096 * 1024, 0)  # bump allocator

    printf("=== ViperSIMD Test Suite ===\n")

    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()

    printf("\n=== Results: %d passed, %d failed ===\n", test_passed, test_failed)
    return 0