import pdb
from tqdm import tqdm
import types

from itertools import combinations, permutations
from scipy.spatial.distance import cosine

import numpy as np
from numpy.linalg import norm
from numpy.random import shuffle, choice
from numpy.testing import assert_array_equal
import pandas as pd

from matplotlib import pyplot as plt
import matplotlib.gridspec as gridspec

from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split

from keras import backend as K
from keras.utils import np_utils
from keras.models import Sequential
from keras.layers import Dense, Activation
from keras.datasets import mnist

import sys
from PyQt5.QtWidgets import QApplication
app = QApplication(sys.argv)
screen = app.screens()[0]
my_dpi = screen.physicalDotsPerInch()
app.quit()

my_dpi = 100

###################################
## 2d dataset generation methods ##
###################################

### Figure generation methods ###

def generate_a_drawing(figsize, U, V, noise=0.0):
	fig = plt.figure(figsize=(figsize,figsize))
	ax = plt.subplot(111)
	plt.axis('Off')
	ax.set_xlim(0,figsize)
	ax.set_ylim(0,figsize)
	ax.fill(U, V, "k")
	fig.canvas.draw()
	imdata = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)[::3].astype(np.float32)
	imdata = 255 - imdata + noise * np.random.random(imdata.size)
	imdata[imdata < 0] = 0
	imdata[imdata > 255] = 255
	plt.close(fig)
	return imdata

def generate_a_rectangle(width_in_pixels= None, noise=0.0, free_location=False):
	figsize = width_in_pixels / my_dpi
	U = np.zeros(4)
	V = np.zeros(4)
	if free_location:
		corners = np.random.random(4) * figsize
		top = max(corners[0], corners[1])
		bottom = min(corners[0], corners[1])
		left = min(corners[2], corners[3])
		right = max(corners[2], corners[3])
		
		theta = np.pi*np.random.random()
		mid_len = (right+left)/2
		mid_hei = (top+bottom)/2
		length = right - left 
		height = top - bottom
		
		U[0] = mid_len - (length/2) - np.sqrt(length*length + height*height)*np.sin(theta)/2
		U[1] = mid_len + (length/2) - np.sqrt(length*length + height*height)*np.sin(theta)/2
		U[2] = mid_len + (length/2) + np.sqrt(length*length + height*height)*np.sin(theta)/2
		U[3] = mid_len - (length/2) + np.sqrt(length*length + height*height)*np.sin(theta)/2

		V[0] = mid_hei + (height/2) - np.sqrt(length*length + height*height)*np.sin(theta)/2
		V[1] = mid_hei + (height/2) + np.sqrt(length*length + height*height)*np.sin(theta)/2
		V[2] = mid_hei - (height/2) + np.sqrt(length*length + height*height)*np.sin(theta)/2
		V[3] = mid_hei - (height/2) - np.sqrt(length*length + height*height)*np.sin(theta)/2

	else:
		side = (0.3 + 0.7 * np.random.random()) * figsize
		top = figsize/2 + side/2
		bottom = figsize/2 - side/2
		left = bottom
		right = top
		U[0] = U[1] = top
		U[2] = U[3] = bottom
		V[0] = V[3] = left
		V[1] = V[2] = right
	return generate_a_drawing(figsize, U, V, noise)

def generate_a_disk(width_in_pixels= None, noise=0.0, free_location=False):
	figsize = width_in_pixels / my_dpi
	if free_location:
		center = (0.1 + 0.8*np.random.random(2)) * figsize/2
	else:
		center = (figsize/2, figsize/2)
	radius = (0.2 + 0.5 * np.random.random()) * figsize/2
	N = 50
	U = np.zeros(N)
	V = np.zeros(N)
	i = 0
	for t in np.linspace(0, 2*np.pi, N):
		U[i] = center[0] + np.cos(t) * radius
		V[i] = center[1] + np.sin(t) * radius
		i = i + 1
	return generate_a_drawing(figsize, U, V, noise)

def generate_a_triangle(width_in_pixels= None, noise=0.0, free_location=False):
	figsize = width_in_pixels / my_dpi
	if free_location:
		U = np.random.random(3)*figsize/2
		V = np.random.random(3)*figsize/2
	else:
		size = (0.3 + 0.7 * np.random.random())*figsize/2
		middle = figsize/2
		U = (middle, middle+size, middle-size)
		V = (middle+size, middle-size, middle-size)
	#imdata = generate_a_drawing(figsize, U, V, noise)
	#return [imdata, [U[0], V[0], U[1], V[1], U[2], V[2]]]
	return generate_a_drawing(figsize, U, V, noise)


#im = generate_a_rectangle(10, True)
#plt.imshow(im.reshape(72,72), cmap='gray')

#im = generate_a_disk(10)
#plt.imshow(im.reshape(72,72), cmap='gray')

#[im, v] = generate_a_triangle(20, False)
#plt.imshow(im.reshape(72,72), cmap='gray')


#######################################
### Generate classification dataset ###
#######################################

def generate_nclasses_dataset(nb_samples, nb_classes, width_in_pixels, noise=0.0, free_location=False, verbose=False, nb_test_samples=None):
	"""
	Generates a dataset that contains nb_classes classes whose template are taken from random patterns in the input space
	Examples are generate from the templates for each class by adding Gaussian noise

	Templates are generated by first choosing a random number of vertices, and then creating shapes out of the edges between those verticies
	"""
	assert nb_classes > 0
	class_templates = []
	for i in range(nb_classes):
		category = np.random.randint(3)
		if category == 0:
			print('Class {0:d} is a rectangle'.format(i))
			class_templates.append(
				generate_a_rectangle(width_in_pixels, noise=noise, free_location=free_location).reshape((width_in_pixels,width_in_pixels))
				)
		elif category == 1:
			print('Class {0:d} is a disk'.format(i))
			class_templates.append(
				generate_a_disk(width_in_pixels, noise=noise, free_location=free_location).reshape((width_in_pixels,width_in_pixels))
				)
		else:
			print('Class {0:d} is a triangle'.format(i))
			class_templates.append(
				generate_a_triangle(width_in_pixels, noise=noise, free_location=free_location).reshape((width_in_pixels,width_in_pixels))
				)
	
	im_size = width_in_pixels*width_in_pixels

	# Draw nb_samples examples from each class:
	X = np.zeros([nb_samples, width_in_pixels, width_in_pixels])
	y = np.zeros(nb_samples, dtype=int)
	for i in range(nb_samples):
		y[i] = np.random.randint(nb_classes)
		X[i] = (class_templates[y[i]] + noise) / (255 + 2 * noise)

	if nb_test_samples is None:
		return (X, y)
	else:
		assert (nb_test_samples<nb_samples)
		test_rate = float(nb_test_samples)/nb_samples
		X_train, X_test, y_train, y_test = train_test_split(
			X, y, test_size=test_rate, random_state=42
			)
		return (X_train, y_train), (X_test, y_test)


def generate_dataset_classification(nb_samples, width_in_pixels, noise=0.0, free_location=False, verbose=False, nb_test_samples=None):
	# Getting im_size:
	im_size = width_in_pixels*width_in_pixels
	X = np.zeros([nb_samples, width_in_pixels, width_in_pixels])
	y = np.zeros(nb_samples)
	print('Creating {} samples...'.format(nb_samples))
	for i in range(nb_samples):
		if verbose and i % 10 == 0:
			print(i)
		category = np.random.randint(3)
		if category == 0:
			X[i] = generate_a_rectangle(width_in_pixels, noise, free_location).reshape((width_in_pixels,width_in_pixels))
		elif category == 1: 
			X[i] = generate_a_disk(width_in_pixels, noise, free_location).reshape((width_in_pixels,width_in_pixels))
		else:
			[X[i], V] = generate_a_triangle(width_in_pixels, noise, free_location).reshape((width_in_pixels,width_in_pixels))
		y[i] = category
	X = (X + noise) / (255 + 2 * noise)
	print('...done')

	if nb_test_samples is None:
		return (X, y)
	else:
		assert (nb_test_samples<nb_samples)
		test_rate = float(nb_test_samples)/nb_samples
		X_train, X_test, y_train, y_test = train_test_split(
			X, y, test_size=test_rate, random_state=42
			)
		return (X_train, y_train), (X_test, y_test)


def generate_test_set_classification(free_location=True):
	np.random.seed(42)
	[X_test, Y_test] = generate_dataset_classification(300, 20, free_location)
	Y_test = np_utils.to_categorical(Y_test, 3)
	return [X_test, Y_test]


###################################
### Generate regression dataset ###
###################################

def generate_dataset_regression(nb_samples, width_in_pixels, noise=0.0, verbose=False, nb_test_samples=None):
	# Getting im_size:
	im_size = width_in_pixels*width_in_pixels
	X = np.zeros([nb_samples,im_size])
	Y = np.zeros([nb_samples, 6])
	print('Creating {} samples...'.format(nb_samples))
	for i in range(nb_samples):
		if verbose and i % 10 == 0:
			print(i)
		[X[i], Y[i]] = generate_a_triangle(noise, True)
	X = (X + noise) / (255 + 2 * noise)
	print('...done')

	if nb_test_samples is None:
		return (X, y)
	else:
		assert (nb_test_samples<nb_samples)
		test_rate = float(nb_test_samples)/nb_samples
		X_train, X_test, y_train, y_test = train_test_split(
			X, y, test_size=test_rate, random_state=42
			)
		return (X_train, y_train), (X_test, y_test)

import matplotlib.patches as patches

def visualize_prediction(x, y):
	fig, ax = plt.subplots(figsize=(5, 5))
	I = x.reshape((72,72))
	ax.imshow(I, extent=[-0.15,1.15,-0.15,1.15],cmap='gray')
	ax.set_xlim([0,1])
	ax.set_ylim([0,1])

	xy = y.reshape(3,2)
	tri = patches.Polygon(xy, closed=True, fill = False, edgecolor = 'r', linewidth = 5, alpha = 0.5)
	ax.add_patch(tri)

	plt.show()

def generate_test_set_regression():
	np.random.seed(42)
	[X_test, Y_test] = generate_dataset_regression(300, 20)
	return [X_test, Y_test]



#########################
## Image dataset class ##
#########################

class ImageDataset:
	def __init__(self, x_train, y_train, x_test, y_test, x_gnr=None, y_gnr=None, filt_labels=None, spl=None):
		"""
		filt_labels is the list of labels to keep, namely useful to work on balanced dichotomies
		"""
		self.train = {}
		self.test = {}

		self.train['x'] = x_train
		self.train['y'] = y_train
		self.test['x'] = x_test
		self.test['y'] = y_test

		if filt_labels is not None:
			in_filt_train = np.isin(y_train, filt_labels)
			x_filt_train = np.where(in_filt_train)[0]
			self.train['x'] = self.train['x'].take(x_filt_train, axis=0)
			self.train['y'] = self.train['y'][in_filt_train]
			in_filt_test = np.isin(y_test, filt_labels)
			x_filt_test = np.where(in_filt_test)[0]
			self.test['x'] = self.test['x'].take(x_filt_test, axis=0)
			self.test['y'] = self.test['y'][in_filt_test]

		self.train['dichs'] = {}
		self.test['dichs'] = {}

		n_train = self.train['x'].shape[0]
		axes_dim = self.train['x'].shape[1:]
		n_test = self.test['x'].shape[0]
		axes_dim_test = self.test['x'].shape[1:]
		assert_array_equal(axes_dim, axes_dim_test)
		self.axes_dim = axes_dim
		self.n_axes = len(self.axes_dim)
		self.tot_dim = np.prod(self.axes_dim)

		if x_gnr is not None:
			self.gnr = {}
			self.gnr['x'] = x_gnr
			self.gnr['y'] = y_gnr
			self.gnr['dichs'] = {}

		if spl is not None:
			train_spl_ids = choice(n_train, int(spl*n_train), replace=False)
			self.train['x'] = self.train['x'][train_spl_ids,:]
			self.train['y'] = self.train['y'][train_spl_ids]
			n_train = int(spl*n_train)

		if x_gnr is not None:
			self.labels = np.sort(np.unique(np.hstack((self.train['y'],self.test['y'],self.gnr['y']))))
		else:
			self.labels = np.sort(np.unique(np.hstack((self.train['y'],self.test['y']))))
		self.n_labels = len(self.labels)
		self.label_map = {lbl: i for i,lbl in enumerate(self.labels)}
		self.n_train = n_train
		self.n_test = n_test

		self.train['x'] = self.train['x'].reshape((self.n_train, self.tot_dim))
		self.test['x'] = self.test['x'].reshape((self.n_test, self.tot_dim))
		if x_gnr is not None:
			self.n_gnr = x_gnr.shape[0]
			self.gnr['x'] = self.gnr['x'].reshape((self.n_gnr, self.tot_dim))

		self.train['y_ohe'] = self.make_ohe(self.train['y'])
		self.test['y_ohe'] = self.make_ohe(self.test['y'])
		if x_gnr is not None:
			self.gnr['y_ohe'] = self.make_ohe(self.gnr['y'])

		spl_size = min(75, self.n_test)
		spl_ids = np.arange(self.n_test)
		shuffle(spl_ids)
		spl_ids = spl_ids[:spl_size]

		self.spl = {
			'x': self.test['x'][spl_ids],
			'y': self.test['y'][spl_ids]
		}
	

	def make_ohe(self, y):
		ohe = np.zeros((len(y), self.n_labels))
		ids = [self.label_map[lbl] for lbl in y]    
		ohe[np.arange(len(y)),ids] = 1
		return ohe

	def generate_gnr_set(self, train_labels, normalize=True):
		gnr_labels = np.setdiff1d(self.labels, train_labels)
		in_filt_train = np.isin(self.train['y'], train_labels)
		x_filt_train = np.where(in_filt_train)[0]
		x_train = self.train['x'].take(x_filt_train, axis=0)
		if normalize:
			x_train = normalize_array(x_train)

		dim_list = [d for d in self.axes_dim]
		dtuple_train = dim_list.copy()
		dtuple_train.insert(0, x_train.shape[0])
		dtuple_train = tuple(dtuple_train)
		x_train = x_train.reshape(dtuple_train)
		y_train = self.train['y'][in_filt_train]
		
		in_filt_test = np.isin(self.test['y'], train_labels)
		x_filt_test = np.where(in_filt_test)[0]
		x_test = self.test['x'].take(x_filt_test, axis=0)
		if normalize:
			x_test = normalize_array(x_test)

		dtuple_test = dim_list.copy()
		dtuple_test.insert(0, x_test.shape[0])
		dtuple_test = tuple(dtuple_test)
		x_test = x_test.reshape(dtuple_test)
		y_test = self.test['y'][in_filt_test]
		
		x_filt_gnrtrain = np.where(~in_filt_train)[0]
		x_filt_gnrtest = np.where(~in_filt_test)[0]
		x_gnr = self.train['x'].take(x_filt_gnrtrain, axis=0)
		x_gnr = np.vstack((x_gnr, self.test['x'].take(x_filt_gnrtest, axis=0)))
		if normalize:
			x_gnr = normalize_array(x_gnr)
			
		dtuple_gnr = dim_list.copy()
		dtuple_gnr.insert(0, x_gnr.shape[0])
		dtuple_gnr = tuple(dtuple_gnr)
		x_gnr = x_gnr.reshape(dtuple_gnr)
		y_gnr = self.train['y'][~in_filt_train]
		y_gnr = np.hstack((y_gnr, self.test['y'][~in_filt_test]))
		
		return ImageDataset(x_train, y_train, x_test, y_test, x_gnr, y_gnr)

	def build_dichLabels(self, dich_set, dich_name):
		"""
		Example: dich_set = [range(0,6), range(6,12), range(12,18)]
		The labels built will range from 0 to 2, with the following mapping:
		y in range(0,6) => output = (1,0,0)
		y in range(6,12) => output = (0,1,0)
		y in range(12,18) => output = (0,0,1)
		"""
		y_train_in_set = []
		for lbl_set in dich_set[:-1]:
			y_train_in_set.append(np.isin(self.train['y'], list(lbl_set)))

		y_train_in_set.append(np.logical_not(np.any(y_train_in_set, 0)))

		self.train['dichs'][dich_name] = np.transpose(
			np.vstack(
				y_train_in_set
			)
		)

		y_test_in_set = []
		for lbl_set in dich_set[:-1]:
			y_test_in_set.append(np.isin(self.test['y'], list(lbl_set)))

		y_test_in_set.append(np.logical_not(np.any(y_test_in_set, 0)))

		self.test['dichs'][dich_name] = np.transpose(
			np.vstack(
				y_test_in_set
			)
		)
		
		if hasattr(self, 'gnr'):
			y_gnr_in_set = []
			for lbl_set in dich_set[:-1]:
				y_gnr_in_set.append(np.isin(self.gnr['y'], list(lbl_set)))

			y_gnr_in_set.append(np.logical_not(np.any(y_gnr_in_set, 0)))

			self.gnr['dichs'][dich_name] = np.transpose(
				np.vstack(
					y_gnr_in_set
				)
			)

	def build_catLabels(self, dich_set, dich_name):
		"""
		Example: dich_set = [range(0,6), range(6,12), range(12,18)]
		The labels built will range from 0 to 2, with the following mapping:
		y in range(0,6) => output = 0
		y in range(6,12) => output = 1
		y in range(12,18) => output = 2
		"""
		y_train_in_set = []
		for lbl_val, lbl_set in enumerate(dich_set):
			y_train_in_set.append(lbl_val*np.isin(self.train['y'], lbl_set))

		self.train['dichs'][dich_name] = np.transpose(
			np.sum(
				y_train_in_set, axis=0
			)
		).reshape((self.n_train, 1))

		y_test_in_set = []
		for lbl_val, lbl_set in enumerate(dich_set):
			y_test_in_set.append(lbl_val*np.isin(self.test['y'], lbl_set))

		y_test_in_set.append(np.logical_not(np.any(y_test_in_set, 0)))

		self.test['dichs'][dich_name] = np.transpose(
			np.sum(
				y_test_in_set, axis=0
			)
		).reshape((self.n_test, 1))
		
		if hasattr(self, 'gnr'):
			y_gnr_in_set = []
			for lbl_val, lbl_set in enumerate(dich_set):
				y_gnr_in_set.append(lbl_val*np.isin(self.gnr['y'], lbl_set))

			y_gnr_in_set.append(np.logical_not(np.any(y_gnr_in_set, 0)))

			self.gnr['dichs'][dich_name] = np.transpose(
				np.sum(
					y_gnr_in_set, axis=0
				)
			).reshape((self.n_gnr, 1))

	def hstack_dichs(self, dich1_name, dich2_name):
		assert dich1_name in self.train['dichs'].keys()
		assert dich2_name in self.train['dichs'].keys()
		self.train['dichs'][dich1_name+'_hstack_'+dich2_name] = np.hstack(
			(self.train['dichs'][dich1_name],
			 self.train['dichs'][dich2_name])
		)

		assert dich1_name in self.test['dichs'].keys()
		assert dich2_name in self.test['dichs'].keys()
		self.test['dichs'][dich1_name+'_hstack_'+dich2_name] = np.hstack(
			(self.test['dichs'][dich1_name],
			 self.test['dichs'][dich2_name])
		)

		if hasattr(self, 'gnr'):
			assert dich1_name in self.gnr['dichs'].keys()
			assert dich2_name in self.gnr['dichs'].keys()
			self.gnr['dichs'][dich1_name+'_hstack_'+dich2_name] = np.hstack(
			(self.gnr['dichs'][dich1_name],
			 self.gnr['dichs'][dich2_name])
		)

	def compstack_dichs(self, dich1_name, dich2_name):
		assert dich1_name in self.train['dichs'].keys()
		assert dich2_name in self.train['dichs'].keys()
		self.train['dichs'][dich1_name+'_compstack_'+dich2_name] = np.transpose(
			np.vstack(
				(self.train['dichs'][dich1_name][:,0],
				 self.train['dichs'][dich2_name][:,0])
			)
		)

		assert dich1_name in self.test['dichs'].keys()
		assert dich2_name in self.test['dichs'].keys()
		self.test['dichs'][dich1_name+'_compstack_'+dich2_name] = np.transpose(
			np.vstack(
				(self.test['dichs'][dich1_name][:,0],
				 self.test['dichs'][dich2_name][:,0])
			)
		)

		if hasattr(self, 'gnr'):
			assert dich1_name in self.gnr['dichs'].keys()
			assert dich2_name in self.gnr['dichs'].keys()
			self.gnr['dichs'][dich1_name+'_compstack_'+dich2_name] = np.transpose(
				np.vstack(
					(self.gnr['dichs'][dich1_name][:,0],
					 self.gnr['dichs'][dich2_name][:,0])
				)
			)

def normalize_array(ar):
	"""
	Each row must be an entry, each column must be a dimension
	"""
	ar_mean = np.mean(ar, axis=0)
	nar = ar - ar_mean
	ar_min = np.min(nar, axis=0)
	ar_max = np.max(nar, axis=0)
	nar = np.divide(2*nar, ar_max-ar_min, out=np.zeros_like(nar), where=(ar_max!=ar_min)) - np.divide(ar_max+ar_min, ar_max-ar_min, out=np.zeros_like(nar), where=(ar_max!=ar_min))
	return nar