import math
import time
import sys
import cv2 as cv
import numpy as np
import scipy.linalg as spla
from functools import reduce
from os import system
from PIL import Image

def read_video(limit = -1):
	video = []
	cap = cv.VideoCapture('D:/Videos/test.mkv')
	frame_num = 0
	while cap.isOpened():
		if limit == frame_num:
			break
		ret, frame = cap.read()
		video.append(frame)
		system('cls')
		print("Loading video: {:.2f}%".format(100*frame_num/limit))
		frame_num += 1
	return video

def change_dimension(data, dim):
	'''
	Video data is in form [Frame][Y Pixel][X Pixel][RGB]
	Maximum size should be 10, but smaller is better.
	Eventually, all dimensional changing should be done in C, programmed
	specifically for each dimension.
	'''
	def cubify(data, s):
		c = tuple(np.array(data.shape)//s)
		cubed = np.zeros(tuple(list(c)+[s]*len(c)), dtype=int)
		trim = [slice(0,s)]
		def permute(n, dim=[]):
			if n == len(c):
				slice_ind = [slice(dim[i]*s, (dim[i]+1)*s) for i in range(n)]
				cubed[tuple(dim)] = data[tuple(slice_ind)]
			else:
				for x in range(c[n]):
					permute(n+1, dim+[x])
		permute(0)

		return cubed

	if dim not in [3,4]:
		raise Exception("Not an available target dimension!")

	if dim == 3:
		# Outer Form: [X Block][Y Block][Z Block]
		# Inner Form: [X Pixel][Y Pixel][Frame+RGB] Size: 6
		s = 6
		blocks = np.swapaxes(np.concatenate(data, axis=2), 0, 1)
		blocks = cubify(blocks, s)
	elif dim == 4:
		# Outer Form [X Block][Y Block][Z Block]
		# Inner Form [Psuedo-Y Block][X Pixel][Y Pizel][Frame+RGB] Size: 6
		s = 6
		blocks = change_dimension(data, 3)
		c = list(blocks.shape)[:3]
		c[2] = c[2]//s
		cubed = np.zeros(tuple(c+[s]*4), dtype=int)
		for x in range(c[0]):
			for y in range(c[1]):
				for z in range(c[2]):
					arr = np.concatenate(blocks[x:x+1][0][y:y+1,s*z:s*(z+1)])
					cubed[x,y,z] = arr
		blocks = cubed

	return blocks

def permutations(arr):
	'''
	The input array will carry possible permutations.
	'''
	final = []
	def permute(n, items=[]):
		if n == len(arr):
			final.append(items)
		else:
			for x in arr[n]:
				permute(n+1, items+[x])
	permute(0)
	return final

def pure_polynomial_fit(blocks, dim): # RESIZE TO 2 PI !!!!!!#
	'''
	A polynomial of P(n-1) is generated with an n dimensional block,
	by calculating the a system of n equations, solved through LU
	decomposition.
	'''
	def permute(c, n, dim=[]):
		if n == len(c):
			d = tuple(dim)
			b = blocks[d].flatten()[::-1]
			pb = p_inv.dot(b)
			y = spla.solve_triangular(l, pb, lower=True, check_finite=False)
			x = spla.solve_triangular(u, y, check_finite=False)
			block_polynomials[d] = x
			if np.sum(dim[1:]) == 0:
				system('cls')
				print("Fit Poly Block {:.2f}%".format(100*dim[0]/c[0]))
		else:
			for x in range(c[n]):
				permute(c, n+1, dim+[x])

	degree = blocks.shape[-1] - 1
	deg_perm = permutations(dim*[list(reversed(range(degree+1)))])
	# Scale the coordinate system from 0 to 2pi, for preparation of DFT.
	block_coord_perm = (2*math.pi/(degree+1)) * deg_perm[:]

	mat_size = int(math.pow(degree + 1, dim))
	a = np.zeros((mat_size, mat_size), dtype=np.uint64)
	for i,perm in enumerate(block_coord_perm):
		a[i] = np.prod(np.power(perm, deg_perm), axis=1, dtype=np.uint64)

	p, l, u = spla.lu(a)
	p_inv = np.linalg.inv(p)
	poly_shape = list(blocks.shape[:-dim]) + [mat_size]
	block_polynomials = np.zeros(tuple(poly_shape), dtype=np.float64)
	permute(blocks.shape[:-dim], 0)

	return block_polynomials


def trig_taylor_fit(block):
	'''
	Assume the polynomial will be a n-dimensional taylor polynomial,
	calculate a polynomial regression that will represent the inputs.
	'''
	pass

def nd_dft(poly, dim):
	'''
	This function calculates a N-dimensional Discrete Fourier Transform,
	specifically where f(t) is a polynomial.
	'''
	def get_integral_mat(n, x, zero, imag):
		'''
		Matrix represents T: P -> K, where P,K have basis' respectively:
		P: x^n, x^(n-1), x^(n-2), ..., c, ix^n, ix^(n-1), ix^(n-2), ..., ic
		K: c, 1/k, 1/k^2, ..., 1/k^n, ic, i/k, i/k^2, ..., i/k^n

		This function returns a matrix to calculate the one-dimensional
		integral. If the argument zero is true, the function will return
		a matrix assuming the definite integral is from 0.
		''' 
		mat_rr, mat_ri, mat_ir, mat_ii = [], [], [], []

		def get_mat_part(zero_row_cond, neg_cond):
			mat_part = []
			for row in range(n+1):
				if row % 2 == zero_row_cond:
					mat_part.append([0]*(n+1))
					continue
				else:
					row_vec = []

				for col in range(n+1):
					x_pow = n-col-row
					if x_pow < 0:
						row_vec.append(0)
					elif x_pow == 0 and zero:
						row_vec.append(0)
					else:
						neg = -1 if (row - neg_cond) % 4 == 0 else 1
						const = math.factorial(n-col)/math.factorial(n-col-row)
						row_vec.append(neg*const*math.pow(x, n-col-row))
				mat_part.append(row_vec)
			return np.array(mat_part, dtype=np.float64)

		mat_r = np.concatenate([get_mat_part(0, 3), get_mat_part(1, 2)])
		if imag:
			mat_i = np.concatenate([get_mat_part(1, 0), get_mat_part(0, 3)])
			return np.concatenate([mat_r, mat_i], axis=1)
		else:
			return mat_r

	def get_mono_mats(block):
		rind, iind = np.nonzero(rblock), np.nonzero(iblock)
		rmono_a, imono_a = rblock[rind], iblock[iind]
		rmono_deg = np.subtract(deg, np.array(deg_perm)[rind].transpose())
		imono_deg = np.subtract(2*deg+1, np.array(deg_perm)[iind].transpose())

		mono_count = len(rmono_a) + len(imono_a)
		mono_mats = np.zeros((dim, mono_count, 2*deg+2))
		for x in range(dim):
			mono_deg = np.concatenate([rmono_deg, imono_deg], axis=1)
			mono_mat[x][np.arange(mono_count), mono_deg[x]]

		return mono_mats

	def precompute_integrals(k_max, dim, degree):
		'''
		This function returns all possible combinations of monomial 
		integrals precomputed for all 'k' until k_max excluding k = 0.
		'''
		integral_k = np.zeros((degree, k_max-1), dtype=np.float64)
		mat = get_integral_mat(degree, 2*math.pi, True, False)
		k_col = np.arange(1, k_max).reshape((k_max-1, 1))
		k_pow = np.power(k_col, -np.arange(1, degree+2).astype('float64'))
		k_b = np.concatenate([k_pow, 1j*k_pow], axis=1).transpose()

		mono_integrals = mat.transpose().dot(k_b)
		# Offset by degree for deg_perm to fit basis x^n, x^(n-1) ... 1.
		deg_perm = permutations(dim*[list(range(degree+1))])
		k_perm = magnitude_spiral(dim*[list(range(1, k_max))])
		poly_integrals = []
		for deg in deg_perm:
			k_perm_prod = []
			for k_vec in k_perm:
				k_perm_prod.append(np.prod(mono_integrals[deg, k_vec]))
			poly_integrals.append(k_perm_prod)

		return np.array(poly_integrals)

	def magnitude_spiral(dim, k_max):

		k_perm = permutations(dim*[list(range(-k_max, k_max+1))])
		k_vecs = {}
		for k_vec in k_perm:
			mag = np.linalg.norm(k_vec)
			try:
				k_vecs[mag].append(k_vec)
			except KeyError:
				k_vecs[mag] = [k_vec]
		k_vec_sort = {}
		for mag in list(sorted(k_vecs.keys())):
			k_vec_sort[mag] = sorted(k_vecs[mag], key=lambda x: list(np.abs(x)))
		
		return np.concatenate(list(k_vec_sort.values()))

	deg = int(math.sqrt(poly.shape[-1], 1/dim)) - 1
	
	# In form [Degree Permutation, K permutation]
	poly_integrals = precompute_integrals(int(math.pow(1000, 1/deg)), dim, deg)


	###### FINISH LOOP
	


def main():
	C_DIM = 4

	video = np.array(read_video(50))
	vid_blocks = change_dimension(video, C_DIM)
	block_poly = np.round(pure_polynomial_fit(vid_blocks, C_DIM), 5)

	print(video_blocks[0])
	c_k = []
	#for poly in block_polynomials:


	#img = Image.fromarray(video[1000])
	#img.save('D:/Videos/test.png')




if __name__ == "__main__":
	main()