"""
test_allay_rt.py  —  AllayRT test suite
========================================
Tests the full pipeline:
  1. Build a toy ARTModel manually
  2. Save → load round-trip
  3. Compile and run inference
  4. Test individual ops
  5. Test the public API (art.read_model / art.compile_model / compiled.run)
"""

import sys
import numpy as np
import tempfile
import os

# ── Add parent dir to path if running from within AllayRT/ ──
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from AllayRT.model import ARTModel, save_art, load_art
from AllayRT.core import run_graph, OP_TABLE
from AllayRT.runtime import AllayRT


PASS = "✓"
FAIL = "✗"
results = []


def check(name, fn):
    try:
        fn()
        print(f"  {PASS} {name}")
        results.append((name, True, None))
    except Exception as e:
        print(f"  {FAIL} {name}  →  {e}")
        results.append((name, False, str(e)))


# ──────────────────────────────────────────────
#  Helper: build a tiny linear model in memory
#  y = ReLU(X @ W.T + b)
# ──────────────────────────────────────────────

def build_linear_model() -> ARTModel:
    m = ARTModel()
    m.meta = {"name": "tiny_linear", "opset_version": 17}
    m.inputs  = [{"name": "input", "dtype": "float32", "shape": [1, 4]}]
    m.outputs = [{"name": "output", "dtype": "float32", "shape": [1, 3]}]

    W = np.array([[1, 0, -1, 0.5],
                  [0, 1,  0, -1],
                  [0.5, 0.5, 0.5, 0.5]], dtype=np.float32)
    b = np.array([0.1, -0.1, 0.2], dtype=np.float32)
    m.tensors = {"W": W, "b": b}

    m.nodes = [
        {
            "op_type": "Gemm",
            "name": "gemm0",
            "domain": "",
            "inputs": ["input", "W", "b"],
            "outputs": ["gemm_out"],
            "attrs": {"transB": 1},
        },
        {
            "op_type": "Relu",
            "name": "relu0",
            "domain": "",
            "inputs": ["gemm_out"],
            "outputs": ["output"],
            "attrs": {},
        },
    ]
    return m


# ──────────────────────────────────────────────
#  Tests
# ──────────────────────────────────────────────

print("\n═══ AllayRT Test Suite ═══\n")

# 1. Build model
print("[1] Model construction")
model = None
def t_build():
    global model
    model = build_linear_model()
    assert model.meta["name"] == "tiny_linear"
    assert len(model.nodes) == 2
    assert "W" in model.tensors

check("Build ARTModel", t_build)


# 2. Save / load round-trip
print("\n[2] Serialization round-trip")
tmp_path = tempfile.mktemp(suffix=".art")

def t_save():
    save_art(model, tmp_path)
    assert os.path.exists(tmp_path)
    assert os.path.getsize(tmp_path) > 0
check("Save .art file", t_save)

loaded = None
def t_load():
    global loaded
    loaded = load_art(tmp_path)
    assert loaded.meta["name"] == "tiny_linear"
    assert len(loaded.nodes) == 2
    np.testing.assert_array_equal(loaded.tensors["W"], model.tensors["W"])
    np.testing.assert_array_equal(loaded.tensors["b"], model.tensors["b"])
check("Load .art file", t_load)

def t_checksum():
    # Corrupt the file and verify checksum failure
    with open(tmp_path, "r+b") as f:
        f.seek(20)
        f.write(b"\xff\xff\xff\xff")
    try:
        load_art(tmp_path)
        raise AssertionError("Should have raised ValueError")
    except ValueError as e:
        assert "Checksum" in str(e)
check("Checksum validation", t_checksum)

os.unlink(tmp_path)


# 3. Individual op tests
print("\n[3] Op correctness")

x = np.array([[1.0, -2.0, 3.0, -4.0]], dtype=np.float32)

def t_relu():
    [y] = OP_TABLE["Relu"]([x], {}, None)
    expected = np.array([[1, 0, 3, 0]], dtype=np.float32)
    np.testing.assert_array_almost_equal(y, expected)
check("Relu", t_relu)

def t_sigmoid():
    [y] = OP_TABLE["Sigmoid"]([np.array([[0.0]])], {}, None)
    np.testing.assert_almost_equal(y, 0.5, decimal=5)
check("Sigmoid", t_sigmoid)

def t_softmax():
    inp = np.array([[1.0, 2.0, 3.0]])
    [y] = OP_TABLE["Softmax"]([inp], {"axis": -1}, None)
    np.testing.assert_almost_equal(y.sum(), 1.0, decimal=5)
    assert y[0, 2] > y[0, 1] > y[0, 0]
check("Softmax", t_softmax)

def t_gemm():
    A = np.eye(3, dtype=np.float32)
    B = np.eye(3, dtype=np.float32) * 2
    [y] = OP_TABLE["Gemm"]([A, B], {}, None)
    np.testing.assert_array_almost_equal(y, B)
check("Gemm (identity * 2)", t_gemm)

def t_matmul():
    A = np.ones((2, 3), dtype=np.float32)
    B = np.ones((3, 4), dtype=np.float32)
    [y] = OP_TABLE["MatMul"]([A, B], {}, None)
    assert y.shape == (2, 4)
    np.testing.assert_array_almost_equal(y, np.full((2,4), 3.0))
check("MatMul", t_matmul)

def t_layernorm():
    inp = np.random.randn(2, 8).astype(np.float32)
    scale = np.ones(8, dtype=np.float32)
    bias  = np.zeros(8, dtype=np.float32)
    [y] = OP_TABLE["LayerNormalization"]([inp, scale, bias], {"axis": -1}, None)
    # Each row should have ~0 mean and ~1 std
    np.testing.assert_almost_equal(y.mean(axis=-1), [0, 0], decimal=5)
check("LayerNormalization", t_layernorm)

def t_reshape():
    arr = np.arange(12, dtype=np.float32)
    shape = np.array([3, 4], dtype=np.int64)
    [y] = OP_TABLE["Reshape"]([arr, shape], {}, None)
    assert y.shape == (3, 4)
check("Reshape", t_reshape)

def t_concat():
    a = np.ones((2, 3), dtype=np.float32)
    b = np.ones((2, 3), dtype=np.float32) * 2
    [y] = OP_TABLE["Concat"]([a, b], {"axis": 0}, None)
    assert y.shape == (4, 3)
check("Concat", t_concat)

def t_reducemean():
    arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
    [y] = OP_TABLE["ReduceMean"]([arr], {"axes": [1], "keepdims": 1}, None)
    np.testing.assert_array_almost_equal(y, [[2.0], [5.0]])
check("ReduceMean", t_reducemean)

def t_clip():
    arr = np.array([-5, 0, 5, 10], dtype=np.float32)
    lo  = np.array(0.0, dtype=np.float32)
    hi  = np.array(6.0, dtype=np.float32)
    [y] = OP_TABLE["Clip"]([arr, lo, hi], {}, None)
    np.testing.assert_array_equal(y, [0, 0, 5, 6])
check("Clip", t_clip)

def t_batchnorm():
    X     = np.random.randn(2, 4, 8, 8).astype(np.float32)
    scale = np.ones(4, dtype=np.float32)
    bias  = np.zeros(4, dtype=np.float32)
    mean  = np.zeros(4, dtype=np.float32)
    var   = np.ones(4, dtype=np.float32)
    [y] = OP_TABLE["BatchNormalization"]([X, scale, bias, mean, var], {}, None)
    assert y.shape == X.shape
check("BatchNormalization", t_batchnorm)

def t_maxpool():
    X = np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4)
    [y] = OP_TABLE["MaxPool"]([X], {"kernel_shape": [2, 2], "strides": [2, 2]}, None)
    assert y.shape == (1, 1, 2, 2)
    np.testing.assert_array_equal(y[0, 0], [[5, 7], [13, 15]])
check("MaxPool 2x2 stride 2", t_maxpool)

def t_global_avgpool():
    X = np.ones((2, 4, 8, 8), dtype=np.float32)
    [y] = OP_TABLE["GlobalAveragePool"]([X], {}, None)
    assert y.shape == (2, 4, 1, 1)
check("GlobalAveragePool", t_global_avgpool)


# 4. Full graph execution
print("\n[4] Graph execution")

model2 = build_linear_model()

def t_graph_run():
    x_in = np.array([[1.0, -1.0, 2.0, -2.0]], dtype=np.float32)
    results = run_graph(
        model2.nodes,
        model2.tensors,
        {"input": x_in},
        ["output"],
    )
    y = results["output"]
    assert y.shape == (1, 3)
    assert (y >= 0).all(), "ReLU output must be non-negative"
check("Full Gemm→Relu graph", t_graph_run)


# 5. Public AllayRT API
print("\n[5] AllayRT public API")
tmp2 = tempfile.mktemp(suffix=".art")
save_art(build_linear_model(), tmp2)

art_rt = AllayRT()

def t_api_read():
    m = art_rt.read_model(tmp2)
    assert isinstance(m, ARTModel)
    assert m.meta["name"] == "tiny_linear"
check("art.read_model()", t_api_read)

def t_api_compile():
    m = art_rt.read_model(tmp2)
    cm = art_rt.compile_model(m, "CPU")
    assert cm.device == "CPU"
    assert cm.input_names == ["input"]
    assert cm.output_names == ["output"]
check("art.compile_model(model, 'CPU')", t_api_compile)

def t_api_run_array():
    m  = art_rt.read_model(tmp2)
    cm = art_rt.compile_model(m, "CPU")
    x  = np.array([[1.0, -1.0, 2.0, -2.0]], dtype=np.float32)
    y  = cm.run(x)
    assert isinstance(y, np.ndarray)
    assert y.shape == (1, 3)
    assert (y >= 0).all()
check("compiled_model.run(np.ndarray)", t_api_run_array)

def t_api_run_dict():
    m  = art_rt.read_model(tmp2)
    cm = art_rt.compile_model(m, "CPU")
    x  = np.array([[1.0, -1.0, 2.0, -2.0]], dtype=np.float32)
    y  = cm.run({"input": x}, return_dict=True)
    assert isinstance(y, dict)
    assert "output" in y
check("compiled_model.run(dict, return_dict=True)", t_api_run_dict)

def t_api_supported_ops():
    ops = AllayRT.supported_ops()
    assert "Relu" in ops
    assert "Conv" in ops
    assert "Gemm" in ops
check("AllayRT.supported_ops()", t_api_supported_ops)

def t_api_custom_op():
    from AllayRT.core import register_op
    def double_op(inputs, attrs, ctx):
        return [inputs[0] * 2]
    AllayRT.register_op("Double", double_op)
    assert "Double" in OP_TABLE
check("AllayRT.register_op()", t_api_custom_op)

def t_api_bad_device():
    m  = art_rt.read_model(tmp2)
    try:
        cm = art_rt.compile_model(m, "QUANTUM")
        raise AssertionError("Should have raised ValueError")
    except ValueError:
        pass
check("Unsupported device raises ValueError", t_api_bad_device)

os.unlink(tmp2)


# ──────────────────────────────────────────────
#  Summary
# ──────────────────────────────────────────────

total = len(results)
passed = sum(1 for _, ok, _ in results if ok)
failed = total - passed

print(f"\n{'═'*40}")
print(f"Results: {passed}/{total} passed", end="")
if failed:
    print(f"  ({failed} failed)")
    for name, ok, err in results:
        if not ok:
            print(f"  {FAIL} {name}: {err}")
else:
    print("  — All tests passed! 🎉")
print(f"{'═'*40}\n")

sys.exit(0 if failed == 0 else 1)
