import math

import torch


def assert_equal(
    x: torch.Tensor,
    y: torch.Tensor,
    check_dtype: bool = True,
    check_shape: bool = True,
    check_stride: bool = True,
) -> None:
    assert not check_dtype or x.dtype == y.dtype, f'Tensor dtypes are not equal: {x.dtype} vs {y.dtype}'
    assert not check_shape or x.shape == y.shape, f'Tensor shapes are not equal: {x.shape} vs {y.shape}'
    assert not check_stride or x.numel() == 0 or x.stride() == y.stride(), (
        f'Tensor strides are not equal: {x.stride()} vs {y.stride()}'
    )
    assert x.device == y.device, f'Tensor devices are not equal: {x.device} vs {y.device}'
    # Hints: The tensor with a size of [32768, 1] and a stride of [1, 32768] is considered contiguous,
    # but using .view will cause an error. Therefore, .flatten is used to ensure the stride of the last dimension is 1.
    x_bytes = x.contiguous().flatten().as_strided((x.numel(),), (1,)).view(torch.uint8)
    y_bytes = y.contiguous().flatten().as_strided((y.numel(),), (1,)).view(torch.uint8)
    if torch.equal(x_bytes, y_bytes):
        return

    try:
        mask = x != y
        detail = f'mask={torch.nonzero(mask)}\n{x[mask]}\nvs\n{y[mask]}'
    except RuntimeError:
        # Some dtypes (e.g. float8_e4m3fn on Ascend) don't support `!=`; compare raw bytes.
        byte_mask = x_bytes != y_bytes
        detail = f'byte_mask={torch.nonzero(byte_mask)}\n{x_bytes[byte_mask]}\nvs\n{y_bytes[byte_mask]}'
    raise AssertionError(f'Tensor values are not equal: {x.shape=} vs {y.shape=}\n{detail}')


def calc_diff(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    if x.device.type in ('npu',):
        x, y = x.to('cpu'), y.to('cpu')
    x, y = x.double(), y.double()
    denominator = (x * x + y * y).sum()
    sim = 2 * (x * y).sum() / denominator
    return 1 - sim if denominator != 0 else 0


def check_bias(x: torch.Tensor, ref_x: torch.Tensor) -> None:
    count = x.numel()
    if count == 0:
        return
    less_count = (x < ref_x).sum()
    equal_count = (x == ref_x).sum()
    less_ratio = (less_count + equal_count / 2) / count

    # Suppose whether a number is larger or smaller after casting is an independent variable
    # and has exact possibility of 0.5 to be larger.
    # Then `less_count` follows a binomial distribution B(count, 0.5)
    # The standard deviation is sqrt(count * 0.5 * 0.5) = sqrt(count) / 2
    # Then `less_ratio` has a standard deviation of sqrt(count) / (2 * count) = 1 / (2 * sqrt(count))
    # When `count` is large enough, the central limit theorem applies, and we have:
    # less_ratio` ~ N(0.5, 1 / (4 * count))
    # So 99.99999% confidence interval should be something like this:
    # (-c / sqrt(count), c / sqrt(count)) around 0.5
    allowed_diff_ratio = 10 / math.sqrt(x.numel())
    assert abs(less_ratio - 0.5) < allowed_diff_ratio, (
        f'Less than ratio not close to 0.5 (size = {x.numel()}): {less_ratio=:.4f}\n'
        f'Expected:\n  {ref_x.view(-1, 4)}\n'
        f'Actual:\n   {x.view(-1, 4)}\n'
        f'  Less than: {less_count}\n  Equal to:  {equal_count}\n  Greater than: {count - less_count - equal_count}\n'
    )


def count_bytes(*tensors) -> int:
    total = 0
    for t in tensors:
        if isinstance(t, tuple | list):
            total += count_bytes(*t)
        elif t is not None:
            total += t.numel() * t.element_size()
    return total