# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

import math
import os
import struct
import time
from typing import Union

import numpy as np
import torch
from loguru import logger
from ttnn.device import Arch
from typing_extensions import deprecated

import ttnn


def get_mesh_device():
    """Fixture to provide mesh device configuration."""
    mesh_device = os.environ.get("MESH_DEVICE", "N150")
    mesh_config = {
        "N150": (1, 1),
        "N300": (2, 1),
        "T3K": (8, 1),
        "TG": (8, 4),
    }.get(mesh_device, (ttnn.get_num_devices(), 1))
    return mesh_config


### Math operations ###
def _nearest_32(x):
    return math.ceil(x / 32) * 32


def nearest_32(
    x,
):  # needs refctoring; to match alias called in some scripts (e.g. test_padding_test in unit tests)
    return _nearest_32(x)


def _nearest_y(x, y):
    return math.ceil(x / y) * y


def nearest_y(x, y):
    return _nearest_y(x, y)


def divup(a, b):
    return (a + b - 1) // b


def roundup(a, b):
    result = divup(a, b) * b
    return result


def roundup32(a):
    return roundup(a, 32)


def float_to_bits(x):
    s = struct.pack(">f", x)
    return struct.unpack(">l", s)[0]


def torch_random(shape, low, high, dtype):
    if dtype in [torch.int64, torch.int32, torch.int16, torch.int8]:
        return torch.randint(low, high, shape, dtype=dtype)
    return torch.zeros(shape, dtype=dtype).uniform_(low, high)


def torch_random_with_zeros(shape, low, high, dtype, zero_fraction=0.1):
    total_elements = torch.prod(torch.tensor(shape)).item()
    num_zeros = int(total_elements * zero_fraction)
    num_random = total_elements - num_zeros

    # Generate random values between low and high
    random_values = torch.empty(num_random).uniform_(low, high)
    zeros = torch.zeros(num_zeros)

    # Combine zeros and random values
    combined = torch.cat([zeros, random_values])

    # Shuffle the tensor
    shuffled = combined[torch.randperm(combined.size(0))]

    # Reshape to the desired shape
    result_tensor = shuffled.view(shape)
    result_tensor.to(dtype)
    return result_tensor


### Profiling ###
class Profiler:
    def __init__(self):
        self.start_times = dict()
        self.times = dict()
        self.disabled = False

    def clear(self):
        self.start_times = dict()
        self.times = dict()
        self.disabled = False

    def enable(self):
        self.disabled = False

    def disable(self):
        self.disabled = True

    def start(self, key, force_enable=False):
        if self.disabled and not force_enable:
            return

        self.start_times[key] = time.time()

    def end(self, key, PERF_CNT=1, force_enable=False):
        if self.disabled and not force_enable:
            return

        if key not in self.start_times:
            return

        diff = time.time() - self.start_times[key]

        if key not in self.times:
            self.times[key] = []

        self.times[key].append(diff / PERF_CNT)

    def get(self, key):
        if key not in self.times:
            return 0

        return sum(self.times[key]) / len(self.times[key])

    def print(self, units="s"):
        for key in self.times:
            average = self.get(key)
            if units == "s":
                pass
            elif units == "ms":
                average *= 1000
            elif units == "us":
                average *= 1000000
            elif units == "ns":
                average *= 1000000000
            else:
                raise ValueError(f"Invalid units: {units}")
            print(f"{key}: {average:.3f}{units}")


profiler = Profiler()


### Turn flags on/off ###
def enable_memory_reports():
    """
    Enables generating reports of memory allocation statistics in .reports/tt_metal dir
    """
    return ttnn.device.EnableMemoryReports()


def disable_memory_reports():
    """
    Disables generating reports of memory allocation statistics
    """
    return ttnn.device.DisableMemoryReports()


### Tensor conversion ###
def torch2tt_tensor(
    py_tensor: torch.Tensor,
    tt_device,
    tt_layout=ttnn.TILE_LAYOUT,
    tt_memory_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED),
    tt_dtype=ttnn.bfloat16,
):
    size = list(py_tensor.size())

    while len(size) < 4:
        size.insert(0, 1)

    tt_tensor = ttnn.Tensor(py_tensor.reshape(size), tt_dtype)
    tt_tensor = tt_tensor.to(tt_layout)

    if tt_device is not None:
        tt_tensor = tt_tensor.to(tt_device, tt_memory_config)
    else:
        tt_tensor = tt_tensor.cpu()

    return tt_tensor


def tt_tensors_to_torch_tensors(
    tt_tensors_device: ttnn.Tensor, mesh_device: Union[ttnn.MeshDevice, ttnn.Device], concat_dim: int = 0
):
    # Convert tensors to interleaved
    if tt_tensors_device.is_sharded():
        tt_tensors_device = ttnn.sharded_to_interleaved(tt_tensors_device)

    # Convert tensors to RM layout
    if tt_tensors_device.layout == ttnn.TILE_LAYOUT:
        # Convert to bfloat16 to ensure untilize works
        if tt_tensors_device.dtype != ttnn.bfloat16:
            tt_tensors_device = ttnn.clone(
                tt_tensors_device, dtype=ttnn.bfloat16, memory_config=ttnn.DRAM_MEMORY_CONFIG
            )
        # Untilize using singlecore since multicore version runs out of l1 memory (Issue #9022)
        tt_tensors_device = ttnn.untilize(tt_tensors_device, use_multicore=False)

    return torch.cat([t.to_torch() for t in ttnn.get_device_tensors(tt_tensors_device.cpu())], dim=concat_dim)


def tt2torch_tensor(tt_tensor):
    tt_output = tt_tensor.cpu()
    if tt_output.get_layout() != ttnn.ROW_MAJOR_LAYOUT:
        tt_output = tt_output.to(ttnn.ROW_MAJOR_LAYOUT)
    return tt_output.to_torch()


def tt_to_torch_tensor(tt_tensor):
    tt_output = tt_tensor.cpu().to(ttnn.ROW_MAJOR_LAYOUT)
    return tt_output.to_torch()


def torch_to_tt_tensor_rm(py_tensor, device, shape=None, put_on_device=True):
    if shape is None:
        shape = list(py_tensor.size())
        while len(shape) < 4:
            shape.insert(0, 1)

    tt_tensor = ttnn.Tensor(py_tensor.reshape(shape), ttnn.bfloat16)
    if put_on_device:
        tt_tensor = tt_tensor.to(device)
    return tt_tensor


def torch_to_tt_tensor(py_tensor, device):
    shape = list(py_tensor.size())
    while len(shape) < 4:
        shape.insert(0, 1)

    tt_tensor = (
        ttnn.Tensor(py_tensor.reshape(shape), ttnn.bfloat16)
        .to(
            ttnn.TILE_LAYOUT
        )  # change memory layout of TT Tensor to TILE (as operation that will use it expects TILE layout)
        .to(device)  # move TT Tensor from host to TT accelerator device (device is of type ttnn.device.Device)
    )

    return tt_tensor


def unpad_from_zero(x, desired_shape):
    if x.padded_shape[-1] == desired_shape[-1] and x.padded_shape[-2] == desired_shape[-2]:
        x = tt2torch_tensor(x)
    else:
        x = x.cpu()
        if x.get_layout() != ttnn.ROW_MAJOR_LAYOUT:
            x = x.to(ttnn.ROW_MAJOR_LAYOUT)
        x = x.unpad(
            (0, 0, 0, 0),
            (
                desired_shape[0],
                desired_shape[1],
                desired_shape[2],
                desired_shape[3],
            ),
        )

        x = x.to_torch()
    return x


def pad_activation(x):
    """
    This function pads an activation with 0s as a pre-preprocessing step to tilization.

    In the 2d case, it pads a vector to the right with 0s, and in the 2+d case,
    it pads the bottom and right corners of the last two dimensions.

    :param x: Input PyTorch Tensor
    :type x: class:`torch.Tensor`

    WARNING: This function should eventually be retired in favour of padding on device
    """
    nearest_32 = _nearest_32

    assert isinstance(x, torch.Tensor), "Input to this function must be an instance of torch.Tensor"
    assert len(x.shape) >= 1 and len(x.shape) <= 4, "Only tensors with dimension 1-4 supported"
    if len(x.shape) == 1:  # (num_features,)
        padded_tensor = torch.zeros(1, 1, 32, nearest_32(x.shape[0]))
        padded_tensor[:, 0, 0, : x.shape[0]] = x
    elif len(x.shape) == 2:  # (batch, num features)
        padded_tensor = torch.zeros(x.shape[0], 1, 32, nearest_32(x.shape[1]))
        padded_tensor[:, 0, 0, : x.shape[1]] = x
    elif len(x.shape) == 3:  # (batch, num features y, num features x)
        padded_tensor = torch.zeros(x.shape[0], 1, nearest_32(x.shape[-2]), nearest_32(x.shape[-1]))
        padded_tensor[..., 0, : x.shape[-2], : x.shape[-1]] = x
    else:  # (batch, num channels, num features y, num features x)
        padded_tensor = torch.zeros(*x.shape[:-2], nearest_32(x.shape[-2]), nearest_32(x.shape[-1]))
        padded_tensor[..., : x.shape[-2], : x.shape[-1]] = x
    return padded_tensor


def pad_weight(x):
    """
    This function pads a weight/bias with 0s as a pre-preprocessing step to tilization.

    tt_tensor = ttnn.Tensor(
        py_tensor.reshape(shape), ttnn.bfloat16
    In the 2d case, it pads a vector to the right with 0s, and in the 2+d case,
    it pads the bottom and right corners of the last two dimensions.

    :param x: Input PyTorch Tensor
    :type x: class:`torch.Tensor`

    WARNING: This function should eventually be retired in favour of padding on device
    """
    nearest_32 = _nearest_32

    assert isinstance(x, torch.Tensor), "Input to this function must be an instance of torch.Tensor"
    assert len(x.shape) >= 1 and len(x.shape) <= 4, "Only tensors with dimension 1-4 supported"

    if len(x.shape) == 1:  # (num_features,)
        padded_tensor = torch.zeros(1, 1, 32, nearest_32(x.shape[0]))
        padded_tensor[:, 0, 0, : x.shape[0]] = x
    elif len(x.shape) == 2:  # (r_features, c_features)
        padded_tensor = torch.zeros(1, 1, nearest_32(x.shape[0]), nearest_32(x.shape[1]))
        padded_tensor[:, 0, : x.shape[0], : x.shape[1]] = x
    else:
        padded_tensor = torch.zeros(*x.shape[:-2], nearest_32(x.shape[-2]), nearest_32(x.shape[-1]))
        padded_tensor[..., : x.shape[-2], : x.shape[-1]] = x

    return padded_tensor


def convert_weights_2d_matrix(weights, w_shape):
    """
    :param weights: Input PyTorch Tensor
    :type weights: class:`torch.Tensor`
    """
    ret_shape = [1, 1, w_shape[0], w_shape[1] * w_shape[2] * w_shape[3]]
    if isinstance(weights, torch.Tensor):
        ret = torch.zeros(np.prod(ret_shape))
    else:
        ret = np.zeros(np.prod(ret_shape))
    idx = 0
    for k in range(w_shape[0]):
        for r in range(w_shape[2]):
            for s in range(w_shape[3]):
                for c in range(w_shape[1]):
                    ret[idx] = weights[k][c][r][s]
                    idx += 1
    assert idx == np.prod(ret_shape)
    return ret.reshape(ret_shape).transpose(2, 3)


def convert_act_2d_matrix(activation, kernel_y, kernel_x, stride_y, stride_x, pad_y, pad_x):
    """
    :param activation: Input PyTorch Tensor
    :type activation: class:`torch.Tensor`
    """
    N = activation.shape[0]
    C = activation.shape[1]
    H = activation.shape[2]
    W = activation.shape[3]

    OH = (int)((H - kernel_y + 2 * pad_y) // stride_y) + 1
    OW = ((W - kernel_x + 2 * pad_x) // stride_x) + 1
    nrows = OH * OW
    ncols = C * kernel_x * kernel_y
    ret_shape = [1, N, nrows, ncols]
    if isinstance(activation, torch.Tensor):
        ret = torch.zeros(np.prod(ret_shape))
    else:
        ret = np.zeros(np.prod(ret_shape))
    idx = 0
    for n in range(N):
        for h in range(-1 * pad_y, H + pad_y - kernel_y + 1, stride_y):
            for w in range(-1 * pad_x, W + pad_x - kernel_x + 1, stride_x):
                for r in range(kernel_y):
                    for s in range(kernel_x):
                        for c in range(C):
                            h_offs = h + r
                            w_offs = w + s
                            pad = h_offs < 0 or h_offs >= H or w_offs < 0 or w_offs >= W
                            ret[idx] = 0 if pad else activation[n][c][h_offs][w_offs]
                            idx += 1
    assert idx == np.prod(ret_shape)
    return ret.reshape(ret_shape)


### Tilizing / Untilizing ###
@deprecated("PyTorch data is handled automatically in tensor infra. This function does nothing now:")
def tilize(x):
    return x


@deprecated("PyTorch data is handled automatically in tensor infra. This function does nothing now:")
def tilize_to_list(x):
    """
    Returns a flattened list of the tensor
    """
    return tilize(x).reshape(-1).tolist()


@deprecated("PyTorch data is handled automatically in tensor infra. This function does nothing now:")
def untilize(x):
    return x


### Measuring accuracy and other metrics ###
def is_close(a, b, rtol=1e-2, atol=1e-2, max_mag=2.0, max_mag_fraction=0.02):
    """
    A variant of np.isclose with logging.
    """
    absdiff = (a - b).abs()
    reldiff1 = (a.abs() / b.abs()) - 1.0
    reldiff2 = (a.abs() + 1.0) / (b.abs() + 1.0) - 1.0  # in case b.abs() is 0
    reldiff_or = torch.logical_or(reldiff1.abs() < rtol, reldiff2.abs() < rtol)
    max_mag_ok = absdiff < max_mag * max_mag_fraction

    or_abs_rel = torch.logical_or(absdiff < atol, reldiff_or)
    or_abs_rel = torch.logical_or(or_abs_rel, max_mag_ok)
    debug_index = or_abs_rel.to(torch.int32).argmin().item()

    if not or_abs_rel.reshape(-1)[debug_index]:
        logger.info(f"isclose mismatch at index={debug_index}")
        logger.info(a.reshape(-1)[debug_index])
        logger.info(b.reshape(-1)[debug_index])
        logger.info(f"reldiff1={reldiff1.reshape(-1)[debug_index]}")
        logger.info(f"reldiff2={reldiff2.reshape(-1)[debug_index]}")
        logger.info(f"absdiff={absdiff.reshape(-1)[debug_index]}")

        HT = a.shape[-2] // 32
        WT = a.shape[-1] // 32
        hwt = debug_index // 1024
        wt = hwt % WT
        ht = hwt // WT
        h = (debug_index % 1024) // 32
        w = (debug_index % 1024) % 32

        logger.info(f"****    at {debug_index} --- HTWT={ht} {wt} HW={h} {w}")

    return torch.all(or_abs_rel)


def _comp_nonfinite(golden, calculated):
    """
    Returns True if tensors contain the same non-finite values (nan, inf, -inf) at the same positions. Also returns True if all elements are finite.
    Returns False if non-finite values differ between both tensors.
    """

    # torch.equal(['nan'], ['nan']] => False
    # For this reason, we check for nan and inf separately
    if torch.not_equal(torch.isnan(golden), torch.isnan(calculated)).any():
        return False

    golden_inf_mask = torch.isinf(golden)
    calculated_inf_mask = torch.isinf(calculated)

    if torch.not_equal(golden_inf_mask, calculated_inf_mask).any():
        return False

    golden_inf = golden[golden_inf_mask]
    calculated_inf = calculated[calculated_inf_mask]
    return torch.equal(golden_inf, calculated_inf)


def comp_allclose(golden, calculated, rtol=1e-05, atol=1e-08):
    if golden.dtype != calculated.dtype:
        calculated = calculated.type(golden.dtype)

    atol_delta = torch.max(torch.abs(golden - calculated)).item()
    rtol_delta = torch.max(torch.abs(golden - calculated) / torch.abs(calculated)).item()
    return (
        torch.allclose(golden, calculated, rtol, atol, True),
        f"Max ATOL Delta: {atol_delta}, Max RTOL Delta: {rtol_delta}",
    )


def comp_pcc(golden, calculated, pcc=0.99, rtol=1e-05, atol=1e-04):
    golden = torch.Tensor(golden)
    calculated = torch.Tensor(calculated)

    if golden.dtype != calculated.dtype:
        calculated = calculated.type(golden.dtype)

    if torch.all(torch.isnan(golden)) and torch.all(torch.isnan(calculated)):
        logger.warning("Both tensors are 'nan'")
        return True, 1.0

    if torch.all(torch.isnan(golden)) or torch.all(torch.isnan(calculated)):
        logger.error("One tensor is all nan, the other is not.")
        return False, 0.0

    # Test if either is completely zero — but a zero tensor is also a constant tensor,
    # so fall back to allclose instead of a hard 0.0: zero-vs-small-constant may be
    # within the caller's tolerances.
    if torch.any(golden.bool()) != torch.any(calculated.bool()):
        logger.warning("One tensor is all zero. PCC undefined; falling back to allclose.")
        result = torch.allclose(golden, calculated, rtol=rtol, atol=atol)
        return result, float(result)

    golden = torch.squeeze(golden).flatten()
    calculated = torch.squeeze(calculated).flatten()

    # For now, mask all infs and nans (to zero) so that we check the rest... TODO
    # Skip this for integer types which don't have NaN/Inf values.
    if golden.dtype.is_floating_point:
        # FP8 doesn't support isfinite/nan_to_num and bfloat16 products lose precision,
        # so correlate these in float32.
        if golden.dtype in (torch.float8_e4m3fn, torch.float8_e5m2, torch.bfloat16):
            golden = golden.to(torch.float32)
            calculated = calculated.to(torch.float32)

        # Zero out NaN/Inf, preserving the historical PCC values. nan_to_num allocates a
        # full-size copy of each tensor, so only do it when invalid values are actually
        # present; on the common all-finite path the tensors stay as views and no copy is
        # made (this short-circuit is what keeps peak memory near 1x of one input).
        if not bool((torch.isfinite(golden) & torch.isfinite(calculated)).all()):
            golden = torch.nan_to_num(golden, nan=0.0, posinf=0.0, neginf=0.0)
            calculated = torch.nan_to_num(calculated, nan=0.0, posinf=0.0, neginf=0.0)

    if torch.equal(golden, calculated):
        return True, 1.0

    # Integer tensors must be correlated in floating point (centering/products would
    # otherwise truncate/overflow). float32 keeps the working set small.
    if not golden.dtype.is_floating_point:
        golden = golden.to(torch.float32)
        calculated = calculated.to(torch.float32)

    # Pearson r with float64 *accumulation* (dtype= on the reductions) over the float32
    # data: no float64 copy of either tensor is materialized, so peak memory stays near
    # 1x of one input on large tensors while matching a full-float64 correlation to
    # |Δ|<1e-9 across the high-PCC (>=0.999) range.
    n = golden.numel()
    g_centered = golden - (golden.sum(dtype=torch.float64) / n).to(golden.dtype)
    c_centered = calculated - (calculated.sum(dtype=torch.float64) / n).to(calculated.dtype)
    cov = (g_centered * c_centered).sum(dtype=torch.float64)
    g_sq_sum = g_centered.pow(2).sum(dtype=torch.float64)
    c_sq_sum = c_centered.pow(2).sum(dtype=torch.float64)
    denom = torch.sqrt(g_sq_sum * c_sq_sum)
    # pow/sum stay in float32 before the reduction; large-magnitude tensors (e.g. ldexp)
    # can overflow to inf here even though float64 accumulation would be finite.
    if not math.isfinite(denom.item()) or not math.isfinite(cov.item()):
        g_centered64 = g_centered.to(torch.float64)
        c_centered64 = c_centered.to(torch.float64)
        cov = (g_centered64 * c_centered64).sum()
        denom = torch.sqrt(g_centered64.pow(2).sum() * c_centered64.pow(2).sum())
    cal_pcc = (cov / denom).item()

    # Zero variance -> denom == 0 -> cal_pcc is nan: PCC is undefined for constant tensors.
    # Fall back to allclose rather than returning a misleading 1.0.
    if math.isnan(cal_pcc):
        logger.warning("PCC is NaN (zero variance / constant tensor). Falling back to allclose check.")
        result = torch.allclose(golden, calculated, rtol=rtol, atol=atol)
        return result, float(result)

    return cal_pcc >= pcc, cal_pcc


def ulp(x: Union[ttnn.Tensor, torch.Tensor]) -> Union[ttnn.Tensor, torch.Tensor]:
    "Return Unit of Least Precision for each element of a given tensor"

    received_ttnn_input = False
    if isinstance(x, ttnn.Tensor):
        x = ttnn.to_torch(x)
        received_ttnn_input = True

    # Notes:
    # - This should be identical to the definition of ULP by Goldberg
    #   "What every computer scientist should know about floating-point arithmetic"
    #   https://docs.oracle.com/cd/E19957-01/806-3568/ncg_goldberg.html
    # - We use torch.abs(x) to ensure symmetry ULP(-x) == ULP(x)
    # - For x powers of 2, x + ULP(x) is not closest number but second closest (previous number is 2x closer)
    #   However, this avoids rounding-to-nearest-tie-to-even issues on addition (i.e. x + ULP(x) != x)
    abs_x = torch.abs(x)
    next = torch.nextafter(
        abs_x, torch.tensor(math.inf, dtype=x.dtype)
    )  # 1 ULP ~ Difference between two consecutive floating point numbers
    ulp_value = next - abs_x

    # Special case: if abs_x == torch.finfo(x.dtype).max, then next == math.inf, which leads to ULP(x) == inf rather than finite number
    # We fix this problem by manually calculating ULP at max value, and masking tensor when input == max
    dtype_max = torch.finfo(x.dtype).max
    max_epsilon = dtype_max - torch.nextafter(
        torch.tensor(dtype_max, dtype=x.dtype), torch.tensor(-math.inf, dtype=x.dtype)
    )
    ulp_value = torch.where(abs_x == dtype_max, max_epsilon, ulp_value)

    if received_ttnn_input:  # Ensures that type(input) == type(output)
        ulp_value = ttnn.from_torch(ulp_value)

    return ulp_value


def comp_ulp(golden, calculated, ulp_threshold, allow_nonfinite=False):
    """
    Compute absolute error between two tensors in Units of Least Precision (ULP)
    """

    # If both tensors are empty, then we can return True
    if torch.numel(golden) == 0 and torch.numel(calculated) == 0:
        return True, "Both tensors are empty"

    if not allow_nonfinite and not torch.all(torch.isfinite(calculated)):
        return False, "Calculated tensor contains non-finite values"

    if not _comp_nonfinite(golden, calculated):
        return False, "Tensors are not finite at the same positions"
    # nonfinite elements can interfere with ULP error calculation
    # To avoid this, replace nan, +inf, -inf with 0
    # (we have already checked that both tensors have the same nonfinite elements)
    mask_finite = ~torch.isfinite(golden)
    golden = golden.clone()
    calculated = calculated.clone()
    golden[mask_finite] = 0
    calculated[mask_finite] = 0

    # ULP is measured according to the golden tensor
    # In most cases, data type of golden tensor should be the same as calculated tensor.
    # However, in some cases, we may want to measure < 1 ULP differences, which requires golden tensor
    # to have higher precision than calculated tensor.
    # If we passed golden tensor to ulp() as is, we would get ULP of higher precision.
    # e.g. ulp of float32 rather bfloat16 calculation, which would give us a wrong value.
    ulp_value = ulp(golden.type(calculated.dtype))

    if golden.dtype != calculated.dtype:  # Note: assumes that golden has higher precision than calculated tensor
        calculated = calculated.type(golden.dtype)
        ulp_value = ulp_value.type(golden.dtype)  # Convert ULP to higher precision (for sub-1 ULP measurements)

    ulp_tensor = torch.abs(calculated - golden) / ulp_value
    ulp_delta = torch.max(ulp_tensor)
    within_threshold = ulp_delta <= ulp_threshold
    message = f"Max ULP Delta: {ulp_delta}"
    if not within_threshold:
        ulp_index = torch.argmax(ulp_tensor)
        ulp_index_tuple = tuple(int(idx) for idx in torch.unravel_index(ulp_index, golden.shape))
        message += (
            f" @ {list(ulp_index_tuple)} = "
            f"|calculated {calculated[ulp_index_tuple]} - golden {golden[ulp_index_tuple]}| "
            f"/ ULP(golden) {ulp_value[ulp_index_tuple]}"
        )
    return (within_threshold, message)


def calculate_detailed_ulp_stats(expected, actual):
    """
    Calculate detailed ULP statistics for analysis.

    Returns:
        dict: Dictionary with ULP statistics including max, mean, std, and percentiles
    """
    if isinstance(actual, ttnn.Tensor):
        actual = ttnn.to_torch(actual)
    if isinstance(expected, ttnn.Tensor):
        expected = ttnn.to_torch(expected)

    # Convert to bfloat16 if not already
    expected = expected.to(torch.bfloat16)
    actual = actual.to(torch.bfloat16)

    # Handle special cases
    if torch.allclose(expected, actual, rtol=0, atol=0, equal_nan=True):
        return {
            "max_ulp": 0.0,
            "mean_ulp": 0.0,
            "median_ulp": 0.0,
            "std_ulp": 0.0,
            "p95_ulp": 0.0,
            "p99_ulp": 0.0,
            "perfect_matches": 1.0,
        }

    # Convert bfloat16 to uint16 representation for bit manipulation
    expected_bits = expected.view(torch.int16).to(torch.int32)
    actual_bits = actual.view(torch.int16).to(torch.int32)

    # Handle sign differences
    expected_sign = expected_bits < 0
    actual_sign = actual_bits < 0
    same_sign = expected_sign == actual_sign

    # Calculate ULP differences
    expected_abs_bits = torch.where(expected_sign, -expected_bits, expected_bits)
    actual_abs_bits = torch.where(actual_sign, -actual_bits, actual_bits)

    ulp_diff = torch.where(same_sign, torch.abs(expected_bits - actual_bits), expected_abs_bits + actual_abs_bits)

    # Handle non-finite values
    expected_finite = torch.isfinite(expected)
    actual_finite = torch.isfinite(actual)
    both_finite = expected_finite & actual_finite

    ulp_diff = torch.where(both_finite, ulp_diff, torch.tensor(float("inf")))

    # Handle same non-finite values
    both_nan = torch.isnan(expected) & torch.isnan(actual)
    both_posinf = torch.isposinf(expected) & torch.isposinf(actual)
    both_neginf = torch.isneginf(expected) & torch.isneginf(actual)
    same_nonfinite = both_nan | both_posinf | both_neginf

    ulp_diff = torch.where(same_nonfinite, torch.tensor(0.0), ulp_diff)

    # Calculate statistics only on finite ULP differences
    finite_ulp = ulp_diff[torch.isfinite(ulp_diff)]

    if len(finite_ulp) == 0:
        return {
            "max_ulp": float("inf"),
            "mean_ulp": float("inf"),
            "median_ulp": float("inf"),
            "std_ulp": float("inf"),
            "p95_ulp": float("inf"),
            "p99_ulp": float("inf"),
            "perfect_matches": 0.0,
        }

    finite_ulp_float = finite_ulp.float()
    perfect_matches = (finite_ulp == 0).float().mean().item()

    return {
        "max_ulp": torch.max(finite_ulp).item(),
        "mean_ulp": torch.mean(finite_ulp_float).item(),
        "median_ulp": torch.median(finite_ulp_float).item(),
        "std_ulp": torch.std(finite_ulp_float).item(),
        "p95_ulp": torch.quantile(finite_ulp_float, 0.95).item(),
        "p99_ulp": torch.quantile(finite_ulp_float, 0.99).item(),
        "perfect_matches": perfect_matches,
    }


def comp_allclose_and_pcc(golden, calculated, rtol=1e-05, atol=1e-08, pcc=0.99):
    # 0-volume tensors are special because they don't have elements, so we can't compute PCC, etc.
    # If one of the tensors is a 0-volume tensor, simply call torch.equal to check if they are equal
    # (i.e. that both are 0-volume tensors and they have equal shapes).
    if golden.numel() == 0 or calculated.numel() == 0:
        return torch.equal(golden, calculated), f"{golden} != {calculated}"

    if golden.dtype != calculated.dtype:
        calculated = calculated.type(golden.dtype)

    passing = True
    output = ""
    passing_allclose, output_allclose = comp_allclose(golden, calculated, rtol, atol)
    passing &= passing_allclose
    output += output_allclose
    if torch.numel(golden) != 1:
        passing_pcc, output_pcc = comp_pcc(golden, calculated, pcc, rtol=rtol, atol=atol)
        passing &= passing_pcc
        output += f", pcc={output_pcc}"

    return passing, output


def comp_equal(golden, calculated):
    if golden.dtype != calculated.dtype:
        calculated = calculated.type(golden.dtype)

    # If either tensor is zero-volume, broadcasting can still yield an empty delta and
    # crash torch.max(); defer entirely to torch.equal (False on shape mismatch).
    if golden.numel() == 0 or calculated.numel() == 0:
        return torch.equal(golden, calculated), f"{golden} != {calculated}"

    atol_delta = torch.max(torch.abs(golden - calculated)).item()
    rtol_delta = torch.max(torch.abs(golden - calculated) / torch.abs(calculated)).item()
    return (
        torch.equal(golden, calculated),
        f"Max ATOL Delta: {atol_delta}, Max RTOL Delta: {rtol_delta}",
    )


def get_oom_of_float(float_lst):
    """
    Given a list of floats, returns a list of the order or magnitudes
    of the floats. Useful when you want to make sure that even if your
    tt outputs don't match pytorch all that well, they are at least
    on the same order of magnitude
    """
    ooms = []
    for el in float_lst:
        str_el = str(el)
        if "e" in str_el:
            oom = int(str_el.split("e")[1])
        elif str_el[:2] == "0.":
            str_el = str_el.split(".")[1]

            oom = -1
            for e in str_el:
                if e != "0":
                    break
                oom -= 1
        else:
            oom = len(str_el.split(".")[0])

        ooms.append(oom)

    return ooms


def print_diff_argmax(a, b, annotation=""):
    """
    Prints out the value of both tensors at a point where the absolute difference is the largest.
    """
    absdiff = (a - b).abs()
    argmax = absdiff.argmax().item()
    diff = absdiff.reshape(-1)[argmax]
    rela = a.abs() / (torch.max(a.abs(), b.abs()))
    relb = b.abs() / (torch.max(a.abs(), b.abs()))
    HT = a.shape[-2] // 32
    WT = a.shape[-1] // 32
    hwt = argmax // 1024
    wt = hwt % WT
    ht = hwt // WT
    h = (argmax % 1024) // 32
    w = (argmax % 1024) % 32
    print(
        "Abs diff=",
        diff,
        " at ",
        argmax,
        " --- ",
        annotation,
        "HTWT=",
        ht,
        wt,
        "HW=",
        h,
        w,
    )
    print("  (a=", a.reshape(-1)[argmax].item(), ")")
    print("  (b=", b.reshape(-1)[argmax].item(), ")")
    print("  Rel a=", rela.reshape(-1)[argmax], " at ", argmax)
    print("  Rel b=", relb.reshape(-1)[argmax], " at ", argmax)
    return diff.item()


def print_diff_tt_pyt(a, b, annotation=""):
    # first convert a pytorch tensor argument b to tt
    padded_b = pad_weight(b)
    pyt_a = tt2torch(a)  # untilizes also
    return print_diff_argmax(pyt_a, padded_b, annotation)


def ttP(x, count=4, offset=0, stride=1):
    if type(x) == torch.Tensor:
        t1 = x.reshape(-1)
    else:
        tt_out = x.cpu()
        torch_out = untilize(tt_out.to_torch())
        t1 = torch_out.reshape(-1)
    print("Tensor vals: (", end="")
    for j in range(offset, offset + count * stride, stride):
        print(t1[j].item(), " ", end="")
    print(")")


### Conv related helpers ###
def read_conv_act_into_mm_act_block(
    conv_act,
    act_address_map_index,
    address_map,
    address_map_this_block_size,
    act_block_h,
    act_block_w,
):
    mm_act_block_shape = [1, 1, act_block_h * 32, act_block_w * 32]
    mm_act_block_size = act_block_h * act_block_w * 1024
    mm_act_block = torch.zeros(mm_act_block_size, dtype=torch.bfloat16).float()
    for i in range(0, address_map_this_block_size, 4):
        src_address = address_map[act_address_map_index]
        dst_address = address_map[act_address_map_index + 1]
        read_size = address_map[act_address_map_index + 2]
        pad = address_map[act_address_map_index + 3]
        for s in range(read_size):
            assert dst_address + s < mm_act_block_size
            if pad:
                mm_act_block[dst_address + s] = 0
            else:
                assert src_address + s < len(conv_act)
                mm_act_block[dst_address + s] = conv_act[src_address + s]
        act_address_map_index += 4
    return (mm_act_block.reshape(mm_act_block_shape), act_address_map_index)


def read_conv_weight_into_mm_weight_block(
    conv_weight,
    weight_address_map_index,
    weight_address_map,
    weight_address_map_this_block_size,
    weight_block_h,
    weight_block_w,
):
    mm_weight_block_shape = [1, 1, weight_block_h * 32, weight_block_w * 32]
    mm_weight_block_size = weight_block_h * weight_block_w * 1024
    mm_weight_block = torch.zeros(mm_weight_block_size, dtype=torch.bfloat16).float()
    for i in range(0, weight_address_map_this_block_size, 4):
        src_address = weight_address_map[weight_address_map_index]
        dst_address = weight_address_map[weight_address_map_index + 1]
        read_size = weight_address_map[weight_address_map_index + 2]
        pad = weight_address_map[weight_address_map_index + 3]
        for s in range(read_size):
            assert dst_address + s < mm_weight_block_size
            if pad:
                mm_weight_block[dst_address + s] = 0
            else:
                assert src_address + s < len(conv_weight)
                mm_weight_block[dst_address + s] = conv_weight[src_address + s]
        weight_address_map_index += 4
    return (mm_weight_block.reshape(mm_weight_block_shape), weight_address_map_index)


def blocked_mm_with_conv_act(
    conv_act,
    mm_weight,
    act_address_map,
    weight_address_map,
    num_blocks_act_h,
    num_blocks_act_w,
    num_blocks_weight_w,
    act_block_h,
    act_block_w,
    weight_block_w,
):
    # act refers to conv activation tensor
    # weight refers to conv weight tensor
    mm_output_shape = [
        1,
        1,
        num_blocks_act_h * act_block_h * 32,
        num_blocks_weight_w * weight_block_w * 32,
    ]
    ret = torch.zeros(mm_output_shape, dtype=torch.bfloat16).float()
    mm_output_block_shape = [1, 1, act_block_h * 32, weight_block_w * 32]
    act_address_map_index = 0
    weight_address_map_index = 0
    weight_block_h = act_block_w
    num_groups = act_address_map[act_address_map_index]
    assert num_groups == num_blocks_act_h * num_blocks_act_w * num_blocks_weight_w
    weight_num_groups = act_address_map[weight_address_map_index]
    assert weight_num_groups == num_groups
    act_address_map_index += 1
    weight_address_map_index += 1
    for block_act_h in range(num_blocks_act_h):
        # Reset weight (weight) to the starting tile in this column
        for block_weight_w in range(num_blocks_weight_w):
            output_block = torch.zeros(mm_output_block_shape, dtype=torch.bfloat16).float()
            for block_act_w in range(num_blocks_act_w):
                address_map_this_block_size = act_address_map[act_address_map_index]
                act_address_map_index += 1
                weight_address_map_this_block_size = weight_address_map[weight_address_map_index]
                weight_address_map_index += 1
                (mm_act_block, act_address_map_index) = read_conv_act_into_mm_act_block(
                    conv_act,
                    act_address_map_index,
                    act_address_map,
                    address_map_this_block_size,
                    act_block_h,
                    act_block_w,
                )
                (
                    mm_weight_block,
                    weight_address_map_index,
                ) = read_conv_weight_into_mm_weight_block(
                    mm_weight,
                    weight_address_map_index,
                    weight_address_map,
                    weight_address_map_this_block_size,
                    weight_block_h,
                    weight_block_w,
                )
                # Untilize weight block (this CPU reference does matmul on untilized blocks)
                mm_weight_block = untilize(mm_weight_block)
                for out_h_block in range(act_block_h * 32):
                    for out_w_block in range(weight_block_w * 32):
                        output_block[0][0][out_h_block][out_w_block] += torch.dot(
                            mm_act_block[0, 0, out_h_block, :].reshape(-1),
                            mm_weight_block[0, 0, :, out_w_block].reshape(-1),
                        )
            start_oh = block_act_h * act_block_h * 32
            start_ow = block_weight_w * weight_block_w * 32
            end_oh = start_oh + (act_block_h * 32)
            end_ow = start_ow + (weight_block_w * 32)
            ret[0, 0, start_oh:end_oh, start_ow:end_ow] = output_block

    return ret


def is_conv_supported_on_device(conv_params):
    K, C, R, S, U, V, P_H, P_W, dilation, groups = [conv_params[i] for i in range(10)]

    if K % 32 != 0 or dilation != 1 or groups != 1:
        logger.warning("DOES NOT HAVE SUPPORT FOR Conv with following parameters -")
        logger.warning(
            "K="
            + str(K)
            + " C="
            + str(C)
            + " R="
            + str(R)
            + " S="
            + str(S)
            + " U="
            + str(U)
            + " V="
            + str(V)
            + " PH="
            + str(P_H)
            + " PW="
            + str(P_W)
            + " dilation="
            + str(dilation)
            + " groups="
            + str(groups)
        )
        return False

    return True


def is_x2_harvested(device):
    grid = device.compute_with_storage_grid_size()
    return device.arch() == Arch.WORMHOLE_B0 and (grid.x, grid.y) == (8, 7)


def is_single_chip():
    return ttnn.GetNumAvailableDevices() == 1


def is_quasar():
    ARCH_NAME = ttnn.get_arch_name()
    return "quasar" in ARCH_NAME


def is_blackhole():
    ARCH_NAME = ttnn.get_arch_name()
    return "blackhole" in ARCH_NAME


def is_wormhole_b0():
    ARCH_NAME = ttnn.get_arch_name()
    return "wormhole_b0" in ARCH_NAME


def is_watcher_enabled():
    watcher = os.environ.get("TT_METAL_WATCHER")
    lightweight_asserts = os.environ.get("TT_METAL_LIGHTWEIGHT_KERNEL_ASSERTS")
    return (watcher is not None and watcher != "") or lightweight_asserts == "1"


def is_llk_assert_enabled():
    llk_assert = os.environ.get("TT_METAL_LLK_ASSERTS")
    return llk_assert == "1"


def is_n300():
    return os.environ.get("MESH_DEVICE", "N150") == "N300"


def is_slow_dispatch():
    return os.environ.get("TT_METAL_SLOW_DISPATCH_MODE") == "1"


def ti_skip(condition, reason="Invalid test parameters"):
    import pytest

    return pytest.mark.skipif(condition, reason="Skipping unsupported case: " + reason)


def skip_for_blackhole(reason_str="not a blackhole test"):
    return ti_skip(is_blackhole(), reason=reason_str)


def skip_for_wormhole_b0(reason_str="not a wormhole test"):
    return ti_skip(is_wormhole_b0(), reason=reason_str)


def skip_with_watcher(reason_str="Test is not passing with watcher enabled"):
    return ti_skip(is_watcher_enabled(), reason=reason_str)


def skip_with_llk_assert(reason_str="Test is not passing with LLK asserts enabled"):
    return ti_skip(is_llk_assert_enabled(), reason=reason_str)


def run_for_blackhole(reason_str="only runs for Blackhole"):
    return ti_skip(not is_blackhole(), reason=reason_str)


def run_for_wormhole_b0(reason_str="only runs for Wormhole B0"):
    return ti_skip(not is_wormhole_b0(), reason=reason_str)


def run_for_wormhole_b0_or_blackhole(reason_str="only runs for Wormhole B0 or Blackhole"):
    return ti_skip(not (is_wormhole_b0() or is_blackhole()), reason=reason_str)


def run_for_n_dev(n, reason_str="Test is not meant for this number of devices"):
    return ti_skip(ttnn.get_num_devices() != n, reason=reason_str)


def skip_for_n_dev(n, reason_str="Test is not meant for this number of devices"):
    return ti_skip(ttnn.get_num_devices() == n, reason=reason_str)


def skip_for_n_or_less_dev(n, reason_str="Test is not meant for this number of devices"):
    return ti_skip(ttnn.get_num_devices() <= n, reason=reason_str)


def skip_for_slow_dispatch(reason_str="not working for slow dispatch"):
    return ti_skip(is_slow_dispatch(), reason=reason_str)


def ttl_complex_2_torch_complex(tt_tensor):
    torch_tensor = tt2torch_tensor(tt_tensor)

    # extract real and imag parts of the complex tensor
    real = torch_tensor[:, :, :, : torch_tensor.shape[-1] // 2].to(torch.bfloat16).to(torch.float)
    imag = torch_tensor[:, :, :, torch_tensor.shape[-1] // 2 :].to(torch.bfloat16).to(torch.float)

    # create torch complex tensor
    result = torch.complex(real, imag)
    return result


def pad_and_fold_conv_filters_for_unity_stride(filter_pyt_nchw_tensor, stride_h, stride_w, align_c=4):
    assert stride_h == stride_w
    assert filter_pyt_nchw_tensor.shape[2] == filter_pyt_nchw_tensor.shape[3]
    assert isinstance(align_c, int) and align_c > 0
    # Fold activation for unity stride
    # Pad channel size to align_c. This keeps L1 read addresses aligned; extra channels become
    # zero-valued weights that contribute nothing to the convolution. align_c=4 is the WH/BH default
    # (16B alignment for bf16 gives C a multiple of 4 with a tiled conv reader). Quasar's row-major
    # fold needs align_c=8 (bf16 row-major shard width must be a multiple of 8) so the first conv
    # folds to groups*8 input channels and consumes the aligned output without per-group padding strip.
    C = _nearest_y(filter_pyt_nchw_tensor.shape[1], align_c)
    # Pad filter to nearest stride
    Padded_filter_height = _nearest_y(filter_pyt_nchw_tensor.shape[2], stride_h)
    Padded_filter_width = _nearest_y(filter_pyt_nchw_tensor.shape[3], stride_w)
    filter_pyt_padded = torch.nn.functional.pad(
        filter_pyt_nchw_tensor,
        (
            0,
            Padded_filter_width - filter_pyt_nchw_tensor.shape[3],
            0,
            Padded_filter_height - filter_pyt_nchw_tensor.shape[2],
            0,
            C - filter_pyt_nchw_tensor.shape[1],
        ),
    )
    # Fold filter for unity stride.
    filter_pyt_padded_folded = torch.zeros(
        [
            filter_pyt_padded.shape[0],
            C * stride_h * stride_w,
            (int)(filter_pyt_padded.shape[2] / stride_h),
            (int)(filter_pyt_padded.shape[3] / stride_w),
        ]
    )
    for h in range(0, filter_pyt_padded.shape[2], stride_h):
        for w in range(0, filter_pyt_padded.shape[3], stride_w):
            folded_h = (int)(h / stride_h)
            folded_w = (int)(w / stride_w)
            for i in range(4):
                start_c = i * C
                filter_pyt_padded_folded[:, start_c : start_c + C, folded_h, folded_w] = filter_pyt_padded[
                    :, :, h + (int)(i / stride_w), w + (int)(i % stride_w)
                ]
    return filter_pyt_padded_folded


# produces a tensor where each element in a page is the page number
# this tensor is easy to debug and visualize
def get_debug_tensor(num_pages_width, num_pages_height, dtype, page_width=32, page_height=32):
    torch_tensor = None
    for row_idx in range(0, int(num_pages_height)):
        tile_row = None
        for col_idx in range(0, int(num_pages_width)):
            tile_idx = col_idx + num_pages_width * row_idx
            tile = torch.full((1, 1, page_width, page_height), tile_idx + 1, dtype=dtype)
            if tile_row == None:
                tile_row = tile
            else:
                tile_row = torch.cat((tile_row, tile), 3)
        if torch_tensor == None:
            torch_tensor = tile_row
        else:
            torch_tensor = torch.cat((torch_tensor, tile_row), 2)

    return torch_tensor


# ── transformers 5.x Cache API compatibility ────────────────────────────────
# transformers 5.x removed the legacy Cache API: DynamicCache no longer exposes
# from_legacy_cache / to_legacy_cache / key_cache / value_cache (per-layer KV now
# lives at cache.layers[i].keys/.values). These helpers work on both 4.x and 5.x.
def hf_cache_layer_kv(cache, layer_idx):
    """Return (key, value) tensors for a layer of a transformers Cache.

    Handles the legacy tuple-of-tuples past_key_values, transformers <5 Cache
    (key_cache/value_cache), and transformers >=5 Cache (layers[i].keys/.values).
    """
    if isinstance(cache, (tuple, list)):  # legacy tuple-of-tuples past_key_values
        return cache[layer_idx][0], cache[layer_idx][1]
    if hasattr(cache, "key_cache"):  # transformers < 5.x Cache
        return cache.key_cache[layer_idx], cache.value_cache[layer_idx]
    layer = cache.layers[layer_idx]  # transformers >= 5.x Cache
    return layer.keys, layer.values


def hf_cache_to_legacy(cache):
    """Export a transformers Cache to the legacy tuple-of-(key, value) format."""
    if hasattr(cache, "to_legacy_cache"):  # transformers < 5.x
        return cache.to_legacy_cache()
    return tuple((layer.keys, layer.values) for layer in cache.layers)  # transformers >= 5.x


def hf_dynamic_cache_from_legacy(layer_kvs):
    """Build a transformers DynamicCache from per-layer (key, value) tuples."""
    from transformers import DynamicCache

    layer_kvs = tuple(layer_kvs)
    if hasattr(DynamicCache, "from_legacy_cache"):  # transformers < 5.x
        return DynamicCache.from_legacy_cache(layer_kvs)
    return DynamicCache(layer_kvs)  # transformers >= 5.x


def hf_cache_num_layers(cache):
    """Number of populated layers in a transformers Cache (version-tolerant)."""
    return len(cache.key_cache) if hasattr(cache, "key_cache") else len(cache.layers)


def hf_empty_encoder_decoder_cache():
    """Create an empty transformers EncoderDecoderCache (version-tolerant)."""
    from transformers import DynamicCache, EncoderDecoderCache

    if hasattr(EncoderDecoderCache, "from_legacy_cache"):  # transformers < 5.x
        return EncoderDecoderCache.from_legacy_cache(None)
    return EncoderDecoderCache(DynamicCache(), DynamicCache())  # transformers >= 5.x


def copy_to_buffer(src: "ttnn.Tensor", dst: "ttnn.Tensor", target_dtype) -> None:
    """Convert ``src`` to ``dst``'s layout/dtype/shape/memcfg and write it into
    ``dst``. ``dst``'s device buffer is preserved (no reallocation) so any
    captured trace and the DRAM prefetcher's recorded buffer addresses remain
    valid. The final ``ttnn.to_memory_config`` with ``output_tensor=dst`` both
    reshards to ``dst``'s memory config and copies into ``dst``'s buffer.
    """
    converted = src

    if converted.layout != dst.layout:
        converted = ttnn.to_layout(converted, layout=dst.layout)

    if converted.dtype != target_dtype:
        converted = ttnn.typecast(converted, dtype=target_dtype)

    if tuple(converted.shape) != tuple(dst.shape):
        converted = ttnn.reshape(converted, list(dst.shape))

    ttnn.to_memory_config(converted, dst.memory_config(), output_tensor=dst)