import typing as T
import logging
import scipy
import os

import numpy as np
import math
from keras import ops
from tsgm.types import Tensor as TensorLike

import tsgm


logger = logging.getLogger('utils')
logger.setLevel(logging.DEBUG)


def _ensure_float32_compatibility(kernel_result):
    """Ensure kernel results are compatible with current backend, especially MPS."""
    if os.environ.get("KERAS_BACKEND") == "torch":
        try:
            import torch
            if torch.backends.mps.is_available():
                # Convert to float32 for MPS compatibility
                if isinstance(kernel_result, torch.Tensor):
                    # Already a tensor, just ensure float32 and MPS compatibility
                    result = kernel_result.cpu().numpy().astype(np.float32)
                else:
                    result = np.asarray(kernel_result, dtype=np.float32)
                return ops.convert_to_tensor(result)
        except ImportError:
            pass
    return ops.convert_to_tensor(kernel_result)


def _tensor_to_numpy(tensor):
    """Convert tensor to numpy array safely across backends."""
    if os.environ.get("KERAS_BACKEND") == "torch":
        try:
            import torch
            if isinstance(tensor, torch.Tensor):
                return tensor.detach().cpu().numpy()
        except ImportError:
            pass
    elif hasattr(tensor, 'numpy'):
        return tensor.numpy()
    return np.asarray(tensor)


#  make guassian kernel fit with both pytorch and tensorflow
def exp_quad_kernel(x: TensorLike, y: TensorLike):
    length_scale, feature_ndims = 1.0, 2
    x_expanded = ops.expand_dims(x, axis=-feature_ndims - 1)
    y_expanded = ops.expand_dims(y, axis=-feature_ndims - 2)
    sq_dist = ops.sum((x_expanded - y_expanded) ** 2, axis=(-2, -1))
    return ops.exp(-0.5 * sq_dist / length_scale ** 2)


def MMD(X: tsgm.types.Tensor, Y: tsgm.types.Tensor, kernel: T.Callable = exp_quad_kernel) -> TensorLike:
    XX = _ensure_float32_compatibility(kernel(X, X))
    YY = _ensure_float32_compatibility(kernel(Y, Y))
    XY = _ensure_float32_compatibility(kernel(X, Y))
    return _tensor_to_numpy(ops.mean(XX) + ops.mean(YY) - 2 * ops.mean(XY))


def kernel_median_heuristic(X1: tsgm.types.Tensor, X2: tsgm.types.Tensor) -> float:
    '''
    Median heuristic (Gretton, 2012) for RBF kernel width.
    '''
    # Ensure compatibility with MPS backend
    X1 = _ensure_float32_compatibility(X1)
    X2 = _ensure_float32_compatibility(X2)

    n = X1.shape[0]
    m = X2.shape[0]

    if n * m >= 10 ** 8:
        logger.warning("n * m >= 10^8, consider subsampling for kernel median heuristic")

    X1_squared = X1_squared = ops.transpose(
        ops.tile(ops.reshape(X1 * X1, [1, -1]), (m, 1))
    )
    X2_squared = ops.tile(ops.reshape(X2 * X2, [1, -1]), (n, 1))

    distances = X1_squared + X2_squared - 2 * ops.tensordot(X1, ops.transpose(X2), axes=1)
    distances_np = _tensor_to_numpy(distances)
    assert np.min(distances_np) >= 0

    non_zero_distances = list(filter(lambda x: x != 0, distances_np.reshape([-1])))
    if non_zero_distances:
        median_distance = np.median(non_zero_distances)
    else:
        median_distance = 0

    return _tensor_to_numpy(ops.sqrt(median_distance / 2))  # 2 * sigma**2


def mmd_diff_var(Kyy: tsgm.types.Tensor, Kzz: tsgm.types.Tensor, Kxy: tsgm.types.Tensor, Kxz: tsgm.types.Tensor) -> float:
    '''
    Computes the variance of the difference statistic MMD_{XY} - MMD_{XZ}
    See http://arxiv.org/pdf/1511.04581.pdf Appendix A for more details.
    '''
    # Ensure compatibility with MPS backend
    Kyy = _ensure_float32_compatibility(Kyy)
    Kzz = _ensure_float32_compatibility(Kzz)
    Kxy = _ensure_float32_compatibility(Kxy)
    Kxz = _ensure_float32_compatibility(Kxz)

    m = Kxy.shape[0]
    n = Kyy.shape[0]
    r = Kzz.shape[0]

    Kyy_nd = Kyy - ops.diag(ops.diagonal(Kyy))  # Kyy - diag[Kyy]
    Kzz_nd = Kzz - ops.diag(ops.diagonal(Kzz))  # Kzz - diag[Kzz]

    # Approximations from Eq. 31
    u_yy = ops.sum(Kyy_nd) / (n * (n - 1))
    u_zz = ops.sum(Kzz_nd) / (r * (r - 1))
    u_xy = ops.sum(Kxy) / (m * n)
    u_xz = ops.sum(Kxz) / (m * r)

    Kyy_nd_T = ops.transpose(Kyy_nd)
    Kxy_T = ops.transpose(Kxy)
    Kzz_nd_T = ops.transpose(Kzz_nd)
    Kxz_T = ops.transpose(Kxz)

    # zeta_1 computation, Eq. 30 & 31 in the paper
    term1 = (1 / (n * (n - 1) ** 2)) * ops.sum(Kyy_nd_T @ Kyy_nd) - u_yy ** 2
    term2 = (1 / (n ** 2 * m)) * ops.sum(Kxy_T @ Kxy) - u_xy ** 2
    term3 = (1 / (m ** 2 * n)) * ops.sum(Kxy @ Kxy_T) - u_xy ** 2
    term4 = (1 / (r * (r - 1) ** 2)) * ops.sum(Kzz_nd_T @ Kzz_nd) - u_zz ** 2
    term5 = (1 / (r * m ** 2)) * ops.sum(Kxz @ Kxz_T) - u_xz ** 2
    term6 = (1 / (m * r ** 2)) * ops.sum(Kxz_T @ Kxz) - u_xz ** 2

    term7 = (1 / (m * n * (n - 1))) * ops.sum(Kyy_nd @ Kxy_T) - u_yy * u_xy
    term8 = (1 / (n * m * r)) * ops.sum(Kxy_T @ Kxz) - u_xz * u_xy
    term9 = (1 / (m * r * (r - 1))) * ops.sum(Kzz_nd @ Kxz_T) - u_zz * u_xz

    zeta1 = (term1 + term2 + term3 + term4 + term5 + term6 - 2 * (term7 + term8 + term9))
    zeta2 = (1 / (m * (m - 1))) * ops.sum((Kyy_nd - Kzz_nd - Kxy_T - Kxy + Kxz + Kxz_T) ** 2) - \
        (u_yy - 2 * u_xy - (u_zz - 2 * u_xz)) ** 2

    var_z1 = (4 * (m - 2) / (m * (m - 1))) * zeta1  # Eq (13)
    var_z2 = (2 / (m * (m - 1))) * zeta2  # Eq (13)

    return _tensor_to_numpy(var_z1 + var_z2)


def mmd_3_test(
        X: tsgm.types.Tensor, Y: tsgm.types.Tensor,
        Z: tsgm.types.Tensor, kernel: T.Callable) -> T.Tuple[float, float, float, float]:
    '''
    Relative MMD test --- returns a test statistic for whether Y is closer to X or than Z.
    See http://arxiv.org/pdf/1511.04581.pdf
    '''

    Kxx = _ensure_float32_compatibility(kernel(X, X))
    Kyy = _ensure_float32_compatibility(kernel(Y, Y))
    Kzz = _ensure_float32_compatibility(kernel(Z, Z))
    Kxy = _ensure_float32_compatibility(kernel(X, Y))
    Kxz = _ensure_float32_compatibility(kernel(X, Z))

    Kxx_nd = Kxx - ops.diag(ops.diagonal(Kxx))
    Kyy_nd = Kyy - ops.diag(ops.diagonal(Kyy))
    Kzz_nd = Kzz - ops.diag(ops.diagonal(Kzz))

    m = Kxy.shape[0]
    n = Kyy.shape[0]
    r = Kzz.shape[0]

    u_xx = ops.sum(Kxx_nd) * (1 / (m * (m - 1)))
    u_yy = ops.sum(Kyy_nd) * (1 / (n * (n - 1)))
    u_zz = ops.sum(Kzz_nd) * (1 / (r * (r - 1)))
    u_xy = ops.sum(Kxy) / (m * n)
    u_xz = ops.sum(Kxz) / (m * r)

    t = u_yy - 2 * u_xy - (u_zz - 2 * u_xz)  # test stat
    diff_var = mmd_diff_var(Kyy, Kzz, Kxy, Kxz)
    sqrt_diff_var = math.sqrt(_tensor_to_numpy(diff_var))

    t_np = _tensor_to_numpy(t)
    pvalue = scipy.stats.norm.cdf(-t_np / sqrt_diff_var)
    tstat = t_np / sqrt_diff_var

    mmd_xy = _tensor_to_numpy(u_xx + u_yy - 2 * u_xy)
    mmd_xz = _tensor_to_numpy(u_xx + u_zz - 2 * u_xz)
    return pvalue, tstat, mmd_xy, mmd_xz
