import time

import numpy

import xsor


def numpy_add() -> None:
    a = numpy.array(range(100000), dtype=numpy.float32)
    b = numpy.array(range(100000), dtype=numpy.float32)
    a + b


def numpy_mul() -> None:
    a = numpy.array(range(100000), dtype=numpy.float32)
    b = numpy.array(range(100000), dtype=numpy.float32)
    a * b


def numpy_sum() -> None:
    a = numpy.array(range(100000), dtype=numpy.float32)
    a.sum()


def numpy_matmul() -> None:
    a = numpy.array(range(100000), dtype=numpy.float32).reshape((100, 1000))
    b = numpy.array(range(100000), dtype=numpy.float32).reshape((1000, 100))
    a @ b


def numpy_add_scalar() -> None:
    a = numpy.array(range(100000), dtype=numpy.float32)
    a + 1.0


def xsor_add() -> None:
    a = xsor.tensor(range(100000))
    b = xsor.tensor(range(100000))
    a + b


def xsor_mul() -> None:
    a = xsor.tensor(range(100000))
    b = xsor.tensor(range(100000))
    a * b


def xsor_sum() -> None:
    a = xsor.tensor(range(100000))
    a.sum()


def xsor_matmul() -> None:
    a = xsor.tensor(range(100000), (100, 1000))
    b = xsor.tensor(range(100000), (1000, 100))
    a @ b


def xsor_add_scalar() -> None:
    a = xsor.tensor(range(100000))
    a + 1.0


N = 1000

for title, numpy_fn, xsor_fn in (
    ("ADD", numpy_add, xsor_add),
    ("MUL", numpy_mul, xsor_mul),
    ("SUM", numpy_sum, xsor_sum),
    ("MATMUL", numpy_matmul, xsor_matmul),
    ("ADD_SCALAR", numpy_add_scalar, xsor_add_scalar),
):
    print(f"[ {title} ]")

    print("  numpy ... ", end="")
    start = time.time()
    for _ in range(N):
        numpy_fn()
    end = time.time()
    print(f"{end - start:.3f}s")

    print("  xsor ... ", end="")
    start = time.time()
    for _ in range(N):
        xsor_fn()
    end = time.time()
    print(f"{end - start:.3f}s")

    print()