from igraph import *
import psycopg2
from psycopg2.extensions import AsIs
from psycopg2.extras import execute_values
import sys
from config import config
import numpy as np
import datetime

from pprint import pprint


def main(argv):
	if (len(sys.argv) < 3 or len(sys.argv) > 3):
		print(
			"Debe ingresar nombre de mapa de caminos y nombre de mapa con puntos a conectar")
	else:
		map_name = argv[1]
		points_name = argv[2]
		conn = None
		try:
			# read connection parameters
			params = config()
			# connect to the PostgreSQL server
			print('Connecting to the PostgreSQL database')
			conn = psycopg2.connect(**params)
			print('Connected!\n')
			# create a cursor
			# conn.cursor will return a cursor object, you can use this cursor to perform queries
			cursor = conn.cursor()
			# execute queries
			currentTime = datetime.datetime.now()
			total_cost = build_paths(map_name, points_name, cursor)
			timeDelta = datetime.datetime.now() - currentTime
			# close communication with the PostgreSQL database server
			cursor.close()
			# commit the changes
			conn.commit()
			print(timeDelta)
			print("Costo total de construcción: ", total_cost)
		except (Exception, psycopg2.DatabaseError) as error:
			print(error)
		finally:
			if conn is not None:
				conn.close()


# create table with segments that need to be built
def build_paths(map_name, points_name, cursor):
	# create a results table
	needed_paths_name = map_name + "_suurballe"
	paths_to_build = map_name + "_construccion"
	cursor.execute(
		"CREATE TABLE IF NOT EXISTS %s  (path_id int, id int generated by default as identity PRIMARY KEY);",
		(AsIs(needed_paths_name),))
	find_paths(map_name, points_name, cursor, needed_paths_name)
	cursor.execute(
		"CREATE TABLE %s AS SELECT distinct m.camino, m.costo, c.path_id FROM %s m, %s c WHERE c.path_id = m.id AND m.costo != 0;",
		(AsIs(paths_to_build), AsIs(map_name), AsIs(needed_paths_name)))
	cursor.execute("select sum(costo) from %s;", (AsIs(paths_to_build),))
	cost = cursor.fetchone()[0]
	return cost


# find two disjoint paths between each pair of points, except for capital cities
def find_paths(map_name, points_name, cursor, needed_paths_name):
	# create graph from map linestrings, in which the vertex name is its postgis geometry, edge weight is the building cost and edge id is the corresponding id in table
	cursor.execute("select origen, fin, id, costo from %s where camino notnull;", (AsIs(map_name),))
	edges = cursor.fetchall()
	graph = Graph.TupleList(edges, directed=False, vertex_name_attr="name", edge_attrs=("id", "weight"))

	# get valid points (comunas)
	points = get_points(cursor, points_name)
	# create a grid to mark found paths
	points_grid = np.zeros((len(points), len(points)))
	# fill the grid's diagonal
	np.fill_diagonal(points_grid, 2)
	# get capital cities and peripheral points
	capitales = get_points(cursor, "capitales_regionales")
	limitrofes = get_points(cursor, "limitrofes")

	# for each point we must connect
	for p in range(len(points)):
		geom = points[p]
		try:
			s = graph.vs.find(name=geom)
		except:
			graph.add_vertex(name=geom)
			s = graph.vs.find(name=geom)
	for p in range(len(points)):
		print("buscando desde comuna ", p)
		geom = points[p]
		# list of target comunas
		target = compute_targets(graph, points, points_grid, p)
		# source vertex associated to p
		try:
			s = graph.vs.find(name=geom)
		except:
			graph.add_vertex(name=geom)
			s = graph.vs.find(name=geom)
		# if source is an isolated point, it only needs 1 path
		if geom in limitrofes:
			paths = get_single_path(graph, s, target)
			all_ids = list(
				process_paths(paths, p, graph, points, target, points_grid, map_name, points_name, cursor, geom, single=True))
		else:
			# list of disjoint paths from s to each target
			paths = suurballe(graph, s, target)
			# process paths to obtain the corresponding table IDs
			all_ids = process_paths(paths, p, graph, points, target, points_grid, map_name, points_name, cursor, geom)
			# capitals must have three paths
			if geom in capitales and geom:
				# check if third connection hasn't been registered already
				if max_connections(points_grid, p) < 3:
					try:
						thrid_path = find_third_path(graph, points, s, p, all_ids, target, points_grid)
						all_ids = list(all_ids) + thrid_path
					except:
						all_ids = list(all_ids) + \
											process_new_path(map_name, points_name, cursor, geom, graph, p, points, points_grid)
		# add IDs to psql table
		plain_ids = [str(x[0]) for x in all_ids]
		log = open('caminos.txt', 'a')
		log.write("\n".join(plain_ids))
		log.write("\n")
		log.close()
		update_table(needed_paths_name, cursor, all_ids)
	# create missing paths
	update_table(needed_paths_name, cursor,
							 create_direct_paths(map_name, points_name, cursor, graph, points_grid, points))
	pprint(points_grid)


# Process the paths obtained, filling the grid accordingly and returning all edges' ID
def process_paths(paths, p, graph, points, target, points_grid, map_name, points_name, cursor, geom, single=False):
	all_ids = set([])
	# for each comuna's paths
	for path in paths:
		# obtain target point corresponding to paths
		target_index = paths.index(path)
		point_index = points.index(graph.vs[target[target_index]]['name'])
		inc = 0
		if single and path:
			# only one path per comuna
			inc = 1
			fill_grid(points_grid, p, point_index, inc)
			joint_paths = path
		elif path:
			# first and second are tuple lists like ((i, j), id)
			first = path[0]
			second = path[1]
			# how many paths were found
			inc = count_paths(first, second)
			# fill points grid accordingly
			fill_grid(points_grid, p, point_index, inc)
			joint_paths = first + second
		# add paths to list of needed paths
		if inc:
			for (i, j), id in joint_paths:
				all_ids.add((id,))
	return all_ids


# returns the number of paths found
def count_paths(first, second):
	if (first and second):
		return 2
	elif first or second:
		return 1
	else:
		return 0


def create_direct_paths(map_name, points_name, cursor, graph, points_grid, points):
	new_edges = set([])
	for i in range(len(points)):
		for j in range(len(points)):
			# we must create a path to nearest connection
			if points_grid[i][j] == 0:
				new_edge = process_new_path(map_name, points_name, cursor, points[i], graph, i, points, points_grid)
				if new_edge:
					new_edges.add((new_edge[0],))
	return new_edges


# creates a new connection between p and the nearest comuna
def process_new_path(map_name, points_name, cursor, geom, graph, p, points, points_grid):
	closest = get_closest(cursor, points_name, geom)
	if points_grid[points.index(geom)][points.index(closest)] == 0:
		point_geom, new_edge_id = create_shortest_connection(map_name, cursor, geom, points_name=points_name)
		graph.add_edge(geom, point_geom, weight=0, id=new_edge_id)
		fill_grid(points_grid, p, points.index(point_geom), 1)
		return [new_edge_id]
	else:
		return []


def get_closest(cursor, points_name, geom):
	cursor.execute(
		"select b.pun_geom, min(ST_DistanceSphere(a.pun_geom, b.pun_geom)) m from %s a,%s b where a.pun_geom=%s and a.pun_id != b.pun_id group by b.pun_geom order by m limit 1;",
		(AsIs(points_name), AsIs(points_name), geom))
	target, distance = cursor.fetchall()[0]
	return target


# creates new connection between two vertex
def new_connection(map_name, cursor, geom, graph, points, target, points_grid):
	target, new_edge_id = create_shortest_connection(map_name, cursor, geom, target=target)
	graph.add_edge(geom, target, weight=0, id=new_edge_id)
	fill_grid(points_grid, points.index(geom), points.index(target), 1)
	return new_edge_id


# Returns list of shortest paths from source to each target in target_list
def get_single_path(graph, source, target_list):
	paths = graph.get_shortest_paths(source, weights="weight", to=target_list, output="epath")
	return get_edges(graph, paths)


# returns the points that don't yet have 2 paths to 'p'
def compute_targets(graph, points, points_grid, p):
	target = []
	# for each point
	for q in range(len(points)):
		# if it hasn't been computed
		if points_grid[p][q] == 0:
			t = graph.vs.find(name=points[q])
			target.append(t.index)
	return target


# fills points grid given paths from source to target according to increment
def fill_grid(points_grid, source_index, target_index, increment):
	points_grid[source_index][target_index] += increment
	points_grid[target_index][source_index] += increment


# returns two disjoint shortest paths from source vertex to each point in target_list
def suurballe(graph, source, target_list, directed=False):
	if directed == False:
		graph.to_directed()
	# get shortest path to target vertices using Dijkstra
	edge_path = graph.get_shortest_paths(source, weights="weight", to=target_list, output="epath")

	# get shortest distance to all vertices using Dijkstra
	distance = graph.shortest_paths(source, weights="weight")[0]

	shortest_path = get_edges(graph, edge_path)
	# problem_vertex contains vertices without two paths
	result = []

	# create transformed graph Gv for each V. Future: one unified graph
	for t in range(len(target_list)):
		v = graph.vs.find(target_list[t]).index
		if (distance[v] != float("inf") and target_list[t] != source):
			gv = create_graph(graph, distance, edge_path, t)
			transformed_path = gv.get_shortest_paths(source, weights="weight", to=target_list[t], output="epath")
			shortest_transformed = get_edges(gv, transformed_path)[0]
			# unify paths discarding edges that are in both paths
			union = unify_paths(shortest_transformed, shortest_path[t])
			# obtain paths from union
			path_1 = get_path(union, source.index, v)
			try:
				path_2 = get_path(union, source.index, v)
				# success.append(v, path_1, path_2)
				pair = (path_1, path_2)
				result.append(pair)
			except Exception:
				result.append((path_1, []))
		# if vertex is unreachable
		elif distance[v] == float("inf"):
			result.append(([], []))

		else:
			result.append([])
	return result


# find a third disjoint path between source and target
def find_third_path(graph, points, source, p, path_ids, target, points_grid):
	# for psql reasons the format is (id,)
	edge_index = graph.es.select(lambda edge: (edge["id"],) in list(path_ids))
	# copy the graph so we can delete the edges
	graph_copy = graph.copy()
	graph_copy.delete_edges(edge_index)
	# find shortest distance to the other vertices
	new_distance = graph_copy.shortest_paths_dijkstra(source, target=target, weights="weight", mode=OUT)
	try:
		# select target with shortest distance
		target_index = target.index(min(new_distance[0]))
		points_index = points.index(graph.vs[target[target_index]]['name'])
		# get path to selected target
		path = graph_copy.get_shortest_paths(source, weights="weight", to=target[target_index], output="epath")[0]
		if path:
			new_path = []
			for s in path:
				# obtain actual edge id
				id = graph_copy.es.find(s)["id"]
				new_path.append(id)
			fill_grid(points_grid, p, points_index, 1)
		else:
			pass
	except Exception:
		# must create new path from scratch
		raise Exception()

	return path


# insert paths into table
def update_table(table_name, cursor, path_ids):
	if path_ids:
		query = "insert into " + table_name + " (path_id) values %s;"
		execute_values(cursor, query, path_ids)


# returns copy of graph in which tree edges are reversed and edge weights are transformed for Suurballe
def create_graph(original_graph, distance, edge_path, target):
	gv = Graph(directed=True)
	gv.add_vertices(len(original_graph.vs))
	reverse_edges(gv, original_graph, edge_path, target)
	gv.es["weight"] = transform_attributes(original_graph, distance)
	return gv


# returns union of paths without complement edges
def unify_paths(path1, path2):
	for ((x, y), z) in path1:
		if ((y, x), z) in path2:
			path2.remove(((y, x), z))
			path1.remove(((x, y), z))
	return path2 + path1


# flips edges in gv that are in edge_path
def reverse_edges(gv, graph, edge_path, v):
	gv_edge_id = []
	for e in range(len(graph.get_edgelist())):
		edge = graph.es[e]
		source = edge.source
		target = edge.target
		if e in edge_path[v]:
			gv.add_edge(target, source)
			gv_edge_id.append(edge["id"])
		else:
			gv.add_edge(source, target)
			gv_edge_id.append(edge["id"])
	gv.es["id"] = gv_edge_id


# returns edges in edge_path as source and target IDs and edge ID (from map)
def get_edges(graph, edge_path):
	edges = graph.get_edgelist()
	edges_as_pairs = []
	for path in edge_path:
		path_pairs = []
		for e in path:
			path_pairs.append((edges[e], graph.es[e]["id"]))
		edges_as_pairs.append(path_pairs)
	return edges_as_pairs


# constructs a path from source to target with given edges
def get_path(edges, source, target):
	path = []
	edges.sort(key=lambda x: (x[1], x[0]))
	try:
		current = search_tuple(edges, target)
		path.append(current)
		while (current[0][0] != source):
			current = search_tuple(edges, current[0][0])
			path.append(current)
	except Exception:
		raise Exception
	return path


# looks for pair in which the target vertix matches with elem
def search_tuple(tups, elem):
	result = list(filter(lambda tup: tup[0][1] == elem, tups))
	try:
		tups.remove(result[0])
		return result[0]
	except IndexError:
		raise Exception


# transform edge_costs according to Suurballe algorithm
def transform_attributes(graph, distance):
	transformed_costs = []
	for i in range(len(graph.es)):
		d_target = distance[graph.es[i].target]
		d_source = distance[graph.es[i].source]
		if d_target == float("inf") or d_source == float("inf"):
			transformed_costs.append(float("inf"))
		else:
			transformed_costs.append(graph.es[i]["weight"] - d_target + d_source)
	return transformed_costs


# returns psycopg result as a python list
def clean_list(psycopg_list):
	clean = []
	for p in psycopg_list:
		clean.append(p[0])
	return clean


# returns points from a table, if map_name is specified, only returns points that intersect map
def get_points(cursor, points_name, map_name=None):
	if map_name == None:
		cursor.execute("select ST_snaptogrid(pun_geom, 0.00001) from %s;", (AsIs(points_name),))
		return clean_list(cursor.fetchall())
	else:
		cursor.execute("select distinct ST_snaptogrid(pun_geom, 0.00001) from %s, %s\
				 where st_intersects(st_snaptogrid(pun_geom, 0.00001), st_snaptogrid(camino, 0.00001));",
									 (AsIs(map_name), AsIs(points_name)))
		points_list = cursor.fetchall()
		return clean_list(points_list)


# returns max number of connections that p has
def max_connections(points_grid, p):
	return max(points_grid[p])


# creates the shortest connection from source to any comuna, returns
def create_shortest_connection(map_name, cursor, geom, target=None, points_name=None):
	if target == None:
		# find building cost and point id of closest comuna
		cursor.execute(
			"select b.pun_geom, min(ST_DistanceSphere(a.pun_geom, b.pun_geom)) m from %s a,%s b where a.pun_geom=%s and a.pun_id != b.pun_id group by b.pun_geom order by m limit 1;",
			(AsIs(points_name), AsIs(points_name), geom))
		target, distance = cursor.fetchall()[0]
	else:
		cursor.execute(
			"SELECT ST_DistanceSphere(%s, %s) ;",
			(geom, target))
		distance = cursor.fetchone()[0]
	cursor.execute("select cost_per_km from costs_table where type='nada';")
	path_cost = round((distance / 1000) * math.sqrt(2) * cursor.fetchone()[0])
	new_path_log = open('caminos_nuevos.txt', 'a')
	new_path_log.write("\t".join((str(geom), str(target), str(path_cost))))
	new_path_log.write("\n")
	new_path_log.close()
	# camino geometry, x1 float, y1 float, x2 float, y2 float, tipo varchar(255), id_origen int, largo float, costo int, tabla_origen varchar(255), origen geometry, fin geometry
	cursor.execute(
		"INSERT INTO %s (camino, x1, y1, x2, y2, tipo, largo, costo, origen, fin) VALUES (ST_makeline(%s, %s), ST_X(%s), ST_Y(%s), ST_X(%s), ST_Y(%s), 'nada' ,%s, %s, %s,  %s ) returning id; ", \
		(AsIs(map_name), geom, target, geom, geom, target, target, distance, path_cost, geom, target))
	edge_id = cursor.fetchone()[0]
	return target, edge_id


if __name__ == "__main__":
	main(sys.argv)
