test_special_elementwise.py

"""Tests for special predicate elementwise ops (isnan, isinf, isfinite).

Covers L1 smoke correctness (fp16, 1M) and L4 edge cases (fp32, 4K).
"""

import pytest
import torch

from tests.test_base import FixtureBase, TestBase, exact_compare
from tileops.ops.elementwise import IsfiniteFwdOp, IsinfFwdOp, IsnanFwdOp
from tileops.utils import get_backend_name

DEVICE = get_backend_name()


class SpecialFixture(FixtureBase):
    """Parametrize over shapes / dtypes for special predicate ops."""

    PARAMS = [
        ("n_total, dtype", [
            pytest.param(1_048_576, torch.float16, marks=pytest.mark.smoke),
            pytest.param(1_048_576, torch.bfloat16, marks=pytest.mark.smoke),
            pytest.param(1_048_576, torch.float32, marks=pytest.mark.smoke),
        ]),
    ]


class SpecialEdgeFixture(FixtureBase):
    """L4 edge-case fixture: fp32, 4K elements."""

    PARAMS = [
        ("n_total, dtype", [
            pytest.param(4096, torch.float32, marks=pytest.mark.smoke),
        ]),
    ]


class SpecialTest(TestBase):
    """Generic test harness for special predicate ops."""

    def __init__(self, n_total: int, dtype: torch.dtype, ref_fn, gen_fn=None):
        self.n_total = n_total
        self.dtype = dtype
        self._ref_fn = ref_fn
        self._gen_fn = gen_fn

    def gen_inputs(self) -> tuple[torch.Tensor]:
        if self._gen_fn is not None:
            return (self._gen_fn(self.n_total, self.dtype),)
        x = torch.randn(self.n_total, device=DEVICE, dtype=self.dtype)
        quarter = self.n_total // 4
        x[:quarter] = float("nan")
        x[quarter:2 * quarter] = float("inf")
        x[2 * quarter:3 * quarter] = float("-inf")
        return (x,)

    def ref_program(self, x: torch.Tensor) -> torch.Tensor:
        return self._ref_fn(x)


def _make_special_test(n_total, dtype, op_cls, ref_fn, gen_fn=None) -> None:
    test = SpecialTest(n_total, dtype, ref_fn=ref_fn, gen_fn=gen_fn)
    op = op_cls(N_total=n_total, dtype=dtype)
    test.check(op, *test.gen_inputs(), compare=exact_compare)


@SpecialFixture
def test_isnan(n_total: int, dtype: torch.dtype) -> None:
    _make_special_test(n_total, dtype, IsnanFwdOp, torch.isnan)


@SpecialFixture
def test_isinf(n_total: int, dtype: torch.dtype) -> None:
    _make_special_test(n_total, dtype, IsinfFwdOp, torch.isinf)


@SpecialFixture
def test_isfinite(n_total: int, dtype: torch.dtype) -> None:
    _make_special_test(n_total, dtype, IsfiniteFwdOp, torch.isfinite)


# ---------------------------------------------------------------------------
# L4 edge-case tests (fp32, 4K)
# ---------------------------------------------------------------------------


@SpecialEdgeFixture
def test_isnan_edge(n_total: int, dtype: torch.dtype) -> None:
    """Edge: all NaN input."""
    def _all_nan(n, dtype):
        return torch.full((n,), float("nan"), device=DEVICE, dtype=dtype)

    _make_special_test(n_total, dtype, IsnanFwdOp, torch.isnan, gen_fn=_all_nan)


@SpecialEdgeFixture
def test_isinf_edge(n_total: int, dtype: torch.dtype) -> None:
    """Edge: mix of +inf and -inf."""
    def _all_inf(n, dtype):
        x = torch.full((n,), float("inf"), device=DEVICE, dtype=dtype)
        x[:n // 2] = float("-inf")
        return x

    _make_special_test(n_total, dtype, IsinfFwdOp, torch.isinf, gen_fn=_all_inf)


@SpecialEdgeFixture
def test_isfinite_edge(n_total: int, dtype: torch.dtype) -> None:
    """Edge: all finite input."""
    def _all_finite(n, dtype):
        return torch.randn(n, device=DEVICE, dtype=dtype)

    _make_special_test(n_total, dtype, IsfiniteFwdOp, torch.isfinite, gen_fn=_all_finite)


@pytest.mark.smoke
def test_special_predicates_reject_non_float_dtype() -> None:
    from tileops.kernels.elementwise import IsnanFwdKernel

    with pytest.raises(ValueError, match="only supports dtypes"):
        IsnanFwdKernel(N_total=16, dtype=torch.int32)


# ===========================================================================
# Independent special ops: where, clamp, masked_fill, nan_to_num,
# alibi, sinusoidal
# ===========================================================================


class IndependentFixture(FixtureBase):
    """Parametrize over shapes / dtypes for independent custom-signature ops."""

    PARAMS = [
        ("n_total, dtype", [
            pytest.param(1_048_576, torch.float16, marks=pytest.mark.smoke),
            pytest.param(1_048_576, torch.bfloat16, marks=pytest.mark.smoke),
            pytest.param(1_048_576, torch.float32, marks=pytest.mark.smoke),
        ]),
    ]


class IndependentEdgeFixture(FixtureBase):
    """L4 edge-case fixture: fp32, 4K elements."""

    PARAMS = [
        ("n_total, dtype", [
            pytest.param(4096, torch.float32, marks=pytest.mark.smoke),
        ]),
    ]


# --- L1: where ---

@IndependentFixture
def test_where(n_total: int, dtype: torch.dtype) -> None:
    from tileops.ops.elementwise import WhereFwdOp

    cond = torch.randint(0, 2, (n_total,), device=DEVICE).bool()
    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    y = torch.randn(n_total, device=DEVICE, dtype=dtype)
    ref = torch.where(cond, x, y)
    op = WhereFwdOp(condition=(n_total,), input=(n_total,), other=(n_total,), dtype=dtype)
    out = op(cond, x, y)
    torch.testing.assert_close(out, ref, atol=0, rtol=0)
    print("All checks passed for WhereFwdOp.")


# --- L1: clamp ---

@IndependentFixture
def test_clamp(n_total: int, dtype: torch.dtype) -> None:
    from tileops.ops.elementwise import ClampScalarFwdOp

    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    ref = torch.clamp(x, -0.5, 0.5)
    op = ClampScalarFwdOp(input=(n_total,), min=-0.5, max=0.5, dtype=dtype)
    out = op(x)
    if dtype == torch.float16:
        tol = {"atol": 1e-3, "rtol": 1e-3}
    elif dtype == torch.bfloat16:
        tol = {"atol": 1.6e-2, "rtol": 1.6e-2}
    else:
        tol = {"atol": 1e-5, "rtol": 1e-5}
    torch.testing.assert_close(out, ref, **tol)
    print("All checks passed for ClampFwdOp.")


# --- L1: masked_fill ---

@IndependentFixture
def test_masked_fill(n_total: int, dtype: torch.dtype) -> None:
    from tileops.ops.elementwise import MaskedFillScalarFwdOp

    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    mask = torch.randint(0, 2, (n_total,), device=DEVICE).bool()
    # Use -100.0 to avoid fp16 overflow (fp16 max ~65504)
    fill_value = -100.0
    ref = x.masked_fill(mask, fill_value)
    op = MaskedFillScalarFwdOp(input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype)
    out = op(x, mask)
    if dtype == torch.float16:
        tol = {"atol": 1e-3, "rtol": 1e-3}
    elif dtype == torch.bfloat16:
        tol = {"atol": 1.6e-2, "rtol": 1.6e-2}
    else:
        tol = {"atol": 1e-5, "rtol": 1e-5}
    torch.testing.assert_close(out, ref, **tol)
    print("All checks passed for MaskedFillFwdOp.")


# --- L1: nan_to_num ---

@IndependentFixture
def test_nan_to_num(n_total: int, dtype: torch.dtype) -> None:
    from tileops.ops.elementwise import NanToNumFwdOp

    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    quarter = n_total // 4
    x[:quarter] = float("nan")
    x[quarter:2 * quarter] = float("inf")
    x[2 * quarter:3 * quarter] = float("-inf")
    ref = torch.nan_to_num(x, nan=0.0, posinf=1e4, neginf=-1e4)
    op = NanToNumFwdOp(N_total=n_total, dtype=dtype, nan=0.0, posinf=1e4, neginf=-1e4)
    out = op(x)
    if dtype == torch.float16:
        tol = {"atol": 1e-3, "rtol": 1e-3}
    elif dtype == torch.bfloat16:
        tol = {"atol": 1.6e-2, "rtol": 1.6e-2}
    else:
        tol = {"atol": 1e-5, "rtol": 1e-5}
    torch.testing.assert_close(out, ref, **tol, equal_nan=True)
    print("All checks passed for NanToNumFwdOp.")


# --- L1: alibi ---

class AlibiFixture(FixtureBase):
    PARAMS = [
        ("seq_len, num_heads, dtype", [
            pytest.param(128, 8, torch.float16, marks=pytest.mark.smoke),
            pytest.param(128, 8, torch.float32, marks=pytest.mark.smoke),
        ]),
    ]


@AlibiFixture
def test_alibi(seq_len: int, num_heads: int, dtype: torch.dtype) -> None:
    from tileops.ops.elementwise import AlibiFwdOp

    op = AlibiFwdOp(seq_len=seq_len, num_heads=num_heads, dtype=dtype)
    out = op()

    # Reference: slope_h = 2^(-8*(h+1)/H), bias = -slope * |i - j|
    positions = torch.arange(seq_len, device=DEVICE, dtype=torch.float32)
    dist = (positions.unsqueeze(1) - positions.unsqueeze(0)).abs()
    slopes = torch.pow(
        2.0,
        -8.0 * torch.arange(1, num_heads + 1, device=DEVICE, dtype=torch.float32) / num_heads,
    )
    ref = (-slopes[:, None, None] * dist[None, :, :]).to(dtype)

    tol = {"atol": 1e-2, "rtol": 1e-2} if dtype == torch.float16 else {"atol": 1e-5, "rtol": 1e-5}
    torch.testing.assert_close(out, ref, **tol)
    print("All checks passed for AlibiFwdOp.")


# --- L1: sinusoidal ---

class SinusoidalFixture(FixtureBase):
    PARAMS = [
        ("seq_len, d_model, dtype", [
            pytest.param(512, 256, torch.float16, marks=pytest.mark.smoke),
            pytest.param(512, 256, torch.float32, marks=pytest.mark.smoke),
        ]),
    ]


@SinusoidalFixture
def test_sinusoidal(seq_len: int, d_model: int, dtype: torch.dtype) -> None:

    from tileops.ops.elementwise import SinusoidalFwdOp

    op = SinusoidalFwdOp(seq_len=seq_len, d_model=d_model, dtype=dtype)
    out = op()

    # Reference: compute fp32 cases on CPU float64 to avoid torch_musa
    # device-math approximation differences becoming the source of truth.
    ref_device = "cpu" if dtype == torch.float32 else DEVICE
    ref_compute_dtype = torch.float64 if dtype == torch.float32 else torch.float32
    pos = torch.arange(seq_len, device=ref_device, dtype=ref_compute_dtype).unsqueeze(1)
    dim_pairs = torch.arange(0, d_model, 2, device=ref_device, dtype=ref_compute_dtype)
    base = torch.tensor(10000.0, device=ref_device, dtype=ref_compute_dtype)
    angles = pos / torch.pow(base, dim_pairs / d_model)
    ref = torch.zeros(seq_len, d_model, device=ref_device, dtype=ref_compute_dtype)
    ref[:, 0::2] = torch.sin(angles)
    ref[:, 1::2] = torch.cos(angles)
    ref = ref.to(device=DEVICE, dtype=dtype)

    if dtype == torch.float16:
        tol = {"atol": 1e-3, "rtol": 1e-3}
    else:
        tol = {"atol": 1e-5, "rtol": 1e-5}
    torch.testing.assert_close(out, ref, **tol)
    print("All checks passed for SinusoidalFwdOp.")


# ===========================================================================
# L2 — Dtype x Size (4 cases for clamp)
# ===========================================================================


class ClampDtypeSizeFixture(FixtureBase):
    PARAMS = [
        ("n_total, dtype", [
            pytest.param(1_048_576, torch.float32, marks=pytest.mark.smoke),
            pytest.param(1_048_576, torch.bfloat16, marks=pytest.mark.smoke),
            pytest.param(4096, torch.float16, marks=pytest.mark.smoke),
            pytest.param(16_777_216, torch.float16, marks=pytest.mark.full),
        ]),
    ]


@ClampDtypeSizeFixture
def test_clamp_dtype_size(n_total: int, dtype: torch.dtype) -> None:
    from tileops.ops.elementwise import ClampScalarFwdOp

    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    ref = torch.clamp(x, -0.5, 0.5)
    op = ClampScalarFwdOp(input=(n_total,), min=-0.5, max=0.5, dtype=dtype)
    out = op(x)
    if dtype == torch.float16:
        tol = {"atol": 1e-3, "rtol": 1e-3}
    elif dtype == torch.bfloat16:
        tol = {"atol": 1.6e-2, "rtol": 1.6e-2}
    else:
        tol = {"atol": 1e-5, "rtol": 1e-5}
    torch.testing.assert_close(out, ref, **tol)
    print("All checks passed for ClampFwdOp dtype/size variant.")


# ===========================================================================
# L4 — Edge Cases (8 cases, fp32, 4K)
# ===========================================================================


@IndependentEdgeFixture
def test_clamp_min_gt_max(n_total: int, dtype: torch.dtype) -> None:
    """Edge: min > max -- PyTorch clamp semantics: min wins (output = min_val)."""
    from tileops.ops.elementwise import ClampScalarFwdOp

    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    # When min > max, PyTorch clamp returns min_val for all elements
    ref = torch.clamp(x, min=0.5, max=-0.5)
    op = ClampScalarFwdOp(input=(n_total,), min=0.5, max=-0.5, dtype=dtype)
    out = op(x)
    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
    print("All checks passed for ClampFwdOp min>max edge case.")


@IndependentEdgeFixture
def test_clamp_upper_only(n_total: int, dtype: torch.dtype) -> None:
    """Edge: min=None, max=0.5 (upper bound only)."""
    from tileops.ops.elementwise import ClampScalarFwdOp

    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    ref = torch.clamp(x, min=None, max=0.5)
    op = ClampScalarFwdOp(input=(n_total,), min=None, max=0.5, dtype=dtype)
    out = op(x)
    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
    print("All checks passed for ClampFwdOp upper-only edge case.")


@IndependentEdgeFixture
def test_clamp_lower_only(n_total: int, dtype: torch.dtype) -> None:
    """Edge: min=-0.5, max=None (lower bound only)."""
    from tileops.ops.elementwise import ClampScalarFwdOp

    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    ref = torch.clamp(x, min=-0.5, max=None)
    op = ClampScalarFwdOp(input=(n_total,), min=-0.5, max=None, dtype=dtype)
    out = op(x)
    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
    print("All checks passed for ClampFwdOp lower-only edge case.")


@IndependentEdgeFixture
def test_masked_fill_all_true(n_total: int, dtype: torch.dtype) -> None:
    """Edge: all True mask -> all values replaced."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp

    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    mask = torch.ones(n_total, device=DEVICE, dtype=torch.bool)
    fill_value = -1e9
    ref = x.masked_fill(mask, fill_value)
    op = MaskedFillScalarFwdOp(input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype)
    out = op(x, mask)
    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
    print("All checks passed for MaskedFillFwdOp all-true edge case.")


@IndependentEdgeFixture
def test_masked_fill_all_false(n_total: int, dtype: torch.dtype) -> None:
    """Edge: all False mask -> input unchanged."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp

    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    mask = torch.zeros(n_total, device=DEVICE, dtype=torch.bool)
    fill_value = -1e9
    ref = x.masked_fill(mask, fill_value)
    op = MaskedFillScalarFwdOp(input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype)
    out = op(x, mask)
    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
    print("All checks passed for MaskedFillFwdOp all-false edge case.")


@IndependentEdgeFixture
def test_where_all_true(n_total: int, dtype: torch.dtype) -> None:
    """Edge: all True cond -> output = x."""
    from tileops.ops.elementwise import WhereFwdOp

    cond = torch.ones(n_total, device=DEVICE, dtype=torch.bool)
    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    y = torch.randn(n_total, device=DEVICE, dtype=dtype)
    ref = torch.where(cond, x, y)
    op = WhereFwdOp(condition=(n_total,), input=(n_total,), other=(n_total,), dtype=dtype)
    out = op(cond, x, y)
    torch.testing.assert_close(out, ref, atol=0, rtol=0)
    print("All checks passed for WhereFwdOp all-true edge case.")


@IndependentEdgeFixture
def test_where_all_false(n_total: int, dtype: torch.dtype) -> None:
    """Edge: all False cond -> output = y."""
    from tileops.ops.elementwise import WhereFwdOp

    cond = torch.zeros(n_total, device=DEVICE, dtype=torch.bool)
    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    y = torch.randn(n_total, device=DEVICE, dtype=dtype)
    ref = torch.where(cond, x, y)
    op = WhereFwdOp(condition=(n_total,), input=(n_total,), other=(n_total,), dtype=dtype)
    out = op(cond, x, y)
    torch.testing.assert_close(out, ref, atol=0, rtol=0)
    print("All checks passed for WhereFwdOp all-false edge case.")


@IndependentEdgeFixture
def test_nan_to_num_edge(n_total: int, dtype: torch.dtype) -> None:
    """Edge: explicit [NaN, Inf, -Inf, 1.0] pattern."""
    from tileops.ops.elementwise import NanToNumFwdOp

    x = torch.zeros(n_total, device=DEVICE, dtype=dtype)
    # Fill pattern: NaN, Inf, -Inf, 1.0, repeating
    for k in range(0, n_total, 4):
        x[k] = float("nan")
        if k + 1 < n_total:
            x[k + 1] = float("inf")
        if k + 2 < n_total:
            x[k + 2] = float("-inf")
        if k + 3 < n_total:
            x[k + 3] = 1.0

    ref = torch.nan_to_num(x, nan=0.0, posinf=1e4, neginf=-1e4)
    op = NanToNumFwdOp(N_total=n_total, dtype=dtype, nan=0.0, posinf=1e4, neginf=-1e4)
    out = op(x)
    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5, equal_nan=True)
    print("All checks passed for NanToNumFwdOp edge case.")


@pytest.mark.smoke
def test_independent_special_rejects_non_float_dtype() -> None:
    from tileops.kernels.elementwise import ClampFwdKernel
    with pytest.raises(ValueError, match="only supports dtypes"):
        ClampFwdKernel(N_total=16, dtype=torch.int32)


# ===========================================================================
# Negative tests: forward() dtype / numel validation
# ===========================================================================


@pytest.mark.smoke
@pytest.mark.parametrize("op_cls, kwargs", [
    pytest.param("EluFwdOp", {"alpha": 1.0}, id="elu"),
    pytest.param("HardtanhFwdOp", {"min_val": -1.0, "max_val": 1.0}, id="hardtanh"),
    pytest.param("SoftplusFwdOp", {"beta": 1.0, "threshold": 20.0}, id="softplus"),
    pytest.param("ClampScalarFwdOp", {"min": -0.5, "max": 0.5}, id="clamp"),
])
def test_forward_rejects_wrong_dtype(op_cls: str, kwargs: dict) -> None:
    """forward() must raise ValueError when input dtype mismatches."""
    import tileops.ops.elementwise as mod
    cls = getattr(mod, op_cls)
    if cls.__name__ == "ClampScalarFwdOp":
        op = cls(input=(1024,), dtype=torch.float16, **kwargs)
    else:
        op = cls(N_total=1024, dtype=torch.float16, **kwargs)
    x = torch.randn(1024, device=DEVICE, dtype=torch.float32)
    with pytest.raises(ValueError, match="dtype"):
        op(x)


@pytest.mark.smoke
@pytest.mark.parametrize("op_cls, kwargs", [
    pytest.param("EluFwdOp", {"alpha": 1.0}, id="elu"),
    pytest.param("HardtanhFwdOp", {"min_val": -1.0, "max_val": 1.0}, id="hardtanh"),
    pytest.param("SoftplusFwdOp", {"beta": 1.0, "threshold": 20.0}, id="softplus"),
])
def test_forward_rejects_wrong_numel(op_cls: str, kwargs: dict) -> None:
    """forward() must raise ValueError when input numel mismatches.

    ClampScalarFwdOp validates the full input.shape (not just numel), so
    its mismatch case is covered by
    test_clamp_scalar_rejects_same_numel_wrong_shape in
    tests/ops/test_special_elementwise_conformance.py.
    """
    import tileops.ops.elementwise as mod
    cls = getattr(mod, op_cls)
    op = cls(N_total=1024, dtype=torch.float16, **kwargs)
    x = torch.randn(512, device=DEVICE, dtype=torch.float16)
    with pytest.raises(ValueError, match="elements"):
        op(x)


@pytest.mark.smoke
@pytest.mark.parametrize("op_cls, kwargs", [
    pytest.param("MaskedFillScalarFwdOp", {"value": -100.0}, id="masked_fill"),
])
def test_masked_fill_forward_rejects_wrong_dtype(op_cls: str, kwargs: dict) -> None:
    """MaskedFillFwdOp forward() must raise ValueError when input dtype mismatches."""
    import tileops.ops.elementwise as mod
    cls = getattr(mod, op_cls)
    op = cls(input=(1024,), mask=(1024,), dtype=torch.float16, **kwargs)
    x = torch.randn(1024, device=DEVICE, dtype=torch.float32)
    mask = torch.ones(1024, device=DEVICE, dtype=torch.bool)
    with pytest.raises(ValueError, match="dtype"):
        op(x, mask)


@pytest.mark.smoke
@pytest.mark.parametrize("op_cls, kwargs", [
    pytest.param("MaskedFillScalarFwdOp", {"value": -100.0}, id="masked_fill"),
])
def test_masked_fill_forward_rejects_wrong_numel(op_cls: str, kwargs: dict) -> None:
    """MaskedFillFwdOp forward() must raise ValueError when input shape mismatches."""
    import tileops.ops.elementwise as mod
    cls = getattr(mod, op_cls)
    op = cls(input=(1024,), mask=(1024,), dtype=torch.float16, **kwargs)
    x = torch.randn(512, device=DEVICE, dtype=torch.float16)
    mask = torch.ones(512, device=DEVICE, dtype=torch.bool)
    with pytest.raises(ValueError, match="input.shape"):
        op(x, mask)


# ===========================================================================
# Negative tests: __init__() scalar parameter validation
# ===========================================================================


@pytest.mark.smoke
def test_elu_rejects_unrepresentable_alpha() -> None:
    """EluFwdOp must reject alpha that overflows the kernel dtype."""
    from tileops.ops.elementwise import EluFwdOp
    with pytest.raises((ValueError, TypeError)):
        EluFwdOp(N_total=1024, dtype=torch.float16, alpha=1e6)


@pytest.mark.smoke
def test_hardtanh_rejects_unrepresentable_min_val() -> None:
    """HardtanhFwdOp must reject min_val that overflows the kernel dtype."""
    from tileops.ops.elementwise import HardtanhFwdOp
    with pytest.raises((ValueError, TypeError)):
        HardtanhFwdOp(N_total=1024, dtype=torch.float16, min_val=1e6)


@pytest.mark.smoke
def test_hardtanh_rejects_unrepresentable_max_val() -> None:
    """HardtanhFwdOp must reject max_val that overflows the kernel dtype."""
    from tileops.ops.elementwise import HardtanhFwdOp
    with pytest.raises((ValueError, TypeError)):
        HardtanhFwdOp(N_total=1024, dtype=torch.float16, max_val=1e6)


@pytest.mark.smoke
def test_softplus_rejects_unrepresentable_beta() -> None:
    """SoftplusFwdOp must reject beta that overflows the kernel dtype."""
    from tileops.ops.elementwise import SoftplusFwdOp
    with pytest.raises((ValueError, TypeError)):
        SoftplusFwdOp(N_total=1024, dtype=torch.float16, beta=1e6)


@pytest.mark.smoke
def test_softplus_rejects_unrepresentable_threshold() -> None:
    """SoftplusFwdOp must reject threshold that overflows the kernel dtype."""
    from tileops.ops.elementwise import SoftplusFwdOp
    with pytest.raises((ValueError, TypeError)):
        SoftplusFwdOp(N_total=1024, dtype=torch.float16, threshold=1e6)


@pytest.mark.smoke
def test_clamp_rejects_unrepresentable_min_val() -> None:
    """ClampFwdOp must reject min_val that overflows the kernel dtype."""
    from tileops.ops.elementwise import ClampScalarFwdOp
    with pytest.raises((ValueError, TypeError)):
        ClampScalarFwdOp(input=(1024,), min=1e6, dtype=torch.float16)


@pytest.mark.smoke
def test_clamp_rejects_unrepresentable_max_val() -> None:
    """ClampFwdOp must reject max_val that overflows the kernel dtype."""
    from tileops.ops.elementwise import ClampScalarFwdOp
    with pytest.raises((ValueError, TypeError)):
        ClampScalarFwdOp(input=(1024,), max=1e6, dtype=torch.float16)


@pytest.mark.smoke
def test_masked_fill_forward_rejects_cpu_mask() -> None:
    """MaskedFillFwdOp forward() must raise ValueError when mask is not on backend."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp
    op = MaskedFillScalarFwdOp(input=(1024,), mask=(1024,), value=-100.0, dtype=torch.float16)
    x = torch.randn(1024, device=DEVICE, dtype=torch.float16)
    mask = torch.ones(1024, dtype=torch.bool)  # CPU mask
    with pytest.raises(ValueError, match=f"Mask must be a {DEVICE.upper()} tensor"):
        op(x, mask)


@pytest.mark.smoke
def test_masked_fill_forward_rejects_non_bool_mask() -> None:
    """MaskedFillFwdOp forward() must raise ValueError when mask dtype is not bool."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp
    op = MaskedFillScalarFwdOp(input=(1024,), mask=(1024,), value=-100.0, dtype=torch.float16)
    x = torch.randn(1024, device=DEVICE, dtype=torch.float16)
    mask = torch.ones(1024, device=DEVICE, dtype=torch.float32)  # wrong dtype
    with pytest.raises(ValueError, match="mask.dtype"):
        op(x, mask)


@pytest.mark.smoke
def test_masked_fill_forward_rejects_wrong_mask_numel() -> None:
    """MaskedFillFwdOp forward() must raise ValueError when mask numel mismatches."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp
    op = MaskedFillScalarFwdOp(input=(1024,), mask=(1024,), value=-100.0, dtype=torch.float16)
    x = torch.randn(1024, device=DEVICE, dtype=torch.float16)
    mask = torch.ones(512, device=DEVICE, dtype=torch.bool)  # wrong shape
    with pytest.raises(ValueError, match="mask.shape"):
        op(x, mask)


# ---------------------------------------------------------------------------
# MaskedFillScalar: int / uint / bool dtype coverage
# ---------------------------------------------------------------------------


_MASKED_FILL_INT_DTYPES = [
    torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64,
]


def _masked_fill_int_inputs(n_total: int, dtype: torch.dtype):
    iinfo = torch.iinfo(dtype)
    lo = max(iinfo.min, -1000)
    hi = min(iinfo.max, 1000) + 1
    x = torch.randint(lo, hi, (n_total,), device=DEVICE, dtype=dtype)
    mask = torch.randint(0, 2, (n_total,), device=DEVICE).bool()
    return x, mask


@pytest.mark.smoke
@pytest.mark.parametrize("dtype", _MASKED_FILL_INT_DTYPES)
def test_masked_fill_int_dtypes(dtype: torch.dtype) -> None:
    """L1: each manifest int dtype matches PyTorch on a representative fill."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp

    n_total = 4096
    fill_value = 7  # arbitrary in-range value; the contract is parity with PyTorch.
    x, mask = _masked_fill_int_inputs(n_total, dtype)
    ref = x.masked_fill(mask, fill_value)
    op = MaskedFillScalarFwdOp(
        input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype,
    )
    out = op(x, mask)
    torch.testing.assert_close(out, ref, atol=0, rtol=0)


@pytest.mark.smoke
def test_masked_fill_uint8_wraps_negative_int() -> None:
    """uint8 wraps a negative Python int via two's complement (PyTorch: -1 -> 255)."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp

    n_total = 4096
    x, mask = _masked_fill_int_inputs(n_total, torch.uint8)
    ref = x.masked_fill(mask, -1)
    op = MaskedFillScalarFwdOp(
        input=(n_total,), mask=(n_total,), value=-1, dtype=torch.uint8,
    )
    torch.testing.assert_close(op(x, mask), ref, atol=0, rtol=0)


@pytest.mark.smoke
def test_masked_fill_int_truncates_fractional_float() -> None:
    """Integer dtypes truncate a float fill toward zero (PyTorch: 1.5 -> 1)."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp

    n_total = 4096
    x, mask = _masked_fill_int_inputs(n_total, torch.int32)
    ref = x.masked_fill(mask, 1.5)
    op = MaskedFillScalarFwdOp(
        input=(n_total,), mask=(n_total,), value=1.5, dtype=torch.int32,
    )
    torch.testing.assert_close(op(x, mask), ref, atol=0, rtol=0)


@pytest.mark.smoke
@pytest.mark.parametrize("fill_value", [True, False])
def test_masked_fill_bool(fill_value) -> None:
    """L1: bool masked_fill coerces non-zero -> True via uint8 storage view."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp

    n_total = 4096
    x = torch.randint(0, 2, (n_total,), device=DEVICE).bool()
    mask = torch.randint(0, 2, (n_total,), device=DEVICE).bool()
    ref = x.masked_fill(mask, fill_value)
    op = MaskedFillScalarFwdOp(
        input=(n_total,), mask=(n_total,), value=fill_value, dtype=torch.bool,
    )
    out = op(x, mask)
    torch.testing.assert_close(out, ref, atol=0, rtol=0)


@pytest.mark.smoke
@pytest.mark.parametrize("dtype, fill_value", [
    pytest.param(torch.float16, float("inf"), id="fp16-inf"),
    pytest.param(torch.bfloat16, float("-inf"), id="bf16-neg-inf"),
    pytest.param(torch.float32, float("nan"), id="fp32-nan"),
])
def test_masked_fill_float_nonfinite(dtype: torch.dtype, fill_value: float) -> None:
    """L4: +/-Inf and NaN fill values pass through unchanged (no clamp)."""
    from tileops.ops.elementwise import MaskedFillScalarFwdOp

    n_total = 4096
    x = torch.randn(n_total, device=DEVICE, dtype=dtype)
    mask = torch.randint(0, 2, (n_total,), device=DEVICE).bool()
    ref = x.masked_fill(mask, fill_value)
    op = MaskedFillScalarFwdOp(
        input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype,
    )
    out = op(x, mask)
    torch.testing.assert_close(out, ref, atol=0, rtol=0, equal_nan=True)


_MASKED_FILL_REJECT_CASES = [
    pytest.param(torch.int32, float("inf"), id="int-inf"),
    pytest.param(torch.int32, float("nan"), id="int-nan"),
]


@pytest.mark.smoke
@pytest.mark.parametrize("dtype, fill_value", _MASKED_FILL_REJECT_CASES)
def test_masked_fill_rejects_when_pytorch_rejects(
    dtype: torch.dtype, fill_value,
) -> None:
    """Op must reject every scalar that PyTorch's own masked_fill rejects.

    The contract is parity, not the error message; assert both call sites
    raise, leaving wording to the implementation.
    """
    from tileops.ops.elementwise import MaskedFillScalarFwdOp

    pytorch_mask = torch.tensor([True], device=DEVICE)
    pytorch_tensor = torch.zeros(1, device=DEVICE, dtype=dtype)
    with pytest.raises(Exception):  # noqa: B017
        pytorch_tensor.masked_fill(pytorch_mask, fill_value)

    with pytest.raises(Exception):  # noqa: B017
        MaskedFillScalarFwdOp(
            input=(1024,), mask=(1024,), value=fill_value, dtype=dtype,
        )


@pytest.mark.smoke
def test_elu_rejects_infinite_alpha() -> None:
    """EluFwdOp must reject infinite alpha."""
    from tileops.ops.elementwise import EluFwdOp
    with pytest.raises(ValueError, match="finite"):
        EluFwdOp(N_total=1024, dtype=torch.float32, alpha=float("inf"))


@pytest.mark.smoke
def test_softplus_rejects_non_numeric_beta() -> None:
    """SoftplusFwdOp must reject non-numeric beta."""
    from tileops.ops.elementwise import SoftplusFwdOp
    with pytest.raises(TypeError, match="int/float"):
        SoftplusFwdOp(N_total=1024, dtype=torch.float32, beta="bad")


if __name__ == "__main__":
    pytest.main([__file__, "-vvs"])