import cupy as cp
import numpy as np 
from scipy.special import rel_entr

# Relative Entropy Kernel
# Ported from SciPy 1.17.1
# https://github.com/scipy/scipy/blob/maintenance/1.17.x/scipy/special/_convex_analysis.pxd
# Fixes overflow issue not addressed in cupyx implementation
# https://github.com/scipy/scipy/issues/20710
_rel_entr_kernel = cp.ElementwiseKernel(
	'T x, T y',
	'T z',
	'''
	// Select the appropriate minimum normal value for the type at compile time.
	// FLT_MIN  ≈ 1.1755e-38  (float32 smallest normal)
	// DBL_MIN  ≈ 2.2251e-308 (float64 smallest normal)
	const T type_min = (sizeof(T) == 4) ? (T)FLT_MIN : (T)DBL_MIN;

	if (isnan(x) || isnan(y)) {
		z = (T)CUDART_NAN;
	} else if (x <= 0 || y <= 0) {
		if (x == 0 && y >= 0) {
			z = 0;
		} else {
			z = (T)CUDART_INF;
		}
	} else {
		T ratio = x / y;
		if (ratio > 0.5 && ratio < 2.0) {
			z = x * log1p((x - y) / y);
		} else if (ratio > type_min && ratio < (T)CUDART_INF) {
			z = x * log(ratio);
		} else {
			z = x * (log(x) - log(y));
		}
	}
	''',
	'rel_entr_robust'
)

def jensenshannondivergence_cpu(p, q, axis=0, keepdims=False, base=None):
	"""
	Compute the Jensen-Shannon divergence between two probability arrays on GPU.
	
	This is a port of scipy.spatial.distance.jensenshannon using a 
	numerically robust rel_entr kernel, except not taking the square root

	https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290
	"""

	p = np.asarray(p)
	q = np.asarray(q)
	
	# Normalize inputs along the specified axis
	p = p / np.sum(p, axis=axis, keepdims=True)
	q = q / np.sum(q, axis=axis, keepdims=True)
	
	# Calculate the midpoint distribution
	m = (p + q) / 2.0
	
	# Calculate relative entropy using the robust CUDA kernel
	# Note: rel_entr results are in nats
	left = rel_entr(p, m)
	right = rel_entr(q, m)
	
	# Sum across the specified axis
	left_sum = np.sum(left, axis=axis, keepdims=keepdims)
	right_sum = np.sum(right, axis=axis, keepdims=keepdims)
	
	# JSDiv = (D_kl(P||M) + D_kl(Q||M)) / 2
	js = (left_sum + right_sum) / 2.0
	
	# Convert units if base is provided (e.g., base=2 for bits)
	if base is not None:
		js /= np.log(base)

	return js / 2.0


def jensenshannondivergence_gpu(p, q, axis=0, keepdims=False, base=None):
	"""
	Compute the Jensen-Shannon divergence between two probability arrays on GPU.
	
	This is a port of scipy.spatial.distance.jensenshannon using a 
	numerically robust rel_entr kernel, except not taking the square root

	https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290
	"""
	
	p = cp.asarray(p)
	q = cp.asarray(q)

	if p.dtype not in (cp.float32, cp.float64) or p.dtype != q.dtype :
		raise TypeError(
			f"Both arrays must be the same dtype and either float32 or float64, "
			f"got p.dtype={p.dtype}, q.dtype={q.dtype}"
		)
	
	# Normalize inputs along the specified axis
	p = p / cp.sum(p, axis=axis, keepdims=True)
	q = q / cp.sum(q, axis=axis, keepdims=True)
	
	# Calculate the midpoint distribution
	m = (p + q) / 2.0
	
	# Calculate relative entropy using the robust CUDA kernel
	# Note: rel_entr results are in nats
	left  = _rel_entr_kernel(p, m)
	right = _rel_entr_kernel(q, m)
	
	# Sum across the specified axis
	left_sum  = cp.sum(left, axis=axis, keepdims=keepdims)
	right_sum = cp.sum(right, axis=axis, keepdims=keepdims)
	
	# JSDiv = (D_kl(P||M) + D_kl(Q||M)) / 2
	js = (left_sum + right_sum) / 2.0
	
	# Convert units if base is provided (e.g., base=2 for bits)
	if base is not None:
		js /= cp.log(base)

	return js / 2.0
	

def jensenshannon_gpu(p, q, axis=0, keepdims=False, base=None):
	"""
	Compute the Jensen-Shannon distance between two probability arrays on GPU.
	
	This is a port of scipy.spatial.distance.jensenshannon using a 
	numerically robust rel_entr kernel.

	https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290
	"""
	js = jensenshannondivergence_gpu( p=p, q=q, axis=axis, keepdims=keepdims, base=base )

	return cp.sqrt(js / 2.0)
	