import numpy as np
import re 

from place_bbs import superimpose
from utils import PBC3DF_sym


def adjust_edges(placed_edges, placed_nodes, sc_unit_cell):

	adjusted_placed_edges = []
	adjusted_placed_edges_extend = adjusted_placed_edges.extend

	placed_edges = np.asarray(placed_edges)
	edge_labels = set(map(int, placed_edges[:, -1]))

	edge_dict = dict((k, []) for k in edge_labels)

	node_connection_points = [list(map(float, i[1:4])) for i in placed_nodes if re.sub('[0-9]', '', i[4]) == 'X']
	
	for edge in placed_edges:
		ty = int(edge[-1])
		edge_dict[ty].append(edge)

	for k in edge_dict:

		edge = np.asarray(edge_dict[k])
		elems = edge[:, 0]
		evecs = [list(map(float, i)) for i in edge[:, 1:4]]
		cp = edge[:, 4]
		ty = edge[:, 5]
		
		xvecs = [list(map(float, i)) for (i, j) in zip(evecs,cp) if re.sub('[0-9]', '', j) == 'X']
		relavent_node_xvecs = []
		relavent_node_xvecs_append = relavent_node_xvecs.append

		for ex in xvecs:

			min_dist = (1e6, [], 0)

			f_ex = np.dot(np.linalg.inv(sc_unit_cell), ex)
			for i in range(len(node_connection_points)):
				nx = node_connection_points[i]
				f_nx = np.dot(np.linalg.inv(sc_unit_cell), nx)

				fdist_vec,sym = PBC3DF_sym(f_ex, f_nx, "adjust_edges") 
				cdist = np.linalg.norm(np.dot(sc_unit_cell, fdist_vec))

				if cdist < min_dist[0]:
					min_dist = (cdist, np.dot(sc_unit_cell,f_nx + sym), i)

			node_connection_points.pop(min_dist[2])
			relavent_node_xvecs_append(min_dist[1])

		ecom = np.average(xvecs, axis=0)
		rnxcom = np.average(relavent_node_xvecs, axis=0)

		evecs = np.asarray(evecs - ecom)
		xvecs = np.asarray(xvecs - ecom)
		relavent_node_xvecs = np.asarray(relavent_node_xvecs)

		trans = rnxcom
		min_dist,rot, _ = superimpose(xvecs, relavent_node_xvecs)
		adjusted_evecs = np.dot(evecs, rot) + trans
		adjusted_edge = np.column_stack((elems,adjusted_evecs, cp, ty))
		adjusted_placed_edges_extend(adjusted_edge)

	return (adjusted_placed_edges)		
