#!/usr/bin/env python3
import pygad
import numpy as np
import networkx as nx
from matplotlib import pyplot as plt
import random
import itertools
import sys
import os
import logging
import math
import portion
from collections import defaultdict
logging.basicConfig(encoding="utf-8", level=logging.__dict__.get(os.getenv("LOGLEVEL", default="DEBUG")))
logging.info("starting")
def remove_cycles (pot): # za task2 poganjaj samo na odsekih!
	pot = np.array(pot)
	for node in set(pot):
		poz = np.where(pot == node)[0]
		if len(poz) > 1:
			pot = np.concatenate([pot[:poz[0]], pot[poz[-1]:]])
	return list(pot)
def naključen (n, p):
	while True:
		g = nx.DiGraph()
		g.add_nodes_from(range(n))
		for i in range(n):
			for j in range(n):
				if i != j and random.random() < p:
					g.add_edge(i, j, weight=round(random.uniform(0.0, 100.0), 1))
		if nx.is_strongly_connected(g):
			return g
def koskromosoma2pot (kromosom):
	# logging.debug(f"koskromosoma2pot called with {kromosom=}")
	if -1 in kromosom:
		return list(map(int, kromosom[0:list(kromosom).index(-1)]))
	return list(kromosom)
def pot2koskromosoma (pot, kromosomlen):
	assert pot.count(postaj-1) == 2
	return np.array(list(pot)+[-1]*(kromosomlen-len(pot)))
def kromosom2poti (kromosom):
	dolžinakosa = len(kromosom)//agentov
	assert agentov*dolžinakosa == len(kromosom)
	return [koskromosoma2pot(kromosom[i*dolžinakosa:(i+1)*dolžinakosa]) for i in range(agentov)]
def poti2kromosom (poti, kromosomlen):
	# logging.debug(f"poti2kromosom received {poti=} {kromosomlen=}")
	dolžinakosa = kromosomlen//agentov
	assert agentov*dolžinakosa == kromosomlen
	geni = []
	for pot in poti:
		dodaj = pot2koskromosoma(pot, dolžinakosa).tolist()
		# logging.debug(f"poti2kromosom: pot2koskromosoma vrnil {len(dodaj)=}")
		geni += dodaj
	# logging.debug(f"poti2kromosom returning {geni=}")
	return np.array(geni)
maxura = 0.0 # hack, da lahko preberemo lokalni spremenljivki v fitness
overlap = 0.0 # seveda samo za zadnjo izračunano vrednost
def intervallen (i):
	dolž = 0
	for c in i:
		if c.lower == -math.inf or c.upper == math.inf:
			return math.inf
		if c.lower == math.inf or c.upper == -math.inf:
			continue
		dolž += c.upper - c.lower
	return dolž
def fitness_function_factory (g):
	def fitness (ga_instance, solution, solution_idx):
		global maxura
		global overlap
		poti = kromosom2poti(solution)
		logging.debug(f"fitness obdeluje {poti=}")
		fitnes_od_dolžin = 0
		for agentpot in poti:
			agentpot = agentpot[:]
			while -2 in agentpot:
				agentpot.remove(-2)
			logging.debug(f"fitness računa uteženo dolžino {agentpot=}")
			fitnes_od_dolžin += -nx.path_weight(g, agentpot, weight='weight')+len(agentpot)/(len(fitnesses)+1)
		if agentov == 1:
			return fitnes_od_dolžin
		ezasedenost = defaultdict(portion.empty) # maps povezava (u, v) to zasedenost (interval set)
		vzasedenost = defaultdict(portion.empty) # maps vozlišče u to zasedenost (interval set)
		overlap = 0.0
		maxura = 0.0
		for agent in range(agentov):
			ura = 0.0
			curnode = poti[agent][0]
			for vozlidx in range(1, len(poti[agent])):
				newnode = poti[agent][vozlidx]
				if newnode == -2:
					dodaj = portion.closed(ura, ura+20)
					overlap += intervallen(vzasedenost[curnode] & dodaj) #(vzasedenost[curnode] & dodaj).upper - (vzasedenost[curnode] & dodaj).lower
					vzasedenost[curnode] |= dodaj
					ura += 20
				else:
					utež = g.get_edge_data(curnode, newnode)["weight"]
					edodaj = portion.closed(ura, ura+utež)
					overlap += intervallen(ezasedenost[(curnode, newnode)] & edodaj) # (ezasedenost[(curnode, newnode)] & edodaj).upper - (ezasedenost[(curnode, newnode)] & edodaj).lower
					ezasedenost[(curnode, newnode)] |= edodaj
					ura += utež
					curnode = newnode
					vdodaj = portion.closed(ura, ura+10)
					overlap += intervallen(vzasedenost[curnode] & vdodaj) # (vzasedenost[curnode] & vdodaj).upper - (vzasedenost[curnode] & vdodaj).lower
					vzasedenost[newnode] |= vdodaj
					ura += 10
			if ura > maxura:
				maxura = ura
		return -maxura-4*overlap # overlap je nezaželjen, zato ga močno obtežimo
	return fitness
def crossover_func (parents, offspring_size, ga_instance):
	print("crossover ... ", end="", file=sys.stderr)
	sys.stderr.flush()
	offspring = []
	idx = 0
	lastadded = 0
	dolžinakromosoma = parents.shape[1]
	while len(offspring) != offspring_size[0]:
		idx += 1
		očepoti = kromosom2poti(parents[idx % parents.shape[0], :])
		if lastadded + len(parents)*5 < idx:
			logging.warn("ADDING parents as offsprings!")
			while len(offspring) != offspring_size[0]:
				offspring.append(poti2kromosom(očepoti, dolžinakromosoma))
			break
		matipoti = kromosom2poti(parents[(idx + 1) % parents.shape[0], :])
		matipoti = list(map(razdeli_po_postajah, matipoti))
		očepoti = list(map(razdeli_po_postajah, očepoti))
		otrokpoti = []
		for agent in range(agentov):
			if agentov == 1:
				odseki = list(range(postaj-1))
			else:
				odseki = list(range(postaj))
			random.shuffle(odseki)
			otrokpot = očepoti[agent][0:odseki[0]] + matipoti[agent][odseki[0]:]
			for pivotni_odsek in odseki:
				presek = np.intersect1d(očepoti[agent][pivotni_odsek], matipoti[agent][pivotni_odsek])
				presek.aslist()
				while -2 in presek:
					presek.remove(-2)
				if len(presek) == 0:
					continue
				ran = random.choice(presek)
				otrok = očepoti[agent][0:pivotni_odsek] + [očepoti[agent][pivotni_odsek][0:očepoti[agent][pivotni_odsek].index(ran)] + matipoti[agent][pivotni_odsek][matipoti[agent][pivotni_odsek].index(ran):]] + matipoti[agent][pivotni_odsek+1:]
				break
			otrokpoti.append(otrokpot)
		kromosom = poti2kromosom(map(združi_po_postajah, otrokpoti), dolžinakromosoma)
		if len(kromosom) <= dolžinakromosoma and kromosom.tolist() not in [x.tolist() for x in offspring]:
			print("crossover made ", end="")
			print(kromosom2poti(kromosom))
			lastadded = idx
			offspring.append(kromosom)
		else:
			print("crossover rejected due to too long or duplicate")
	print("done", file=sys.stderr)
	return np.array(offspring[0:offspring_size[0]])
explore_more = 0
def razdeli_po_postajah (pot):
	ret = []
	if agentov == 1: # task 1 and task 2
		for postaja in range(postaj-1):
			ret.append(pot[0:pot.index(postaja+1)+1])
			pot = pot[pot.index(postaja+1):]
	else:
		for postaja in range(postaj):
			ret.append(pot[0:pot.index(postaja)+1])
			pot = pot[pot.index(postaja+1):]
	return ret
def združi_po_postajah (odseki):
	ret = [odseki[0][0]]
	if agentov == 1:
		postaja = 0
	else:
		postaja = -1
	for odsek in odseki:
		logging.debug(f"{odsek=} {postaja=}")
		if postaja != -1:
			assert odsek[0] == postaja
		assert odsek[-1] == postaja+1
		# assert len(odsek[1:-1]) == len(set(odsek[1:-1])) # mutacija lahko naredi cikle, ko je remove_cycles_func identiteta
		if True: # and postaj > 2:
			assert len(set(range(0, postaj)).intersection(set(odsek[1:-1]))) == 0
		ret += odsek[1:]
		postaja += 1
	return ret
def g_brez (g, vozlišča):
	ret = g.copy()
	for vozlišče in vozlišča:
		ret.remove_node(vozlišče)
	return ret
def mutation_func_factory (g):
	def mutation_func (vhod, ga_instance):
		global explore_more
		logging.info("mutacija")
		sys.stderr.flush()
		ret = []
		# for chromosome_idx in range(offspring.shape[0]):
		idx = 0
		while len(ret) < len(vhod):
			allow_simple_paths = False
			izletpoti = kromosom2poti(vhod[idx % vhod.shape[0]])
			kromolen = len(vhod[idx % vhod.shape[0]])
			logging.debug(f"mutiram {izlet=}")
			novepoti = []
			for izletpotidx, izletpot in enumerate(izletpoti):
				odseki = []
				for odsekidx, odsek in enumerate(razdeli_po_postajah(izletpot)):
					opcije = ["naive", "keep", "node"]
					if agentov > 1:
						opcije.append("wait")
					def identity (arg):
						return arg
					remove_cycles_func = identity
					if explore_more > 0 or (len(fitnesses) >= 4 and len(set(fitnesses[-5:-1])) == 1):
						if explore_more > 0:
							explore_more -= 1
						else:
							explore_more = 3*postaj*agentov
						opcije.append("initial")
						remove_cycles_func = remove_cycles
					odsekbrezwait = odsek
					while -2 in odsekbrezwait:
						odsekbrezwait.remove(-2)
					if 1 <= len(odsekbrezwait)-1-1:
						odseka = None
						odsekb = None
						while odseka == None or odsekb == None or odsek[odseka] == -2 or odsek[odsekb] == -2:
							odseka = random.randint(0, len(odsek)-1-1)
							odsekb = random.randint(odseka+1, len(odsek)-1)
					else:
						opcije = ["initial", "keep", "wait"]
					if -2 in list(map(int, odsek)):
						opcije.append("unwait")
					sporna = set(range(0, postaj))-set([odsek[0], odsek[-1]])
					gbsp = g_brez(g, sporna) # g brez spornih postaj
					match random.choice(opcije):
						case "wait":
							novodsek = remove_cycles_func(odsek)
							novodsek = novodsek[odseka:] + [-2] + novodsek[odseka:]
							odseki.append(novodsek)
						case "unwait":
							novodsek = remove_cycles_func(odsek)[:]
							waiti = np.where(np.array(novodsek) == -2)[0]
							random.shuffle(waiti)
							novodsek = novodsek[0:waiti[0]] + novodsek[waiti[0]+1:]
							odseki.append(novodsek)
						case "keep":
							odseki.append(remove_cycles_func(odsek))
						case "naive":
							if random.choice([allow_simple_paths, False]):
								obvoz = list(random.choice(list(itertools.islice(nx.all_simple_paths(gbsp, odsek[odseka], odsek[odsekb], int((odsekb-odseka)*2)), len(g)))))
							else:
								logging.debug(f"{gbsp=} {odsek=} {odseka=} {odsekb=} {sporna=}")
								agobvoz = list(random.choice(list(itertools.islice(nx.all_shortest_paths(gbsp, odsek[odseka], odsek[odsekb]), len(g)))))
							odseki.append(remove_cycles_func(odsek[0:odseka] + obvoz + odsek[odsekb+1:]))
						case "node":
							for poskus in range(5):
								try:
									vozl = random.choice(list(set(range(postaj-1, len(g)))))
									if random.choice([allow_simple_paths, False]):
										obvoz1 = list(random.choice(list(itertools.islice(nx.all_simple_paths(gbsp, odsek[odseka], vozl, len(g)//2), len(g)))))
									else:
										obvoz1 = list(random.choice(list(itertools.islice(nx.all_shortest_paths(gbsp, odsek[odseka], vozl), len(g)))))
									if random.choice([allow_simple_paths, False]):
										obvoz2 = list(random.choice(list(itertools.islice(nx.all_simple_paths(gbsp, vozl, odsek[odsekb], len(g)//2), len(g)))))
									else:
										obvoz2 = list(random.choice(list(itertools.islice(nx.all_shortest_paths(gbsp, vozl, odsek[odsekb]), len(g)))))
									odseki.append(remove_cycles_func(odsek[0:odseka] + obvoz1 + obvoz2[1:] + odsek[odsekb+1:]))
									break
								except nx.exception.NetworkXNoPath:
									logging.debug(f"no path found in node type mutation, {poskus=}")
						case "initial":
							odseki.append(razdeli_po_postajah(inicialna(True)[agentidx])[odsekidx])
					if odsekidx in odseki[-1][1:-1] or odsekidx+1 in odseki[-1][1:-1]: # and postaj > 2:
						odseki[-1] = remove_cycles(odseki[-1])
				novepoti.append(združi_po_postajah(odseki))
			kromosom = poti2kromosom(poti, kromolen)
			if len(kromosom) == kromolen and kromosom.tolist() not in [x.tolist() for x in ret]:
				logging.debug(f"mutacija generirala nove {kromosom2poti(kromosom)=}!")
				ret.append(kromosom)
			else:
				logging.warn(f"mutacija: zavržena predolge ali podvojene {novepoti=}")
			idx += 1
		logging.info("konec mutacije")
		return np.array(ret)
	return mutation_func
last_fitness = -99999999
fitnesses = []
def on_generation (ga_instance):
	global last_fitness
	global fitnesses
	print(f"Generation = {ga_instance.generations_completed}")
	print(f"Fitness    = {ga_instance.best_solution(pop_fitness=ga_instance.last_generation_fitness)[1]}")
	print(f"Change     = {ga_instance.best_solution(pop_fitness=ga_instance.last_generation_fitness)[1] - last_fitness}")
	last_fitness = ga_instance.best_solution(pop_fitness=ga_instance.last_generation_fitness)[1]
	fitnesses.append(last_fitness)
inicialne_cache_dfs = []
inicialne_cache_bfs = []
def inicialna (bfs):
	ret = inicialna_impl(bfs)
	ret = [list(x) for x in ret]
	# logging.debug(f"inicialna returning {ret=}")
	return ret
def inicialna_impl (bfs):
	global inicialne_cache_dfs
	global inicialne_cache_bfs
	if bfs:
		if len(inicialne_cache_bfs) > 0:
			return inicialne_cache_bfs.pop()
		else:
			inicialne_cache_bfs = inicialne(True)
			return inicialne_cache_bfs.pop()
	else:
		if len(inicialne_cache_dfs) > 0:
			return inicialne_cache_dfs.pop()
		else:
			inicialne_cache_dfs = inicialne(False)
			return inicialne_cache_dfs.pop()
def inicialne (bfs):
	agentdo0 = []
	for agent in range(agentov):
		sporna = set(range(1, postaj))
		gbsp = g_brez(g, sporna)
		potke = []
		if bfs:
			potke += (list(itertools.islice(nx.all_simple_paths(gbsp, postaj+agent, 0), len(g))))
			potke += list(itertools.islice(nx.all_shortest_paths(gbsp, postaj+agent, 0), len(g)))
		else:
			potke += (list(itertools.islice(nx.all_simple_paths(gbsp, postaj+agent, 0), len(g))))
		potke = list(map(list, set(map(tuple, potke))))
		random.shuffle(potke)
		agentdo0.append(potke)
	# logging.debug(f"inicialne generated {agentdo0=}")
	assert len(agentdo0) == agentov
	pathsets = []
	for postaja in range(postaj-1):
		sporna = set(range(0, postaj))-set([postaja, postaja+1])
		gbsp = g_brez(g, sporna) # g brez spornih postaj
		if bfs:
			pathsets.append(list(itertools.islice(nx.all_simple_paths(gbsp, postaja, postaja+1), len(g))))
			pathsets[-1] += list(itertools.islice(nx.all_shortest_paths(gbsp, postaja, postaja+1), len(g)))
		else:
			pathsets.append(list(itertools.islice(nx.all_simple_paths(gbsp, postaja, postaja+1), len(g))))
		pathsets[-1] = list(map(list, set(map(tuple, pathsets[-1]))))
		random.shuffle(pathsets[-1])
	r = set()
	for i in range(len(g)):
		ret = []
		for agent in range(agentov):
			if agentov == 1:
				pot = [0]
			else:
				pot = random.choice(agentdo0[agent])[:]
			logging.debug(f"inicialne: pot samo agentdo0 je {pot=}")
			for postaja in range(postaj-1):
				dodaj = random.choice(pathsets[postaja])[1:]
				logging.debug(f"inicialne: dodajam poti {dodaj=}")
				pot += dodaj
			ret.append(tuple(pot))
			assert pot.count(postaj-1) == 1
		r.add(tuple(ret))
	vrni = list(map(list, r))
	# logging.debug(f"inicialne returning {vrni=}")
	return vrni
logging.info("generiram graf")
vozlišč = int(sys.argv[1])
verjetnost = float(sys.argv[2])
postaj = int(sys.argv[3])
agentov = int(sys.argv[4])
g = naključen(vozlišč, verjetnost)
logging.info("generiral graf!")
logging.info("generiram inicialne poti")
shuffled_paths = []
while len(shuffled_paths) < len(g):
	dolžinakromosoma = len(g)*(postaj)*agentov
	dodaj = poti2kromosom(inicialna(random.choice([True, False])), dolžinakromosoma)
	if len(dodaj) == dolžinakromosoma:
		shuffled_paths.append(dodaj)
	# logging.debug(f"adding to shuffled_paths: {len(shuffled_paths[-1])=}")
random.shuffle(shuffled_paths)
logging.info("generiral inicialne poti")
ga_instance = pygad.GA(num_generations=len(g),
	num_parents_mating=len(g),
	sol_per_pop=len(g)*2,
	fitness_func=fitness_function_factory(g),
	initial_population=np.array(shuffled_paths),
	crossover_type=crossover_func,
	mutation_type=mutation_func_factory(g),
	keep_elitism=len(g)//2,
	on_generation=on_generation)
ga_instance.run()
ga_instance.plot_fitness()
solution, solution_fitness, solution_idx = ga_instance.best_solution(ga_instance.last_generation_fitness)
print(f"Parameters of the best solution : {solution}")
print(f"Fitness value of the best solution = {solution_fitness}")
print(f"Index of the best solution : {solution_idx}")
if ga_instance.best_solution_generation != -1:
    print(f"Best fitness value reached after {ga_instance.best_solution_generation} generations.")
# ga_instance.save(filename='genetic')
# loaded_ga_instance = pygad.load(filename='genetic')
# loaded_ga_instance.plot_fitness()
if agentov == 1:
	pos = nx.spring_layout(g)
	barve = ["red"]*postaj + ["blue"]*(len(g)-postaj)
	nx.draw(g, pos, with_labels=True, node_color=barve)
	nx.draw_networkx_edge_labels(g, pos, edge_labels=nx.get_edge_attributes(g, 'weight'))
	pot = kromosom2poti(solution)[0]
	if len(g) <= 50:
		plt.show()
	plt.clf()
	nx.draw_networkx_nodes(g, pos, node_color=barve)
	nx.draw_networkx_labels(g, pos)
	nx.draw_networkx_edge_labels(g, pos, edge_labels=nx.get_edge_attributes(g, 'weight'))
	nx.draw_networkx_edges(g, pos=pos, edgelist=[(pot[x], pot[x+1]) for x in range(len(pot)-1)], edge_color="red")
	shortest = [0]
	for postaja in range(postaj-1):
		shortest += nx.shortest_path(g, postaja, postaja+1, weight="weight")[1:]
	nx.draw_networkx_edges(g, pos=pos, edgelist=[(shortest[x], shortest[x+1]) for x in range(len(shortest)-1)], edge_color="green")
	solution_len = nx.path_weight(g, pot, weight='weight')
	shortest_len = nx.path_weight(g, shortest, weight='weight')
	print(f"{solution_len=}, {shortest_len=}")
	print(f"{pot==shortest=}")
	if len(g) <= 100:
		plt.show()
else:
	# TODO vizualiziraj in ovrednoti rešitev za task3
	poti = kromosom2poti(solution)
	print("================ NAJBOLJŠA NAJDENA REŠITEV =====================")
	for agentidx, agentpot in enumerate(poti):
		print(f"AGENT {agentidx}\t")
		for vozl in agentpot:
			if vozl == -2:
				vozl == "wait"
			print(f"{vozl}, ", end="")
		print("")
	print(f"Vsi agenti bodo svoje delo zaključili znotraj {maxura} časovnih enot.")
	print(f"Napaka rešitve (prekrivanje) je {overlap} časovnih enot.")