from copy import deepcopy
import geopandas as gpd
from shapely.ops import transform
from shapely.geometry import MultiLineString, Point, LineString
from shapely.ops import linemerge
import numpy as np
import pandas as pd
from itertools import islice
from scipy.sparse.csgraph import dijkstra
import math
from osm_importer.graph_utils import (
	get_edge_path,
	get_nodes_degree,
	sparse_matrix,
	get_path,
	get_strongly_connected_component,
)
from typing import Any, Callable
import logging

log = logging.getLogger(__name__)


def process_list_in_col(col_values: Any, new_type: Any, function: Callable[[Any], Any] = np.nanmean) -> Any:
	"""
	ex: links['maxspeed'].apply(lambda x: process_list_in_col(x,float,np.nanmean)
	apply the function to the list. if its not a list. apply new type.
	"""
	if isinstance(col_values, list):
		return function([new_type(val) for val in col_values])
	else:
		return new_type(col_values)


def remove_list_in_col(col_values: Any, method: str = 'first'):
	"""
	remove list. keep first item if first, last if last. return any other value that are not a list
	"""
	if isinstance(col_values, list):
		if method == 'first':
			return col_values[0]
		else:
			return col_values[-1]
	else:
		return col_values


def get_epsg(lat: float, lon: float) -> int:
	"""
	return EPSG in meter for a given (lat,lon)
	lat is north south
	lon is est west
	"""

	return int(32700 - round((45 + lat) / 90, 0) * 100 + round((183 + lon) / 6, 0))


def get_intersections(gdf):
	count = {}
	for geometry in gdf['geometry'].values:
		for p in geometry.coords:
			count[p] = count.get(tuple(p), 0) + 1
	return {k for k, v in count.items() if v > 1}


def split_feature(geom, nodes=(), start=0):
	length = len(geom)
	if length == 2:
		return [geom]
	if length > 2:
		cut = 0
		for p in geom[start + 1 : length - 1]:
			cut += 1
			if tuple(p) in nodes:
				left = geom
				right = geom
				left = left[start : cut + 1]
				right = right[cut:]
				return [left] + split_feature(right, nodes=nodes, start=0)
	return [geom]


def get_split_features(geojson_dict):
	features = []
	intersections = get_intersections(geojson_dict)
	for feature in geojson_dict['features']:
		features += split_feature(feature, nodes=intersections)
	return features


def split_features(gdf):
	intersections = get_intersections(gdf)
	gdf['coords'] = gdf['geometry'].apply(lambda x: x.coords)
	df = gdf.drop(columns=['geometry'])
	df['coords'] = df['coords'].apply(lambda x: split_feature(x, intersections))
	df = df.explode('coords')
	geometry = df['coords'].apply(LineString)

	return gpd.GeoDataFrame(df.drop(columns=['coords']), geometry=geometry, crs=gdf.crs)


def split_directions(geojson_dict):
	features = geojson_dict['features']
	reversed_features = [deepcopy(f) for f in features if not f['properties']['oneway']]
	for f in reversed_features:
		f['geometry']['coordinates'] = list(reversed(f['geometry']['coordinates']))
	geojson_dict['features'] = features + reversed_features


def get_links_and_nodes(links: gpd.GeoDataFrame) -> tuple[gpd.GeoDataFrame, gpd.GeoDataFrame]:
	"""
	read overpass way.geojson file and return links and nodes as gpd.GeoDataFrames.
	Split direction dangerous here as Oneway tag is still rough (not just true / false...)

	parameters
	----------
	geojson_file (str) : path of the file
	split_direction (bool) : if we want to duplicate oneway=False and reverse the geom.

	returns
	----------
	links, nodes
	"""
	links = split_features(links)

	# get nodes
	links['a'] = links['geometry'].apply(lambda x: x.coords[0])
	links['b'] = links['geometry'].apply(lambda x: x.coords[-1])

	node_coordinates = set(links['a']).union(links['b'])
	node_index = {coord: i for i, coord in enumerate(node_coordinates)}
	nodes = gpd.GeoDataFrame(geometry=[Point(pt) for pt in node_coordinates], crs=links.crs)
	nodes.index.name = 'index'

	links['a'] = links['a'].apply(node_index.get)
	links['b'] = links['b'].apply(node_index.get)

	if links.index.name == 'id':
		links = links.reset_index()
	else:
		links = links.reset_index(drop=True)
	links = links.rename(columns={'id': 'osmid'})
	links.index.name = 'index'

	return links, nodes


def main_strongly_connected_component(
	links: gpd.GeoDataFrame, nodes: gpd.GeoDataFrame = None, split_direction=True
) -> gpd.GeoDataFrame:
	"""
	use split direction if the links are not duplicated (if there is 1 link for oneway=False)
	remove cul-de-sac
	"""
	edges = links[['a', 'b']].values
	if 'oneway' in links.columns and split_direction:
		reversed_edges = links.loc[~links['oneway'].astype(bool)][['b', 'a']].values.tolist()
		edges = np.concat([edges, reversed_edges])
	main_scc = set(get_strongly_connected_component(edges))
	_links = links[links['a'].isin(main_scc) & links['b'].isin(main_scc)]
	if nodes is not None:
		_nodes = nodes.loc[list(main_scc)]
		return _links, _nodes
	return _links


def reverse_geom(geom):
	def _reverse(x, y, z=None):
		if z:
			return x[::-1], y[::-1], z[::-1]
		return x[::-1], y[::-1]

	return transform(_reverse, geom)


def rectify_geometry_direction(links, nodes):
	node_dict = nodes['geometry'].to_dict()
	links['node_geom'] = links['a'].apply(lambda x: node_dict.get(x))
	sliced = links['geometry'].apply(lambda x: x.coords[0]) != links['node_geom'].apply(lambda x: x.coords[0])
	log.info(len(links[sliced]), 'geometry to inverse')
	links.loc[sliced, 'geometry'] = links.loc[sliced]['geometry'].apply(lambda x: reverse_geom(x))
	links = links.drop(columns='node_geom')
	return links


def batched(iterable, n):
	"Batch data into lists of length n. The last batch may be shorter."
	# batched('ABCDEFG', 3) --> ABC DEF G
	it = iter(iterable)
	while True:
		batch = list(islice(it, n))
		if not batch:
			return
		yield batch


def cut(line, distance, normalized=True):
	if normalized:
		distance *= line.length

	# Cuts a line in two at a distance from its starting point
	if distance <= 0.0:
		return [None, LineString(line)]
	if distance >= line.length:
		return [LineString(line), None]
	coords = list(line.coords)
	pd = 0
	for i, p in enumerate(coords):
		if i == 0:
			continue
		pd += math.dist(p, coords[i - 1])
		if pd == distance:
			return [LineString(coords[: i + 1]), LineString(coords[i:])]
		if pd > distance:
			cp = line.interpolate(distance)
			return [LineString(coords[:i] + [(cp.x, cp.y)]), LineString([(cp.x, cp.y)] + coords[i:])]


def simplify(links, cutoff=10):
	# create a graph and find all nodes with deg >2 (sources)
	deg_dict = get_nodes_degree(links[['a', 'b']].values)
	sources = [node for node, degree in deg_dict.items() if degree > 2]
	links['weight'] = 1
	log.info(len(sources), 'deg 2 nodes')
	path_to_merge = []
	# graph on oneway and highway as unique as we dont want to aggregate highway together or one way
	for col1, col2 in [(a, b) for a in links['oneway'].unique() for b in links['highway'].unique()]:
		filtered_links = links[(links['oneway'] == col1) & (links['highway'] == col2)]
		if len(filtered_links) < 2:
			continue
		nodes_set = set(filtered_links['a']).union(set(filtered_links['b']))
		mat, node_index = sparse_matrix(filtered_links[['a', 'b', 'weight']].values)
		sparse_deg_dict = {node_index[key]: val for key, val in deg_dict.items() if key in nodes_set}
		index_node = {v: k for k, v in node_index.items()}
		filtered_sources = [s for s in sources if s in nodes_set]
		unfounds_origins = []

		# 1) for each nodes with deg > 2
		# 2) take all non empty path
		# 3) take all path with destination a node with deg > 2
		#    and only deg 2 nodes in between ex: [3,2,2,4]
		# 4) take all path with destination a node of deg == 2
		#    if the path length is the cutoff.

		# batch the dijkstra, my computer was crashing with ~40 000 sources
		# it it not that much slower.
		for origins in batched(filtered_sources, 1000):
			origin_sparse = [node_index[x] for x in origins]
			dist_matrix, predecessors = dijkstra(
				csgraph=mat, directed=True, indices=origin_sparse, return_predecessors=True, limit=cutoff
			)

			# from all source (source are deg>2 by construction)
			for oi in range(len(origins)):
				# destinations are non inf paths
				destinations = np.where(dist_matrix[oi, :] != np.inf)[0]
				# for each destination with deg >= 2
				for di in (d for d in destinations if sparse_deg_dict[d] >= 2):
					path = get_path(predecessors, oi, di)  # get path, with origin and destinations in the list
					path_deg = [*map(sparse_deg_dict.get, path)]  # get deg of each nodes in path
					# if destination is not 2. keep path if every nodes in between are 2
					if path_deg[-1] != 2:
						if set(path_deg[1:-1]) == {2}:
							path_to_merge.append([*map(index_node.get, path)])
							if len(path_to_merge) > 1:
								if path_to_merge[-1] == path_to_merge[-2]:
									log.info(path)
					# if destination == 2
					elif len(path) == cutoff + 1:
						# keep path if len(path) == cutoff+1. at the end redo on those with large cutoff
						if set(path_deg[1:]) == {2}:
							unfounds_origins.append(origins[oi])

		# redo with large cutoff the path thats were longer than the cutoff
		# ex:cutoff = 5, path = [3,2,2,2,2]. the real path may be [3,2,2,2,2,2,2,2,4]. and we will find it now.
		if len(unfounds_origins) > 0:
			log.info('find path with large cutoff for', len(unfounds_origins), ' origins')
			origin_sparse = [node_index[x] for x in unfounds_origins]
			dist_matrix, predecessors = dijkstra(
				csgraph=mat, directed=True, indices=origin_sparse, return_predecessors=True, limit=100
			)
			for oi in range(len(unfounds_origins)):
				# destinations are non inf paths
				destinations = np.where(dist_matrix[oi, :] != np.inf)[0]
				# for each destination with deg >= 2
				for di in (d for d in destinations if sparse_deg_dict[d] >= 2):
					path = get_path(predecessors, oi, di)  # get path, with origin and destinations in the list
					path_deg = [*map(sparse_deg_dict.get, path)]  # get deg of each nodes in path
					# if destination is not 2. keep path if every nodes in between are 2
					# those were not find in cutoff, only check for larger cutoff.
					if (path_deg[-1] != 2) & (len(path) >= cutoff + 1):
						if set(path_deg[1:-1]) == {2}:
							path_to_merge.append([*map(index_node.get, path)])

	# transform sparse nodes path to links path ([rlinks_2, rlinks_40, ...])
	links_dict = links.reset_index().set_index(['a', 'b'])['index'].to_dict()
	links_paths = [[*map(links_dict.get, get_edge_path(path))] for path in path_to_merge]

	# apply a group number to each sequence to merge. group 0 is nothing to merge.
	links['group'] = 0
	idx = 1
	for path in links_paths:
		links.loc[path, 'group'] = idx
		idx += 1
	tlinks = links[links['group'] > 0].copy()
	# sort tlinks links to merge in the right order
	flat_index = [l for sublist in links_paths for l in sublist]
	tlinks = tlinks.loc[flat_index]

	def list_or_first(x: pd.Series):
		return list(x) if x.nunique() > 1 else x.iloc[0]

	def merge_lines(x):
		return linemerge(MultiLineString(list(x)))

	# merge links
	agg_dict = {col: list_or_first for col in tlinks.columns}
	agg_dict['geometry'] = merge_lines
	agg_dict['a'] = 'first'
	agg_dict['b'] = 'last'
	agg_dict['index'] = 'first'
	tlinks = tlinks.reset_index()

	tlinks = tlinks.groupby('group').agg(agg_dict)
	tlinks = tlinks.set_index('index')

	# remove agg that change oneway
	log.info(
		len(tlinks[tlinks['oneway'].apply(lambda x: type(x) == list)]),
		'links were not merge because the oneway field is not the same',
	)
	tlinks = tlinks[~tlinks['oneway'].apply(lambda x: type(x) == list)]

	log.info(
		len(tlinks[tlinks['highway'].apply(lambda x: type(x) == list)]),
		'links were not merge because the highway field is not the same',
	)
	tlinks = tlinks[~tlinks['highway'].apply(lambda x: type(x) == list)]

	# TODO: on a des multiLinestring. il faudrait mieux merger.
	log.info(
		len(tlinks[tlinks['geometry'].apply(lambda x: x.geom_type != 'LineString')]),
		'merged_links unmerged because the geometry became a multilinestring',
	)
	tlinks = tlinks[tlinks['geometry'].apply(lambda x: x.geom_type == 'LineString')]

	# merge tlinks back into links
	tlinks = gpd.GeoDataFrame(tlinks, crs=links.crs)
	links = links[~links['group'].isin(tlinks['group'].unique())]
	links = pd.concat([links, tlinks])
	links = links.drop(columns=['weight', 'group'])

	return links


def split_oneway(links: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
	"""
	duplicated all links with oneway == False. reverse there geometry, reverse a,b and rename with _r
	"""
	nlinks = links[~links['oneway']].copy()
	nlinks.index = nlinks.index.astype(str)  # index must be string
	nlinks.index = nlinks.index + '_r'
	nlinks['geometry'] = nlinks['geometry'].apply(lambda x: reverse_geom(x))
	nlinks = nlinks.rename(columns={'a': 'b', 'b': 'a'})
	links = pd.concat([links, nlinks])
	links['oneway'] = True
	return links


def split_duplicated_links(
	links: gpd.GeoDataFrame, nodes: gpd.GeoDataFrame
) -> tuple[gpd.GeoDataFrame, gpd.GeoDataFrame]:
	crs = nodes.crs
	# index must be string
	links.index = links.index.astype(str)
	nodes.index = nodes.index.astype(str)
	links['a'] = links['a'].apply(str)
	links['b'] = links['b'].apply(str)

	links = links.reset_index()
	duplicated_links = links[links.duplicated(subset=['a', 'b'])].copy()
	log.info(f'Splitting {len(duplicated_links)} duplicated links')
	duplicated_links['geometries'] = duplicated_links['geometry'].apply(cut, distance=0.5)
	duplicated_links['index'] = duplicated_links['index'].apply(lambda x: [x + '_a', x + '_b'])
	middle_nodes = duplicated_links['geometries'].apply(lambda g: Point(g[0].coords[-1]))
	middle_nodes.index = duplicated_links['a'] + '_m'
	middle_nodes.name = 'geometry'

	duplicated_links['middle_node'] = middle_nodes.index
	duplicated_links['a'] = duplicated_links[['a', 'middle_node']].values.tolist()
	duplicated_links['b'] = duplicated_links[['middle_node', 'b']].values.tolist()

	cols = ['index', 'a', 'b', 'geometries']
	other_cols = [c for c in duplicated_links.columns if c not in cols]
	duplicated_links = duplicated_links.set_index(other_cols).apply(pd.Series.explode).reset_index()
	duplicated_links = duplicated_links.drop(columns=['geometry', 'middle_node']).rename(
		columns={'geometries': 'geometry'}
	)

	links = pd.concat([links[~links.duplicated(subset=['a', 'b'])].copy(), duplicated_links])
	links = links.set_crs(crs, allow_override=True)
	nodes = pd.concat([nodes, middle_nodes])
	nodes = nodes.set_crs(crs, allow_override=True)

	return links.set_index('index'), gpd.GeoDataFrame(nodes, crs=crs)


def drop_duplicated_links(
	links: gpd.GeoDataFrame, sort_column: str = 'maxspeed', ascending: bool = False
) -> gpd.GeoDataFrame:
	"""
	drop duplicated links (a,b) with condition sort_column. if maxspeed and ascending=False, keep faster road
	"""
	before = len(links)
	links = links.sort_values(sort_column, ascending=ascending).drop_duplicates(subset=['a', 'b']).sort_index()
	log.info(before - len(links), 'links dropped')
	return links


def clean_maxspeed(links: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
	"""
	remove kph, remove and convert mph. transport everything to float. NaN if string (ex: walk)
	"""
	links['maxspeed'] = links['maxspeed'].str.lower().str.replace('kph', '')
	try:
		mph_index = links['maxspeed'].astype(str).str.lower().str.contains('mph')
		links.loc[mph_index, 'maxspeed'] = (
			links.loc[mph_index, 'maxspeed'].str.lower().str.replace('mph', '').astype('float') * 1.60934
		)
	except:
		log.info('fail to convert mph in maxspeed to float ')
	# this remove all strings lefts
	links['maxspeed'] = links['maxspeed'].apply(pd.to_numeric, errors='coerce')
	return links


def clean_oneway(links: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
	"""
	clean oneways. return True or False
	"""
	# only yes is oneway. all other tags set to false  'no' '-1', 'alternating' 'reversible'
	oneway_dict = {'yes': True}
	links['oneway'] = links['oneway'].apply(lambda x: oneway_dict.get(x, False))
	if len(links['oneway'].unique()) > 2:
		log.info('WARNING: some oneway tags are not defined', links.unique())
	return links


def clean_lanes(links: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
	"""
	remove semicolon (ex lanes = 1;2) and average the result (1;2 -> 1.5).
	transform everything to float. NaN if string
	"""
	semicolon_index = links['lanes'].astype(str).str.contains(';')
	try:
		links.loc[semicolon_index, 'lanes'] = links.loc[semicolon_index, 'lanes'].apply(
			lambda x: np.mean([float(el) for el in x.split(';')])
		)
	except:
		pass
	links['lanes'] = links['lanes'].apply(pd.to_numeric, errors='coerce')
	return links


def fill_na_col(
	links: gpd.GeoDataFrame, group: str = 'highway', col: str = 'maxspeed', agg_func: Callable[[list], Any] = np.mean
) -> gpd.GeoDataFrame:
	"""
	fill the nan value with the average of its group (if np.mean)
	for example: with group=highway, col=maxspeed and agg_func = np.mean. the mean maxpeed per highway is compute.
	then its apply to every NaN value of each highway.
	Note that a group with only NaN will stay a nan (no agg_func will be applied.)

	parameters
	----------
	links: gpd.GeoDataFrame = links to apply the function
	group: group to do the groupby and find the average (or any other function)
	col: column to fill
	agg_func : function to apply

	returns
	----------
	{index: elevation} : Dict[Any,int] = elevataion for each index in the GeoDataFrame
	"""

	val_dict = links[[group, col]].dropna().groupby(group)[col].agg(agg_func).to_dict()
	links.loc[~np.isfinite(links[col]), col] = links.loc[~np.isfinite(links[col]), group].apply(
		lambda x: val_dict.get(x)
	)
	return links


def get_columns_with_list(df: pd.DataFrame) -> list:
	"""
	return the list of column that contain list in them.
	"""

	def _contain_list(df, col):
		for val in df[col].values:
			if type(val) == list:
				return True
		else:
			return False

	object_columns = df.columns[(df.dtypes == object).values]
	res = []
	for col in object_columns:
		if _contain_list(df, col):
			res.append(col)
	return res
