#This file contains functions to find information on the complexity of the methods.
import numpy as np
import scipy as sp
import R3
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import time
from scipy.optimize import root
from skimage import measure

#Plot the time of an approximation and a marching cubes sampling as we increase the amount of input points:
#Parameters: f, the domain, weighing function theta, the range of amount of equidistant points per axis and the order.
def time_input(f, domain, theta, sizes, order):
	X, Y, Z = domain
	approx_times = [None] * len(sizes)
	sample_times = [None] * len(sizes)
	inputs = [None] * len(sizes)
	i = 0
	for s in sizes:
		test = np.mgrid[-4:4:s*1j, -4:4:s*1j, -4:4:s*1j].reshape(3,-1).T #Grid size can be changed here if needed.
		inputs[i] = len(test)
		start = time.time()
		#Get the aproximation (comment/uncomment to choose the method).
		approx = R3.least_squares_R3(test, f(test), order) 
		#approx = R3.weighted_least_squares_R3(test, f(test), order, theta)
		#approx = R3.moving_least_squares_R3(test, f(test), order, theta)
		approx_times[i] = time.time() - start
			
		#Take samples:
		start = time.time()
		vertices, simplices,normals, values = measure.marching_cubes_lewiner(approx(X,Y,Z), -0.3)
		sample_times[i] = time.time() - start
		i+=1

	total_times = [None] * len(sizes)
	for i in range(len(sizes)):
		total_times[i] = approx_times[i] + sample_times[i]
		
		
	plt.plot(inputs, approx_times, color='b', label="Approximation time")
	plt.scatter(inputs, approx_times, color='b')
	plt.plot(inputs, sample_times, color='r', label="Sample time")
	plt.scatter(inputs, sample_times, color='r')
	plt.plot(inputs, total_times, color='#000000', label="Total time")
	plt.scatter(inputs, total_times, color='#000000')
	plt.xlabel('Amount of input points')
	plt.ylabel('Time')
	plt.legend(loc='upper left')
	plt.savefig('input.png')
	plt.show()

#Plot the time of an approximation and a marching cubes sampling as we increase the order:	
#Parameters: input points, f, the domain, weighing function theta and the range of orders of the polynomial space.
def time_order(points, f, domain, theta, ms):
	X, Y, Z = domain
	approx_times = [None] * len(ms)
	sample_times = [None] * len(ms)
	i = 0
	for m in ms:
		start = time.time()
		#Get the aproximation (comment/uncomment to choose the method).
		approx = R3.least_squares_R3(points, f(points), m) 
		#approx = R3.weighted_least_squares_R3(points, f(points), m, theta)
		#approx = R3.moving_least_squares_R3(points, f(points), m, theta)
		approx_times[i] = time.time() - start
			
		#Take samples:
		start = time.time()
		vertices, simplices,normals, values = measure.marching_cubes_lewiner(approx(X,Y,Z), -0.3)
		sample_times[i] = time.time() - start
		i+=1
		
	width = 0.2
	plt.bar(ms, approx_times, width, color='b', label="Approximation time")
	plt.bar(ms, sample_times, width, color='r', label="Sample time", bottom=approx_times)
	plt.xlabel('Order of the polynomial space')
	plt.ylabel('Time')
	plt.legend(loc='upper left')
	plt.savefig('order.png')
	plt.show()
	
domain = np.mgrid[-4:4:50j, -4:4:50j, -4:4:50j]
test = np.mgrid[-4:4:6j, -4:4:6j, -4:4:6j].reshape(3,-1).T
ms = range(2, 10) 
theta = lambda d : ((1-d/8)**4)*(4*d/8 + 1)

#Define the 3D function:
def func(arr):
	res = np.zeros((len(arr),1))
	i = 0
	for p in arr:
		res[i] = np.cos(p[0]) + np.cos(p[1]) + np.cos(p[2])
		i+=1
	return res
	
# time_order(test, func, domain, theta, ms)

# sizes = range(4, 16) 
# time_input(func, domain, theta, sizes, 3)