# Copyright © 2024, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE.md
import hashlib
import logging
from copy import deepcopy
from math import ceil
from os.path import join
from tempfile import gettempdir

import numpy as np
import pandas as pd
import shapely.wkt
from shapely.geometry import LineString
from shapely.geometry import Point
from shapely.ops import substring

from polaris.network.tools.geo_index import GeoIndex
from polaris.network.transit.transit_elements.mode_correspondence import mode_correspondence
from polaris.utils.database.db_utils import read_and_close
from polaris.utils.optional_deps import check_dependency
from polaris.utils.signals import SIGNAL

GRAPH_VERSION = 1
CONNECTOR_SPEED = 1


class MMGraph:
    """Build specialized map-matching graphs. Not designed to be used by the final user"""

    signal = SIGNAL(object)

    def __init__(self, geotool, stops, agency, mtmm):
        self.geotool = geotool
        self.stops = stops
        self.__mtmm = mtmm
        self._idx = None
        self.max_link_id = -1
        self.max_node_id = -1
        self.mode = ""
        self.modename = ""
        self.mode_id = -1
        self.__mode = ""
        self.__df_file = ""
        self.__agency = agency
        self.__centroids_file = ""
        self.__mm_graph_file = ""
        self.node_corresp = []
        self.__all_links = {}  # type: Dict[pd.Series]
        self.distance_to_project = -1
        self.df = pd.DataFrame([])
        self.logger = logging.getLogger("polaris")
        self._path_to_file = geotool._network_file

    def build_graph_with_broken_stops(self, mode_id: int, distance_to_project=50):
        """Build the graph for links for a certain mode while splitting the closest links at stops' projection

        Args:
            *mode_id* (:obj:`int`): Mode ID for which we will build the graph for
            *distance_to_project* (:obj:`float`, **Optional**): Radius search for links to break at the stops. Defaults to 50m
        """
        self.logger.debug(f"Called build_graph_with_broken_stops for mode_id={mode_id}")
        self.mode_id = mode_id
        self.distance_to_project = distance_to_project
        self.__mode = mode_correspondence[self.mode_id]  # type: ignore
        self.__mm_graph_file = join(gettempdir(), f"map_matching_graph_{self.__agency}_{self.__mode}.csv")
        modename = self.geotool._mode_name(self.__mode)

        with read_and_close(self._path_to_file, spatial=True) as conn:
            # When it is rail we don't care about the number of lanes to define link directionality
            extra_clause = " " if self.__mode == "RAIL" else "AND lanes_ab+lanes_ba>0"
            get_qry = f"""Select link link_id, 0 direction, node_a a_node, node_b b_node, max(fspd_ab, fspd_ba) speed,
                                 "length" distance, lanes_ab, lanes_ba, ST_AsText(geo) wkt from Link
                          INNER JOIN Link_Type ON Link.type = Link_Type.link_type
                          WHERE INSTR(Link_Type.use_codes, "{self.__mode}")>0 {extra_clause}"""

            self.geotool.logger.debug("  Reading links from physical network")
            self.df = pd.read_sql(get_qry, conn)
            if self.__mode != "RAIL":
                self.df.loc[self.df.lanes_ab == 0, "direction"] = -1
                self.df.loc[self.df.lanes_ba == 0, "direction"] = 1
            self.df.drop(columns=["lanes_ab", "lanes_ba"], inplace=True)

        if not self.df.shape[0]:
            check_dependency("aequilibrae")
            from aequilibrae.paths import Graph

            return Graph()
        if modename not in self.geotool.mode_link_idx:
            self.geotool.build_mode_link_index(self.__mode)

        # We do some wrangling to save the graph to disk, in case we need to run this more than once
        authvalue = hashlib.md5()
        authvalue.update(str(distance_to_project).encode())
        authvalue.update(str(GRAPH_VERSION).encode())
        authvalue.update("".join([str(x) for x in self.stops.keys()]).encode())
        authvalue.update("".join([str(x) for x in self.df.link_id.values]).encode())
        authvalue.update(modename.encode())
        graph_hash = authvalue.hexdigest()

        self.__df_file = join(gettempdir(), f"{graph_hash}_df.zip")
        self.__centroids_file = join(gettempdir(), f"{graph_hash}_centroids.zip")
        # if isfile(self.__df_file) and isfile(self.__centroids_file):
        #     return self.__build_graph_from_cache()

        return self.__build_graph_from_scratch()

    def __build_graph_from_cache(self):
        self.logger.info(f"Loading map-matching graph from disk for mode_id={self.mode_id}")
        net = pd.read_csv(self.__df_file)
        centroid_corresp = pd.read_csv(self.__centroids_file)
        centroids = np.copy(centroid_corresp.centroid_id.values)
        centroid_corresp.set_index("node_id", inplace=True)
        for stop in self.stops.values():
            stop.___map_matching_id__[self.mode_id] = centroid_corresp.loc[stop.stop_id, "centroid_id"]
        return self.__graph_from_broken_net(centroids, net)

    def __build_graph_from_scratch(self):
        self.logger.info(f"Creating map-matching graph from scratch for mode_id={self.mode_id}")
        self.df = self.df.assign(original_id=self.df.link_id, is_connector=0, geo=self.df.wkt.apply(shapely.wkt.loads))
        self.max_link_id = self.df.link_id.max() + 1
        self.max_node_id = self.df[["a_node", "b_node"]].max().max() + 1
        # Build initial index
        self.signal.emit(["start", "secondary", self.df.shape[0], f"Indexing links - {self.__mode}", self.__mtmm])
        self._idx = GeoIndex()
        for counter, (_, record) in enumerate(self.df.iterrows()):
            self.signal.emit(["update", "secondary", counter + 1, f"Indexing links - {self.__mode}", self.__mtmm])
            self._idx.insert(feature_id=record.link_id, geometry=record.geo)
        # We will progressively break links at stops' projection
        # But only on the right side of the link (no boarding at the opposing link's side)
        centroids = []
        self.node_corresp = []
        self.signal.emit(["start", "secondary", len(self.stops), f"Breaking links - {self.__mode}", self.__mtmm])
        self.df = self.df.assign(free_flow_time=np.inf, wrong_side=0, closest=1, to_remove=0)
        self.__all_links = {rec.link_id: rec for _, rec in self.df.iterrows()}
        for counter, (stop_id, stop) in enumerate(self.stops.items()):
            self.signal.emit(["update", "secondary", counter + 1, f"Breaking links - {self.__mode}", self.__mtmm])
            stop.___map_matching_id__[self.mode_id] = self.max_node_id
            self.node_corresp.append([stop_id, self.max_node_id])
            centroids.append(stop.___map_matching_id__[self.mode_id])
            self.max_node_id += 1
            self.connect_node(stop)
        self.df = pd.concat([pd.DataFrame(rec).transpose() for rec in self.__all_links.values()])

        self.df = self.df[self.df.to_remove == 0]
        fltr = self.df.speed > 0
        self.df.loc[fltr, "free_flow_time"] = self.df.distance[fltr] / self.df.speed[fltr]

        cols = [
            "link_id",
            "a_node",
            "b_node",
            "direction",
            "free_flow_time",
            "distance",
            "is_connector",
            "closest",
            "original_id",
        ]
        net = self.df[cols]
        # Caches the graph outputs
        net.reset_index(inplace=True, drop=True)
        net = net.drop_duplicates(subset=["a_node", "b_node"])
        net.to_csv(self.__df_file, index=False)
        cols.append("wkt")
        self.df[cols].to_csv(self.__mm_graph_file, index=False)
        pd.DataFrame(self.node_corresp, columns=["node_id", "centroid_id"]).to_csv(self.__centroids_file)
        return self.__graph_from_broken_net(centroids, net)

    def connect_node(self, stop) -> None:
        list_nearest = list(self._idx.nearest(stop.geo, 30))

        is_closest = 1
        conn_found = 0
        distances = [stop.geo.distance(self.__all_links[lid].geo) for lid in list_nearest]

        # Sort by distance to the stop
        nearest = pd.DataFrame({"dist": distances, "link_id": list_nearest}).sort_values(by="dist")
        criterium = min(5 * nearest.dist.min(), self.distance_to_project)
        criterium = criterium if nearest.dist.min() < 250 else ceil(nearest.dist.min() * 100) / 100
        nearest = nearest[nearest.dist <= max(criterium, nearest.dist.min())].link_id.tolist()

        for link_id in nearest:
            link = self.__all_links[link_id]

            link_geo = link.geo

            # We disregard links beyond the threshold, but maintain the closest link to ensure connectivity

            if link_geo.boundary.is_empty:
                first = Point([link_geo.coords.xy[0][0], link_geo.coords.xy[1][0]])
                last = Point([link_geo.coords.xy[0][-1], link_geo.coords.xy[1][-1]])
            else:
                first = link_geo.boundary.geoms[0]
                last = link_geo.boundary.geoms[1]

            proj_point = link_geo.project(stop.geo)
            break_point = link_geo.interpolate(proj_point)

            if proj_point <= 1.0:  # If within one meter of the end of the link, let's go with the existing node
                break_point = first
                intersec_node = link.a_node
            elif proj_point >= link_geo.length - 1.0:
                break_point = last
                intersec_node = link.b_node
            else:
                link.to_remove = 1
                self._idx.delete(link_id, link_geo)
                intersec_node = self.max_node_id
                self.max_node_id += 1

                # Create the first portion of the link
                fp = deepcopy(link)
                fp.link_id = self.max_link_id
                fp.b_node = intersec_node
                fp.to_remove = 0
                fp.geo = substring(link_geo, 0, proj_point)
                fp.wkt = fp.geo.wkt
                fp.distance = fp.geo.length
                self._idx.insert(fp.link_id, fp.geo)
                self.max_link_id += 1
                self.__all_links[fp.link_id] = fp

                # Create the second portion of the link
                lp = deepcopy(link)
                lp.link_id = self.max_link_id
                lp.a_node = intersec_node
                lp.geo = substring(link_geo, proj_point, link_geo.length)
                lp.wkt = lp.geo.wkt
                lp.distance = lp.geo.length
                lp.to_remove = 0
                self._idx.insert(lp.link_id, lp.geo)
                self.max_link_id += 1
                self.__all_links[lp.link_id] = lp

            # Add the connector
            connector_geo = LineString([stop.geo, break_point])
            connector = deepcopy(link)
            connector.link_id = self.max_link_id
            connector.original_id = -1
            connector.a_node = stop.___map_matching_id__[self.mode_id]
            connector.b_node = intersec_node
            connector.direction = 0
            connector.to_remove = 0
            connector.geo = connector_geo
            connector.wkt = connector_geo.wkt
            connector.distance = connector_geo.length
            connector.is_connector = 1  # We make sure that the closest connector cannot be turned off
            connector.closest = is_closest
            self.max_link_id += 1
            is_closest = 0
            conn_found += 1
            self.__all_links[connector.link_id] = connector

    def __graph_from_broken_net(self, centroids, net):
        check_dependency("aequilibrae")
        from aequilibrae.paths import Graph

        net_data = pd.DataFrame(
            {
                "distance": net.distance.astype(float),
                "direction": net.direction.astype(np.int8),
                "a_node": net.a_node.astype(int),
                "b_node": net.b_node.astype(int),
                "link_id": net.link_id.astype(int),
                "is_connector": net.is_connector.astype(int),
                "original_id": net.original_id.astype(int),
                "closest": net.closest.astype(int),
                "free_flow_time": net.free_flow_time.astype(float),
            }
        )

        fltr = net.is_connector == 1
        net_data.loc[fltr, "distance"] = 1.2 * (net_data[fltr].distance ** 1.3)  # Already penalize it a bit
        net_data.loc[fltr, "free_flow_time"] = ((net_data[fltr].distance / 1000) / CONNECTOR_SPEED) * 60

        g = Graph()
        g.network_ok = True
        g.status = "OK"
        g.network = net_data
        g.prepare_graph(np.array(centroids))
        g.set_graph("free_flow_time")
        g.set_skimming(["distance"])
        g.set_blocked_centroid_flows(True)
        return g

    def finished(self):
        self.signal.emit(["finished_building_mm_graph_procedure"])
