"""
On-device verification for the cross-compiled numpy wheel (OpenBLAS).

Run after installing:
    pip install numpy-2.5.2-cp312-cp312-linux_<ABI>.whl
        where <ABI> = aarch64 (real device) or x86_64 (emulator)
    (this Android build uses the "linux" platform tag for numpy)

Usage:
    python Test_NumPy.py [--quick]

Exit code 0 = everything required PASSed.
Sections marked [SKIP] are optional (e.g. need Pillow installed).

Generated by RIMI
"""
import os
import sys
import tempfile

RESULTS = []


def test(name, fn):
    try:
        fn()
        RESULTS.append((name, "PASS", None))
    except NotImplementedError as exc:
        RESULTS.append((name, "SKIP", str(exc)))
    except Exception as exc:
        RESULTS.append((name, "FAIL", "%s: %s" % (type(exc).__name__, exc)))
        print("    ! %s -> %s: %s" % (name, type(exc).__name__, exc))


def section(title):
    print("=" * 60)
    print(title)
    print("=" * 60)


WORKDIR = None


def workdir():
    global WORKDIR
    if WORKDIR is None:
        candidates = [os.environ.get("TMPDIR") or "", tempfile.gettempdir(),
                      "/storage/emulated/0/Download", os.getcwd()]
        for base in candidates:
            if not base:
                continue
            try:
                d = os.path.join(base, "test_numpy_tmp")
                os.makedirs(d, exist_ok=True)
                with open(os.path.join(d, "_probe"), "w") as fh:
                    fh.write("ok")
                WORKDIR = d
                break
            except OSError:
                continue
        if WORKDIR is None:
            WORKDIR = "."
    return WORKDIR


# ---------------------------------------------------------------------------
# 1. import / version
# ---------------------------------------------------------------------------
def import_numpy():
    import numpy as np
    print("    numpy", np.__version__)
    assert np.__version__.split(".")[0] == "2", np.__version__
    assert callable(np.show_config)


def array_basics():
    import numpy as np
    a = np.array([[1, 2, 3], [4, 5, 6]])
    assert a.shape == (2, 3)
    assert a.ndim == 2
    assert a.size == 6
    assert a.dtype == np.dtype("int64")
    assert a.itemsize == 8
    assert a.nbytes == 48


# ---------------------------------------------------------------------------
# 2. array creation
# ---------------------------------------------------------------------------
def creation():
    import numpy as np
    assert np.array([1, 2, 3]).tolist() == [1, 2, 3]
    assert np.zeros((2, 2)).sum() == 0
    assert np.ones((2, 2)).sum() == 4
    assert np.full((2,), 7.5).tolist() == [7.5, 7.5]
    assert np.eye(3).shape == (3, 3)
    assert np.arange(5).tolist() == [0, 1, 2, 3, 4]
    assert np.linspace(0, 1, 5).shape == (5,)
    assert len(np.logspace(1, 3, 3)) == 3


def random_rng():
    import numpy as np
    rng = np.random.default_rng(42)  # seeded -> reproducible
    r1 = np.random.default_rng(42)
    r2 = np.random.default_rng(42)
    assert (r1.random(5) == r2.random(5)).all()   # same seed, same stream
    assert rng.random((3, 3)).shape == (3, 3)
    assert rng.integers(0, 10, size=(2, 5)).shape == (2, 5)
    assert rng.normal(0, 1, size=(4,)).shape == (4,)


# ---------------------------------------------------------------------------
# 3. dtypes / casting
# ---------------------------------------------------------------------------
def dtypes():
    import numpy as np
    assert np.array([1, 2, 3], dtype=np.uint8).dtype == np.dtype("uint8")
    assert np.array([1.0, 2.0]).astype(np.float32).dtype == np.dtype("float32")
    assert np.array([1, 2, 3]).astype("f4").dtype == np.dtype("float32")
    for s in ("i1", "i2", "i4", "i8", "u1", "u2", "u4", "u8", "f4", "f8"):
        assert np.dtype(s)


def overflow():
    import numpy as np
    # uint8 arithmetic wraps around
    assert (np.array([200], np.uint8) + np.array([100], np.uint8))[0] == 44
    # int division floors, true division gives float
    assert np.array([5]) // 2 == np.array([2])
    assert np.array([5]) / 2 == np.array([2.5])


# ---------------------------------------------------------------------------
# 4. indexing / slicing / masking
# ---------------------------------------------------------------------------
def indexing():
    import numpy as np
    a = np.arange(12).reshape(3, 4)
    assert a[0].tolist() == [0, 1, 2, 3]
    assert a[0, 2] == 2
    assert a[:, 1].tolist() == [1, 5, 9]
    assert a[1:, :2].tolist() == [[4, 5], [8, 9]]
    assert a[-1].tolist() == [8, 9, 10, 11]
    assert a[::2].tolist() == [[0, 1, 2, 3], [8, 9, 10, 11]]


def masking():
    import numpy as np
    a = np.arange(12).reshape(3, 4)
    assert (a[a > 5] > 5).all()
    assert len(a[(a > 2) & (a < 8)]) == 5
    assert (a[a % 2 == 0] % 2 == 0).all()
    m = a.copy()
    m[m < 5] = 0
    assert m.min() == 0
    m[:, 0] = -1
    assert (m[:, 0] == -1).all()


def fancy_indexing():
    import numpy as np
    a = np.arange(12).reshape(3, 4)
    assert a[[0, 2]].shape == (2, 4)
    assert a[:, np.array([3, 1])].shape == (3, 2)


# ---------------------------------------------------------------------------
# 5. shapes / broadcasting
# ---------------------------------------------------------------------------
def reshaping():
    import numpy as np
    a = np.arange(24)
    assert a.reshape(4, 6).shape == (4, 6)
    assert a.reshape(2, 3, 4).shape == (2, 3, 4)
    assert a.reshape(-1, 6).shape == (4, 6)
    assert a.ravel().shape == (24,)
    assert a.flatten().shape == (24,)
    assert a.reshape(4, 6).T.shape == (6, 4)
    v = np.array([1, 2, 3])
    assert v[np.newaxis, :].shape == (1, 3)
    assert v[:, np.newaxis].shape == (3, 1)


def broadcasting():
    import numpy as np
    m = np.ones((3, 4))
    assert (m + 1 == 2).all()
    assert (m * np.array([10, 20, 30, 40])).shape == (3, 4)
    assert (m + np.array([[1], [2], [3]])).shape == (3, 4)
    # (3,1) * (1,4) -> (3,4)
    out = np.array([[1], [2], [3]]) * np.array([[1, 2, 3, 4]])
    assert out.shape == (3, 4)


# ---------------------------------------------------------------------------
# 6. math / reductions
# ---------------------------------------------------------------------------
def elementwise():
    import numpy as np
    a = np.array([1., 2., 3., 4.])
    assert (a + 1).tolist() == [2., 3., 4., 5.]
    assert (a ** 2).tolist() == [1., 4., 9., 16.]
    assert np.sqrt(np.array([4., 9.])).tolist() == [2., 3.]
    assert np.clip(a, 1.5, 3.5).tolist() == [1.5, 2., 3., 3.5]
    assert np.maximum(a, 2).tolist() == [2., 2., 3., 4.]


def reductions():
    import numpy as np
    a = np.array([1., 2., 3., 4.])
    assert a.sum() == 10
    assert a.mean() == 2.5
    assert a.min() == 1 and a.max() == 4
    assert a.prod() == 24
    assert a.argmax() == 3 and a.argmin() == 0
    assert np.median(a) == 2.5
    assert np.percentile(a, 50) == 2.5
    m = np.arange(6).reshape(2, 3)
    assert m.sum(axis=0).tolist() == [3, 5, 7]
    assert m.sum(axis=1).tolist() == [3, 12]


def comparisons():
    import numpy as np
    a = np.array([1., 2., 3., 4.])
    assert (a > 2).tolist() == [False, False, True, True]
    assert bool(np.any(a > 2)) is True
    assert bool(np.all(a > 2)) is False
    assert np.count_nonzero(a > 2) == 2


# ---------------------------------------------------------------------------
# 7. linear algebra (OpenBLAS accelerated)
# ---------------------------------------------------------------------------
def matmul():
    import numpy as np
    a = np.array([[1., 2.], [3., 4.]])
    b = np.array([[5., 6.], [7., 8.]])
    assert (a @ b).tolist() == [[19., 22.], [43., 50.]]
    assert np.matmul(a, b).tolist() == (a @ b).tolist()
    assert a.dot(b).tolist() == (a @ b).tolist()


def linalg():
    import numpy as np
    a = np.array([[4., 2.], [1., 3.]])
    inv = np.linalg.inv(a)
    ident = inv @ a
    assert np.allclose(ident, np.eye(2), atol=1e-10)
    assert abs(np.linalg.det(a) - 10.0) < 1e-10
    x = np.linalg.solve(a, np.array([6., 4.]))
    assert np.allclose(a @ x, [6., 4.])
    assert np.linalg.norm(np.array([3., 4.])) == 5.0
    w, v = np.linalg.eig(a)
    assert w.shape == (2,)
    assert v.shape == (2, 2)


def point_transform():
    import numpy as np
    M = np.array([[1., 0., 10.], [0., 1., 20.], [0., 0., 1.]])
    p = np.array([5., 6., 1.])
    out = M @ p
    assert out.tolist() == [15., 26., 1.]


# ---------------------------------------------------------------------------
# 8. stacking / splitting
# ---------------------------------------------------------------------------
def stacking():
    import numpy as np
    a = np.array([1, 2, 3])
    b = np.array([4, 5, 6])
    assert np.concatenate((a, b)).tolist() == [1, 2, 3, 4, 5, 6]
    assert np.stack((a, b)).shape == (2, 3)
    assert np.vstack((a, b)).shape == (2, 3)
    assert np.hstack((a, b)).shape == (6,)
    m1 = np.ones((2, 2))
    m2 = np.zeros((2, 2))
    assert np.vstack((m1, m2)).shape == (4, 2)
    assert np.hstack((m1, m2)).shape == (2, 4)


def splitting():
    import numpy as np
    x = np.arange(10)
    parts = np.split(x, 2)
    assert len(parts) == 2 and parts[0].tolist() == [0, 1, 2, 3, 4]
    assert len(np.array_split(x, 3)) == 3
    m = np.ones((4, 4))
    assert len(np.hsplit(m, 2)) == 2
    assert len(np.vsplit(m, 2)) == 2


# ---------------------------------------------------------------------------
# 9. save / load files
# ---------------------------------------------------------------------------
def save_load_npy():
    import numpy as np
    a = np.arange(12).reshape(3, 4)
    p = os.path.join(workdir(), "a.npy")
    np.save(p, a)
    b = np.load(p)
    assert (b == a).all()


def save_load_npz():
    import numpy as np
    a = np.arange(12).reshape(3, 4)
    p = os.path.join(workdir(), "data.npz")
    np.savez(p, x=a, y=a * 2)
    d = np.load(p)
    assert (d["x"] == a).all()
    assert (d["y"] == a * 2).all()
    d.close()


def save_load_text():
    import numpy as np
    a = np.arange(12).reshape(3, 4)
    p = os.path.join(workdir(), "a.csv")
    np.savetxt(p, a, delimiter=",")
    c = np.loadtxt(p, delimiter=",")
    assert c.dtype == np.float64
    assert c.shape == (3, 4)


def save_load_binary():
    import numpy as np
    a = np.arange(12).reshape(3, 4)
    p = os.path.join(workdir(), "a.bin")
    a.tofile(p)
    b = np.fromfile(p, dtype=np.int64)
    assert b.tolist() == list(range(12))


# ---------------------------------------------------------------------------
# 11. terminal printing
# ---------------------------------------------------------------------------
def printing():
    import numpy as np
    a = np.arange(12).reshape(3, 4)
    a.tolist()  # nested python lists
    prev = np.get_printoptions()
    np.set_printoptions(precision=2, threshold=20, edgeitems=3, linewidth=120,
                        suppress=True)
    print(a)
    np.set_printoptions(**prev)


# ---------------------------------------------------------------------------
# 12. everyday snippets
# ---------------------------------------------------------------------------
def snippets():
    import numpy as np
    x = np.array([3., 1., 2., 0.])
    n = (x - x.min()) / (x.max() - x.min())
    assert n.min() == 0 and n.max() == 1
    z = (x - x.mean()) / x.std()
    assert abs(z.mean()) < 1e-12
    cats = np.array([0, 2, 1, 2, 0])
    onehot = np.eye(3)[cats]
    assert onehot.shape == (5, 3)
    assert np.diag(np.arange(9).reshape(3, 3)).tolist() == [0, 4, 8]
    rng = np.random.default_rng(7)
    values, edges = np.histogram(rng.normal(size=1000), bins=20)
    assert len(values) == 20 and len(edges) == 21
    m = rng.random((5, 8))
    assert m.argmax(axis=1).shape == (5,)
    signal = np.array([1., 2., 3., 2., 1.])
    kernel = np.ones(3) / 3
    smooth = np.convolve(signal, kernel, mode="same")
    assert smooth.shape == signal.shape


def elapsed_time():
    import numpy as np
    import time
    t0 = time.perf_counter()
    big = np.arange(1_000_000)
    out = big * 2
    elapsed = time.perf_counter() - t0
    assert out.shape == big.shape
    print("    %.4f s for 1M element multiply" % elapsed)


# ---------------------------------------------------------------------------
def main():
    quick = "--quick" in sys.argv

    section("1. numpy import / version")
    test("import numpy (2.x)", import_numpy)
    test("array basics (shape/ndim/size/dtype)", array_basics)

    section("2. array creation")
    test("creation helpers", creation)
    test("default_rng seeded random", random_rng)

    section("3. dtypes / casting")
    test("dtypes and casting", dtypes)
    test("uint8 overflow / division rules", overflow)

    section("4. indexing / masking")
    test("indexing and slicing", indexing)
    test("boolean masking + assignment", masking)
    test("fancy indexing", fancy_indexing)

    section("5. shapes / broadcasting")
    test("reshape / ravel / T / newaxis", reshaping)
    test("broadcasting rules", broadcasting)

    section("6. math / reductions")
    test("element-wise ufuncs", elementwise)
    test("reductions + axes", reductions)
    test("comparisons / any / all", comparisons)

    section("7. linear algebra")
    test("matrix multiply @", matmul)
    test("inv/det/solve/eig/norm", linalg)
    test("homography point transform", point_transform)

    section("8. stacking / splitting")
    test("concatenate / stack / vstack / hstack", stacking)
    test("split / array_split / hsplit / vsplit", splitting)

    section("9. save / load files")
    test("npy roundtrip", save_load_npy)
    test("npz roundtrip", save_load_npz)
    test("savetxt / loadtxt", save_load_text)
    test("tofile / fromfile", save_load_binary)

    section("10. terminal printing")
    test("print options + tolist", printing)

    section("11. everyday snippets")
    test("normalize / zscore / one-hot / histogram", snippets)
    test("large-array perf sanity", elapsed_time)

    print()
    print("=" * 60)
    print("SUMMARY")
    print("=" * 60)
    fails = 0
    skips = 0
    for name, status, why in RESULTS:
        mark = "  OK" if status == "PASS" else (" SKIP" if status == "SKIP" else "FAIL")
        print("%s %s" % (mark, name))
        if why:
            print("        -> %s" % why)
        if status == "FAIL":
            fails += 1
        elif status == "SKIP":
            skips += 1
    print()
    passed = len(RESULTS) - fails - skips
    print("passed=%d skipped=%d failed=%d" % (passed, skips, fails))
    if fails:
        print("RESULT: FAILED")
    elif skips and not quick:
        print("RESULT: PASSED (with informational skips)")
    else:
        print("RESULT: PASSED")
    sys.exit(1 if fails else 0)


if __name__ == "__main__":
    main()