from __future__ import print_function
import itertools
import numpy as np
from numpy.linalg import norm
from bbcif_properties import X_vecs
from place_bbs import superimpose, mag_superimpose
from ciftemplate2graph import node_vecs
import warnings

def vertex_assign(TG, TVT, node_cns, unit_cell, SYM_TOL, ALL_NODE_COMBINATIONS):

	node_dict = dict((k,[]) for k in TVT)

	for node in node_cns:
		for k in TVT:
			if node[0] == k[0]:
				node_dict[k].append(node[1])

	print('RMSD of the compatible node BBs with assigned vertices:')
	print()
	
	sym_assign = []
	sym_assign_append = sym_assign.append

	for k in node_dict:

		print('vertex', k[1], '('+str(k[0]) + ' connected)')

		matched = 0
		unmatched = 0
		
		if len(node_dict[k]) == 0:
			continue
		coord_num = k[0]

		for n in TG.nodes(data=True):
			name,ndict = n
			distances = []
			distances_append = distances.append

			if ndict['type'] == k[1]:
				for cif in node_dict[k]:

					nvec = np.array([v/np.linalg.norm(v) for v in node_vecs(name, TG, unit_cell, False)])
					bbxvec = np.array([v/np.linalg.norm(v) for v in X_vecs(cif, False)])
					rmsd, _, _ = superimpose(bbxvec, nvec)
					distances_append((rmsd,cif))

				for d in distances:
					disp, cif = d
					if d[0] < SYM_TOL[coord_num]:
						matched += 1
						matches = '(within tolerance)'
					else:
						unmatched += 1
						matches = '(outside tolerance)'
					print('    ', cif, 'deviation =', np.round(disp,5), matches)

				for d in distances:
					if d[0] < SYM_TOL[coord_num]:
						sym_assign_append((k[1], d[1], d[0]))
				break
		print('*', matched, 'compatible building blocks out of', len(node_dict[k]), 'available for node', k[1], '*')
	print()
	
	rearrange = dict((k[1], []) for k in TVT)
	for a in sym_assign:
		rearrange[a[0]].append((a[0], a[1], a[2]))

	va_uncomb = [rearrange[a] for a in rearrange]
	
	for i in range(len(va_uncomb)):
		va_uncomb[i] = sorted(va_uncomb[i], key=lambda x:x[-1])

	va = []
	va_append = va.append
	used = []
	used_append = used.append
	for l in itertools.product(*va_uncomb):

		cifs = sorted(tuple([c[1] for c in l]))
		if cifs in used and not ALL_NODE_COMBINATIONS:
			continue

		choice_dict = dict((i[0],i[1]) for i in l)
		va_temp = []
		va_temp_append = va_temp.append

		for n in TG.nodes(data=True):
			name,ndict = n
			va_temp_append((name, choice_dict[ndict['type']]))

		va_append(va_temp)
		used_append(cifs)
					
	return va

def assign_node_vecs2edges(TG, unit_cell, template_name):
	
	edge_assign_dict = dict((k,{}) for k in TG.nodes())

	for n in TG.nodes(data=True):

		name, ndict = n
		cif = ndict['cifname']
		
		bbxlabels = np.array([l[0] for l in X_vecs(cif, True)])
		nodlabels = np.array([l[0] for l in node_vecs(n[0], TG, unit_cell, True)])

		bbxvec = X_vecs(cif, False)
		nodvec = node_vecs(n[0], TG, unit_cell, False)
		
		_,rot,tran = mag_superimpose(bbxvec, nodvec)
		aff_b = np.dot(bbxvec, rot) + tran
		laff_b = np.c_[bbxlabels, aff_b]
		lnodvec = np.c_[nodlabels, nodvec]
		
		asd = []
		asd_append = asd.append

		distance_matrix = np.zeros((len(laff_b), len(laff_b)))
		nrow = ncol = len(laff_b)
		
		for i in range(nrow):
			for j in range(ncol):
		
				v1 = laff_b[i]
				v1vec = np.array([float(q) for q in v1[1:]])
				v1vec /= norm(v1vec)
		
				v2 = lnodvec[j]
				v2vec = np.array([float(q) for q in v2[1:]])
				v2vec /= norm(v2vec)
		
				dist = np.linalg.norm(v1vec - v2vec)
				distance_matrix[i,j] += dist
		
		distances = []
		for i in range(nrow):
			for j in range(ncol):
				distances.append((distance_matrix[i,j], i, j))
		distances = sorted(distances, key=lambda x:x[0])
		
		used_edges = []
		
		for dist in distances:
		
			v1 = laff_b[dist[1]]
			v1vec = np.array([float(q) for q in v1[1:]])
			mag = np.linalg.norm(v1vec)
		
			v2 = lnodvec[dist[2]]
			ind = int(v2[0])
		
			edge_assign = ind
		
			if edge_assign not in used_edges:

				used_edges.append(edge_assign)
				asd_append([ind, v1[0], mag, v1vec, dist[0]])
				
				if dist[0] > 0.60:
					message = "There is a nodular building block vector that deviates from its assigned edge by more large\nthis may be fixed during scaling, but don't count on it!\n"
					message = message + "the deviation is for " + cif + " assigned to " + name + " for template " + template_name
					warnings.warn(message)
			
			if len(used_edges) == ncol:
				break

		elad = dict((k[0], (k[1], k[2], k[3])) for k in asd)
		edge_assign_dict[name] = elad
					
	return edge_assign_dict
