from __future__ import absolute_import, division, print_function, unicode_literals

from builtins import range
from future.utils import viewitems

from collections import defaultdict
from typing import TYPE_CHECKING

import numpy as np

from .numba import opjit

if TYPE_CHECKING:
	from typing import Iterator, Tuple

RGB_YELLOW = (255, 255, 0)
RGB_WHITE = (255, 255, 255)

def normalized_choice(p_ind, p_prob):  # is this neccessary? I think `np.random.choice` can handle unnormalized probabilities
	return np.random.choice(p_ind, p=p_prob/np.sum(p_prob))

def is_rgb(img):
	# type: (np.ndarray, ) -> bool

	return len(img.shape) >= 1 and img.shape[-1] == 3

#@opjit() np.errstate doesn't work in numba...
def rgb_to_hsi(image):
	# type: (np.ndarray, ) -> np.ndarray

	""" Converts an array [..., channels] of RGB values to HSI color values (H in rad).
		RGB values are assumed to be normalized to (0, 1).
	"""

	if not is_rgb(image):
		raise ValueError("Input needs to be an array of RGB values")

	r = image[...,0]
	g = image[...,1]
	b = image[...,2]

	out = np.zeros_like(image)

	#allequal = (img == img[:, :, 0, np.newaxis]).all(axis=-1)

	with np.errstate(invalid="ignore"): 
		tmp = (2.*r - g - b) / 2. / np.sqrt((r-g)**2 + (r-b)*(g-b)) # if r==g==b then 0/0

		theta = np.arccos(np.clip(tmp, -1., +1.))

		out[...,0] = np.where(b <= g, theta, 2*np.pi - theta) # H
		out[...,2] = np.sum(image, axis=-1) / 3. # I
		out[...,1] = 1 - np.amin(image, axis=-1) / out[...,2] # S if r==g==b==0 then 0/0

	np.nan_to_num(out[...,0:2], copy=False)

	return out

#@opjit() np.dot with more than 2 dimensions is not supported...
def rgb_to_ycbcr(image):
	# type: (np.ndarray, ) -> np.ndarray

	""" Converts an array [..., channels] of RGB values to Digital Y'CbCr (0-255).
		RGB values are assumed to be normalized to (0, 1).
		Don't forget to cast to uint8 for pillow.
	"""


	"""  from RGB (0-1).
	"""

	if not is_rgb(image):
		raise ValueError("Input needs to be an array of RGB values")

	m = np.array([
		[+065.481, +128.553, +024.966],
		[-037.797, -074.203, +112.000],
		[+112.000, -093.786, -018.214],
	])
	a = np.array([16, 128, 128])

	return np.dot(image, m.T) + a

def random_triangular_matrix(size, lower=True):
	# type: (int, bool) -> np.ndarray

	""" Returns a triangular matrix with random value between 0 and 1 uniformly.
	"""

	a = np.random.uniform(0, 1, (size, size))
	if lower:
		ind = np.triu_indices(5, 1)
	else:
		ind = np.tril_indices(5, 1)
	a[ind] = 0

	return a

def issquare(A):
	# type: (np.ndarray, ) -> bool

	if len(A.shape) != 2:
		return False

	if A.shape[0] != A.shape[1]:
		return False

	return True

def batch_vTAv(A, v):
	# type: (float[B, X, X], float[B, X]) -> float[B]

	""" Performs batched calculation of `v^T A v` transform.
		Special case of bilinear form `x^T A y`
	"""

	""" Faster than
		Av = np.matmul(A, v[...,:,None]) # [B, X, 1]
		return np.matmul(v[...,None,:], Av).squeeze((-2, -1)) # [B]
	"""

	return np.einsum("...k,...kl,...l->...", v, A, v)

# was: inner1d
def batch_inner(a, b):
	# type: (np.ndarray, np.ndarray) -> np.ndarray

	""" Performs a batched inner product over the last dimension.
		Replacement for deprecated `from numpy.core.umath_tests import inner1d`.
		Shapes: (A, B), (A, B) -> (A, )
	"""

	if a.shape != b.shape:
		raise ValueError("All dimensions have to be equal")

	return np.einsum("...i,...i->...", a, b) # faster than np.sum(a * b, axis=-1)

def batch_outer(a, b):
	# type: (np.ndarray, np.ndarry) -> np.ndarray

	""" Performs a batched outer product over the last dimension.
		Shapes: (A, B), (A, C) -> (A, B, C)
	"""

	if a.shape[:-1] != b.shape[:-1]:
		raise ValueError("All except the last dimension have to be equal")

	return np.einsum("...i,...j->...ij", a, b) # slightly faster than np.multiply(a[...,:,None], b[...,None,:])

def batchtopk(probs, k=None, axis=-1, reverse=False):
	# type: (np.ndarray, np.ndarray, int) -> np.ndarray

	""" `probs` values ndarray
		`k` take the smallest `k` elements, if `reverse` is False
			and the largest `k` if `reverse` is True
		`axis` sorting and selection axis.
	"""

	assert k is None or k > 0, "k must be larger than zero. Use None to chose all elements."
	assert axis == -1, "Only last axis supported atm"
	assert len(probs.shape) > 1

	if reverse:
		sign = -1
	else:
		sign = 1

	indices = np.argsort(sign * probs, axis=-1) # use argpartition?
	probs = np.take_along_axis(probs, indices[...,:k], axis=-1)

	return indices, probs

def logtrace(m):
	""" Calcuates the sum of the logs of the diagonal elements (batchwise if neccessary)
		m: [..., x, x]
	"""

	return np.sum(np.log(np.diagonal(m, axis1=-2, axis2=-1)), axis=-1)

def shiftedexp(pvals):
	# type: (np.ndarray, ) -> np.ndarray

	""" Prevents overflow. Can be used if probabilities are normalized again later.
	"""

	return np.exp(pvals - np.max(pvals))

class Sampler(object):

	def __init__(self, cdf):
		self.cdf = cdf
		self.psum = cdf[-1]

	def __call__(self):
		rand = np.random.uniform(0, self.psum)
		return np.searchsorted(self.cdf, rand, side="right")

def sample_probabilities(pvals):
	# type: (np.ndarray, ) -> Callable[[], int]

	""" Sample from list of probabilities `pvals` with replacement.
		The probabilities don't need to be normalized.
	"""

	return Sampler(np.cumsum(pvals))

class UnboundedSparseMatrix(object):

	def __init__(self, dtype=float):
		# type: (type, ) -> None

		self.dtype = dtype
		self.zero = self.dtype(0)
		self.m = dict()
		self.cols = 0
		self.rows = 0

	def __getitem__(self, slice):
		# type: (tuple, ) -> T

		return self.m.get(slice, self.zero)

	def __setitem__(self, slice, value):
		# type: (tuple, T) -> None

		c, r = slice
		self.cols = max(self.cols, c+1)
		self.rows = max(self.rows, r+1)
		self.m[slice] = value

	def todense(self):
		# type: () -> np.ndarray

		ret = np.zeros((self.cols, self.rows), dtype=self.dtype)

		for slice, value in viewitems(self.m):
			ret[slice] = value

		return ret

def normalize(pvals):
	# type: (np.ndarray, ) -> np.ndarray

	return pvals / np.sum(pvals)

def categorical(pvals):
	# type: (np.ndarray, ) -> int

	""" Requires normalized inputs: sum(pvals) ~= 1 """

	return np.argmax(np.random.multinomial(1, pvals))

def inf_matrix_power(pm):
	w, v = np.linalg.eig(pm) # scipy.linalg.eig would probably by faster as it can return the left and right eigen vectors

	if not np.isclose(w[0], 1.):
		raise ValueError("The first eigenvalue is not none. Is this a right stochastic matrix?")

	vi = np.linalg.inv(v)
	d = np.zeros(m.shape[0], dtype=np.float)
	d[0] = 1.

	return np.matmul(v, np.matmul(np.diag(d), vi))

def decompress(selectors, data, default):
	# type: (np.ndarray[bool], np.ndarray[T], T) -> None

	""" Same result as:
		`np.array(list(genutility.iter.decompress(selectors, iter(data), default)))`
		but faster
	"""

	out = np.full(len(selectors), default)
	out[selectors] = data
	return out

def unblock(arr, n1, n2, axis1=-1, axis2=-2, blocksize=False):

	""" Inverse of np.block.
		Set axis to (-2, -1) to modify the order of the result.
	"""

	""" test (stackoverflow): Ok, so considering I have N block matrices with bm x bn dimension and want to stack them in a m x n matrix, provided N = m x n, I would then have x.reshape(m,n,bm,bn).swapaxes(1,2).reshape(bm*m,-1)
	"""

	s = np.array(arr.shape)
	assert s[axis1] % n1 == 0 and s[axis2] % n2 == 0, "{}x{} does not divide by {}x{}".format(s[axis1], s[axis2], n1, n2)

	if blocksize:
		n1 = s[axis1] // n1
		n2 = s[axis2] // n2

	arr = np.array(np.split(arr, n1, axis1))
	arr = np.array(np.split(arr, n2, axis2))

	inv_blocksize = n1 * n2
	total = s[axis1] * s[axis2]
	s[axis2] = inv_blocksize
	s[axis1] = total // inv_blocksize

	return np.reshape(arr, s)

def remove_color(img, ratio, neutral_color=RGB_WHITE):
	# type: (np.ndarray, float, Tuple[int, int, int]) -> np.ndarray

	""" Replace colored pixels with a `neutral_color`. The `ratio` defines the 'colorfulness' above
		which level the pixel should be replace.
		I.e. if the `ratio` is 1 nothing will be replaced,
		if `ratio` is 0 only strict greys are kept unmodified.
	"""

	channels = img.shape[-1]
	assert channels == 3, "Not a 3 channel color image"

	norm = np.std(np.array(RGB_YELLOW)) # this is the same for all pure colors

	sd = np.std(img, axis=-1)
	img[sd > ratio*norm] = neutral_color

#@opjit(cache=False) doesn't help a lot, because it's a generator
def sliding_window_2d(image, window_size, step_size=(1, 1)):
	# type: (np.ndarray, Tuple[int, int], Tuple[int, int]) -> Iterator[np.ndarray]

	win_x, win_y = window_size
	step_x, step_y = step_size
	height, width = image.shape[0:2]

	for y in range(0, height - win_y + 1, step_y):
		for x in range(0, width - win_x + 1, step_x):
			yield image[y:y + win_y, x:x + win_x, ...]

def histogram_correlation(hist1, hist2):
	# type: (np.ndarray, np.ndarray) -> np.ndarray

	""" Input shape of `hist1` and `hist2`: [batch..., levels]. The correlation is calculated over `levels`
		and then batched over the remaining dimensions.
	"""

	assert len(hist1.shape) >= 1 and hist1.shape == hist2.shape

	h1norm = (hist1.T - np.mean(hist1, axis=-1)).T
	h2norm = (hist2.T - np.mean(hist2, axis=-1)).T

	num = np.sum(h1norm * h2norm, axis=-1)
	denom = np.sqrt(np.sum(h1norm**2, axis=-1)*np.sum(h2norm**2, axis=-1))
	return num / denom

if __name__ == "__main__":
	import timeit
	from .numpy import random_triangular_matrix

	image = np.random.randint(0, 255, (1000, 1000))

	list(sliding_window_2d(image, (100, 100), (1, 1))) # warmup
	print(min(timeit.repeat('list(sliding_window_2d(image, (100, 100), (1, 1)))', number=10, repeat=5, globals=globals())))
