# Copyright (c) 2025, Alibaba Group;
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#    http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gc
import unittest
from typing import Optional

import torch
from hypothesis import Verbosity, given
from hypothesis import strategies as st

from tzrec.ops import Kernel
from tzrec.utils.test_util import (
    generate_sparse_seq_len,
    get_test_dtypes,
    gpu_unavailable,
)
from tzrec.utils.test_util import hypothesis_settings as settings


class JaggedTensorsTest(unittest.TestCase):
    def teardown_example(self, example):
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.memory_summary()  # prevent oom

    @unittest.skipIf(*gpu_unavailable)
    # pyre-ignore
    @given(
        batch_size=st.integers(2, 8),
        max_len_a=st.integers(20, 100),
        max_len_b=st.integers(20, 100),
        D=st.integers(10, 30),
        is_dense_a=st.sampled_from([True, False]),
        is_dense_b=st.sampled_from([True, False]),
        dtype=st.sampled_from(get_test_dtypes([torch.bfloat16, torch.float32])),
    )
    @settings(
        verbosity=Verbosity.verbose,
        max_examples=20,
        deadline=None,
    )
    # pyre-ignore[2]
    def test_split_2D_jagged_triton(self, *args, **kwargs) -> None:
        self._test_split_2D_jagged(
            *args,
            **kwargs,
            test_backward=True,
            ref_kernel=Kernel.PYTORCH,
            real_kernel=Kernel.TRITON,
        )

    def _test_split_2D_jagged(
        self,
        batch_size: int,
        max_len_a: int,
        max_len_b: int,
        D: int,
        is_dense_a: bool,
        is_dense_b: bool,
        test_backward: bool,
        ref_kernel: Kernel,
        real_kernel: Kernel,
        dtype: torch.dtype = torch.float32,
        skip_comparisons: bool = False,
    ) -> None:
        from tzrec.ops.jagged_tensors import split_2D_jagged

        max_seq_len = max_len_a + max_len_b
        if not is_dense_a:
            lengths_a = torch.randint(
                1, max_len_a + 1, size=(batch_size,), device=torch.device("cuda")
            )
            offsets_a = torch.zeros(
                (batch_size + 1,), dtype=torch.int64, device=torch.device("cuda")
            )
            offsets_a[1:] = torch.cumsum(lengths_a, dim=0)
            total_len_a = int(offsets_a[-1].item())
        else:
            offsets_a = None
            total_len_a = batch_size * max_len_a
            is_dense_b = False
        if not is_dense_b:
            lengths_b = torch.randint(
                1, max_len_b + 1, size=(batch_size,), device=torch.device("cuda")
            )
            offsets_b = torch.zeros(
                (batch_size + 1,), dtype=torch.int64, device=torch.device("cuda")
            )
            offsets_b[1:] = torch.cumsum(lengths_b, dim=0)
            total_len_b = int(offsets_b[-1].item())
        else:
            offsets_b = None
            total_len_b = batch_size * max_len_b
        values = (
            torch.empty(
                (total_len_a + total_len_b, D),
                dtype=dtype,
                device=torch.device("cuda"),
            )
            .uniform_(-1.0, 1.0)
            .requires_grad_()
        )

        ref_values_a, ref_values_b = split_2D_jagged(
            max_seq_len=max_len_a + max_len_b,
            values=values,
            max_len_left=max_len_a if is_dense_a else None,
            max_len_right=max_len_b if is_dense_b else None,
            offsets_left=offsets_a,
            offsets_right=offsets_b,
            kernel=ref_kernel,
        )
        d_values_a = torch.randn_like(ref_values_a)
        d_values_b = torch.randn_like(ref_values_b)
        ref_values_a.backward(d_values_a, retain_graph=True)
        ref_values_b.backward(d_values_b)
        if skip_comparisons:
            return

        assert values.grad is not None
        ref_d_values, values.grad = values.grad.clone(), None

        values = values.detach().clone().requires_grad_()
        real_values_a, real_values_b = split_2D_jagged(
            max_seq_len=max_seq_len,
            values=values,
            max_len_left=max_len_a if is_dense_a else None,
            max_len_right=max_len_b if is_dense_b else None,
            offsets_left=offsets_a,
            offsets_right=offsets_b,
            kernel=real_kernel,
        )
        torch.testing.assert_close(ref_values_a, real_values_a)
        torch.testing.assert_close(ref_values_b, real_values_b)

        if test_backward:
            d_values_a = d_values_a.detach().clone()
            d_values_b = d_values_b.detach().clone()
            real_values_a.backward(d_values_a, retain_graph=True)
            real_values_b.backward(d_values_b)
            real_d_values = values.grad.clone()
            torch.testing.assert_close(ref_d_values, real_d_values)

    @unittest.skipIf(*gpu_unavailable)
    # pyre-ignore
    @given(
        batch_size=st.integers(2, 8),
        max_len_a=st.integers(20, 100),
        max_len_b=st.integers(20, 100),
        D=st.integers(10, 30),
        is_dense_a=st.sampled_from([True, False]),
        is_dense_b=st.sampled_from([True, False]),
        dtype=st.sampled_from(get_test_dtypes([torch.bfloat16, torch.float32])),
    )
    @settings(
        verbosity=Verbosity.verbose,
        max_examples=20,
        deadline=None,
    )
    # pyre-ignore[2]
    def test_concat_2D_jagged_triton(self, *args, **kwargs) -> None:
        self._test_concat_2D_jagged(
            *args,
            **kwargs,
            test_backward=True,
            ref_kernel=Kernel.PYTORCH,
            real_kernel=Kernel.TRITON,
        )

    @unittest.skipIf(*gpu_unavailable)
    # pyre-ignore
    @given(
        batch_size=st.sampled_from([130]),
        max_len_a=st.sampled_from([32768]),
        max_len_b=st.sampled_from([10]),
        D=st.sampled_from([256]),
        is_dense_a=st.sampled_from([True, False]),
        is_dense_b=st.sampled_from([True, False]),
        dtype=st.sampled_from(get_test_dtypes([torch.bfloat16, torch.float32])),
    )
    @settings(
        verbosity=Verbosity.verbose,
        max_examples=20,
        deadline=None,
    )
    # pyre-ignore[2]
    def test_concat_2D_jagged_large_tensor(self, *args, **kwargs) -> None:
        self._test_concat_2D_jagged(
            *args,
            **kwargs,
            test_backward=True,
            skip_comparisons=True,
            ref_kernel=Kernel.TRITON,
            real_kernel=Kernel.TRITON,
        )

    def _test_concat_2D_jagged(
        self,
        batch_size: int,
        max_len_a: int,
        max_len_b: int,
        D: int,
        is_dense_a: bool,
        is_dense_b: bool,
        test_backward: bool,
        ref_kernel: Kernel,
        real_kernel: Kernel,
        dtype: torch.dtype = torch.float32,
        skip_comparisons: bool = False,
    ) -> None:
        from tzrec.ops.jagged_tensors import concat_2D_jagged

        if not is_dense_a:
            lengths_a = torch.randint(
                1, max_len_a + 1, size=(batch_size,), device=torch.device("cuda")
            )
            offsets_a = torch.zeros(
                (batch_size + 1,), dtype=torch.int64, device=torch.device("cuda")
            )
            offsets_a[1:] = torch.cumsum(lengths_a, dim=0)
            total_len_a = int(offsets_a[-1].item())
        else:
            offsets_a = None
            total_len_a = batch_size * max_len_a
        values_a = (
            torch.empty(
                (total_len_a, D),
                dtype=dtype,
                device=torch.device("cuda"),
            )
            .uniform_(-1.0, 1.0)
            .requires_grad_()
        )

        if not is_dense_b:
            lengths_b = torch.randint(
                1, max_len_b + 1, size=(batch_size,), device=torch.device("cuda")
            )
            offsets_b = torch.zeros(
                (batch_size + 1,), dtype=torch.int64, device=torch.device("cuda")
            )
            offsets_b[1:] = torch.cumsum(lengths_b, dim=0)
            total_len_b = int(offsets_b[-1].item())
        else:
            offsets_b = None
            total_len_b = batch_size * max_len_b
        values_b = (
            torch.empty(
                (total_len_b, D),
                dtype=dtype,
                device=torch.device("cuda"),
            )
            .uniform_(-1.0, 1.0)
            .requires_grad_()
        )

        ref_values = concat_2D_jagged(
            values_left=values_a,
            values_right=values_b,
            max_len_left=max_len_a,
            max_len_right=max_len_b,
            offsets_left=offsets_a,
            offsets_right=offsets_b,
            kernel=ref_kernel,
        )
        dout = torch.randn_like(ref_values)
        ref_values.backward(dout)
        if skip_comparisons:
            return

        assert values_a.grad is not None
        ref_d_a, values_a.grad = values_a.grad.clone(), None
        assert values_b.grad is not None
        ref_d_b, values_b.grad = values_b.grad.clone(), None

        values_a = values_a.detach().clone().requires_grad_()
        values_b = values_b.detach().clone().requires_grad_()
        dout = dout.detach().clone()
        real_values = concat_2D_jagged(
            values_left=values_a,
            values_right=values_b,
            max_len_left=max_len_a,
            max_len_right=max_len_b,
            offsets_left=offsets_a,
            offsets_right=offsets_b,
            kernel=real_kernel,
        )
        torch.testing.assert_close(ref_values, real_values)

        if test_backward:
            real_values.backward(dout)
            real_d_a = values_a.grad.clone()
            real_d_b = values_b.grad.clone()
            torch.testing.assert_close(ref_d_a, real_d_a)
            torch.testing.assert_close(ref_d_b, real_d_b)

    @unittest.skipIf(*gpu_unavailable)
    # pyre-ignore
    @given(
        batch_size=st.integers(4, 8),
        max_seq_len=st.integers(50, 500),
        D=st.integers(20, 200),
        K=st.integers(30, 200),
        dtype=st.sampled_from(
            get_test_dtypes([torch.float32, torch.bfloat16, torch.float16])
        ),
        contiguous=st.booleans(),
    )
    @settings(
        verbosity=Verbosity.verbose,
        max_examples=20,
        deadline=None,
    )
    # pyre-ignore[2]
    def test_jagged_dense_bmm_broadcast_add_triton(self, *args, **kwargs) -> None:
        self._test_jagged_dense_bmm_broadcast_add(
            *args,
            **kwargs,
            test_backward=True,
            atol=None,
            rtol=None,
            ref_kernel=Kernel.PYTORCH,
            real_kernel=Kernel.TRITON,
        )

    @unittest.skipIf(*gpu_unavailable)
    # pyre-ignore
    @given(
        batch_size=st.sampled_from([130]),
        max_seq_len=st.sampled_from([32768]),
        D=st.sampled_from([128]),
        K=st.sampled_from([128]),
        dtype=st.sampled_from([torch.float32, torch.bfloat16]),
        contiguous=st.booleans(),
    )
    @settings(
        verbosity=Verbosity.verbose,
        max_examples=2,
        deadline=None,
    )
    def test_jagged_dense_bmm_broadcast_add_triton_large_tensor(
        self,
        # pyre-fixme[2]: Parameter must be annotated.
        *args,
        **kwargs,  # pyre-ignore[2]
    ) -> None:
        self._test_jagged_dense_bmm_broadcast_add(
            *args,
            **kwargs,
            test_backward=True,
            atol=None,
            rtol=None,
            ref_kernel=Kernel.TRITON,
            real_kernel=Kernel.TRITON,
        )

    def _test_jagged_dense_bmm_broadcast_add(
        self,
        batch_size: int,
        max_seq_len: int,
        D: int,
        K: int,
        dtype: torch.dtype,
        ref_kernel: Kernel,
        real_kernel: Kernel,
        test_backward: bool,
        contiguous: bool = True,
        atol: Optional[float] = None,
        rtol: Optional[float] = None,
        sparsity: float = -1,
    ) -> None:
        torch.backends.cudnn.allow_tf32 = False
        torch.backends.cuda.matmul.allow_tf32 = False
        from tzrec.ops.jagged_tensors import (
            jagged_dense_bmm_broadcast_add,
        )

        if sparsity > 0.0:
            lengths = generate_sparse_seq_len(
                size=batch_size,
                max_seq_len=max_seq_len,
                sparsity=sparsity,
                device=torch.device("cuda"),
            ).to(torch.int64)
        else:
            lengths = torch.randint(
                max_seq_len + 1, size=(batch_size,), device=torch.device("cuda")
            )
        # Test the edge case with an empty row
        seq_offsets = torch.zeros(
            (batch_size + 1,), dtype=torch.int64, device=torch.device("cuda")
        )
        seq_offsets[1:] = torch.cumsum(lengths, dim=0)
        jagged_size = int(seq_offsets[-1].item())
        jagged = (
            torch.empty((jagged_size, D), dtype=dtype, device=torch.device("cuda"))
            .uniform_(-1.0, 1.0)
            .requires_grad_()
        )
        dense = (
            torch.empty((batch_size, D, K), dtype=dtype, device=torch.device("cuda"))
            .uniform_(-1.0, 1.0)
            .requires_grad_()
        )
        bias = (
            torch.empty((batch_size, K), dtype=dtype, device=torch.device("cuda"))
            .uniform_(-1.0, 1.0)
            .requires_grad_()
        )

        if not contiguous:
            dense = (
                dense.transpose(1, 2)
                .contiguous()
                .transpose(1, 2)
                .detach()
                .clone()
                .requires_grad_()
            )

        ref_out = jagged_dense_bmm_broadcast_add(
            max_seq_len=max_seq_len,
            seq_offsets=seq_offsets,
            jagged=jagged,
            dense=dense,
            bias=bias,
            kernel=ref_kernel,
        ).to(jagged.dtype)
        if test_backward:
            dout = torch.randn_like(ref_out) * 0.01
            ref_out.backward(dout)
            # pyre-ignore
            ref_d_jagged, jagged.grad = jagged.grad.clone(), None
            ref_d_dense, dense.grad = dense.grad.clone(), None
            ref_d_bias, bias.grad = bias.grad.clone(), None

        jagged = jagged.detach().clone().requires_grad_()
        dense = dense.detach().clone().requires_grad_()
        bias = bias.detach().clone().requires_grad_()
        real_out = jagged_dense_bmm_broadcast_add(
            max_seq_len=max_seq_len,
            seq_offsets=seq_offsets,
            jagged=jagged,
            dense=dense,
            bias=bias,
            kernel=real_kernel,
        )
        torch.testing.assert_close(
            ref_out,
            real_out,
            atol=atol,
            rtol=rtol,
        )
        if test_backward:
            real_out.backward(dout)  # pyre-ignore
            real_d_jagged = jagged.grad.clone()
            real_d_dense = dense.grad.clone()
            real_d_bias = bias.grad.clone()
            torch.testing.assert_close(
                ref_d_jagged,  # pyre-ignore
                real_d_jagged,
                atol=atol,
                rtol=rtol,
            )
            torch.testing.assert_close(
                ref_d_dense,  # pyre-ignore
                real_d_dense,
                atol=atol,
                rtol=rtol,
            )
            torch.testing.assert_close(
                ref_d_bias,  # pyre-ignore
                real_d_bias,
                atol=atol,
                rtol=rtol,
            )


if __name__ == "__main__":
    unittest.main()
