from __future__ import print_function
import numpy as np
from scipy.optimize import minimize
import random
from ase.geometry import cellpar_to_cell


def metric_tensor(l):
	cell = cellpar_to_cell(l)
	Z = cell @ cell.T
	return Z

def objective(V, ncra, ncca, Alpha, ne, nv, Bstar_inv, SBU_IP):
	
	P = V[0:6]
	Z = metric_tensor(P)
	X = np.array(V[6:])
	X = X.reshape(ncra,ncca)
	var_alpha = np.r_[Alpha[0:ne-nv+1,:], X]

	omp = np.dot(Bstar_inv, var_alpha)
	g = np.dot(omp, np.dot(Z,omp.T))

	O1 = 0
	for sbu in SBU_IP:
		O2 = 0
		for i in sbu[1]:
			O2 = O2 + (i[2] - g[i[0] - 1][i[1] - 1])**2
		O1 = O1 + O2

	return O1

def scale(all_SBU_coords,
		  a,b,c,ang_alpha,ang_beta,ang_gamma,
		  max_le,
		  num_vertices,
		  Bstar,
		  alpha,
		  num_edges,
		  fix_uc,
		  n_iters,
		  pre_scale,
		  min_cell_length):

	max_length = 0
	for line in all_SBU_coords:
		for length in [np.linalg.norm(s[1]) for s in line[1]]:
			if length > max_length:
				max_length = length

	if pre_scale == 'none':
		scale_guess = 1.0
	else:
		scale_guess = (max_length/max_le) * pre_scale

	all_SBU_ip = []
	all_SBU_ip_append = all_SBU_ip.append
	for sbu in all_SBU_coords:
		SBU_ip = []
		SBU_ip_append = SBU_ip.append
		w = len(sbu[1])
		for i in range(w):
			ivec = sbu[1][i][1]
			iind = sbu[1][i][0]
			for j in range(i, w):
				jvec = sbu[1][j][1]
				jind = sbu[1][j][0]
				dot  = np.dot(ivec,jvec)
				SBU_ip_append([iind,jind,dot])
		all_SBU_ip_append((sbu[0], SBU_ip))
	ncra = num_vertices - 1
	ncca = 3

	covars_values = []
	covars_values_append = covars_values.append
	for i in range(ncra):
		for j in range(ncca):
			covars_values_append(0)
			
	ucvars = [a,b,c,ang_alpha,ang_beta,ang_gamma]

	if np.any(fix_uc):

		uc_bounds = []
		uc_bounds_append = uc_bounds.append
		for f,p in zip(fix_uc, ucvars):
			if f:
				uc_bounds_append((p,p))
			else:
				uc_bounds_append((0,None))
		uc_bounds = tuple(uc_bounds)
	else:
		max_a = 2 * scale_guess * a
		max_b = 2 * scale_guess * a
		max_c = 2 * scale_guess * a
		uc_bounds = ((min_cell_length, max_a),
					(min_cell_length, max_b),
					(min_cell_length, max_c),
					(20,160),
					(20,160),
					(20,160))

	init_variables = [scale_guess * a,
						scale_guess * b,
						scale_guess * c,
						ang_alpha,
						ang_beta,
						ang_gamma] + covars_values
	x_bounds = tuple([(-10.0,10.0) for _ in covars_values])
	bounds = uc_bounds + x_bounds

	Bstar_inv = np.linalg.inv(Bstar)

	uc_press = 0.05
	covars_perturb = 0.0001
	callbackresults = [[init_variables, objective(init_variables,
													ncra,
													ncca,
													alpha,
													num_edges,
													num_vertices,
													Bstar_inv,
													all_SBU_ip)]]

	def callbackF(X):
		var_string = ''
		for val in X:
			var_string += str(val) + '   '
		callbackresults.append([X, objective(X,ncra,ncca,alpha,num_edges,num_vertices,Bstar_inv,all_SBU_ip)])

	for it in range(n_iters):
		
		res = minimize(objective,
						init_variables,
						args=(ncra,ncca,alpha,num_edges,num_vertices,Bstar_inv,all_SBU_ip),
						method='L-BFGS-B',
						bounds=bounds,
						options={'disp':False},
						callback=callbackF)

		uc_params = res.x[0:6]
		uc_params = [i - (i * uc_press) for i in uc_params]
		mult = [random.choice([-1, 1]) for i in range(len(res.x[6:]))]
		covars = [i + j * (i * covars_perturb) for i,j in zip(res.x[6:], mult)]
		init_variables = uc_params + covars

	if n_iters != 0:	
		sc_a,sc_b,sc_c,sc_alpha,sc_beta,sc_gamma = res.x[0:6]
		sc_covar = res.x[6:].reshape(ncra,ncca)
	else:
		init_variables = np.asarray(init_variables)
		sc_a,sc_b,sc_c,sc_alpha,sc_beta,sc_gamma = init_variables[0:6]
		sc_covar = init_variables[6:].reshape(ncra,ncca)

	ab = [a/b, sc_a/sc_b]
	ac = [a/c, sc_a/sc_c]
	bc = [b/c, sc_b/sc_c]
	alpha = [ang_alpha, sc_alpha]
	beta = [ang_beta, sc_beta]
	gamma = [ang_gamma, sc_gamma]
	covar = [np.average(abs(np.array(callbackresults[0][0][6:]) - np.array(callbackresults[-1][0][6:])))]
	final_obj = [callbackresults[-1][1]]

	print('The final objective function value is', np.round(final_obj[0],3))
	print()

	scaling_data = [ab, ac, bc, alpha, beta, gamma, covar, final_obj]

	return(sc_a,sc_b,sc_c,sc_alpha,sc_beta,sc_gamma,
		sc_covar,Bstar_inv,max_length,callbackresults,ncra,ncca,scaling_data)

