import torch


def cal_diff(x: torch.Tensor, y: torch.Tensor) -> tuple[float, float, float, float]:
    if torch.equal(x, y):
        return 0.0, 0.0, 0.0, 0.0
    x, y = x.double(), y.double()
    rmse = ((x - y) * (x - y)).mean().sqrt().item()
    cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12)
    chebyshev_dis = (x - y).abs().max().item()
    kl = ((torch.exp(x - y) - 1) - (x - y)).mean().item()
    return cos_diff, rmse, chebyshev_dis, kl


def assert_close(x: torch.Tensor, y: torch.Tensor, name: str, eps: float = 1e-4) -> None:
    assert x.shape == y.shape, f"{x.shape=}, {y.shape=}"
    assert not x.isnan().any() and not y.isnan().any()
    x, y = x.clone(), y.clone()
    x[x.isfinite().logical_not()] = 0
    y[y.isfinite().logical_not()] = 0
    cos_diff, rmse, chebyshev_dis, kl = cal_diff(x, y)
    print(f"{name}: {cos_diff=}, {rmse=}, {chebyshev_dis=}, {kl=}")
    assert cos_diff < eps, f"{name}: {cos_diff=}, {rmse=}, {chebyshev_dis=}, {kl=}"