from enum import Enum
import random, time
from typing import Dict, Optional, Union
from feasytools import ReadOnlyTable, CDDiscrete, PDDiscrete, PDGamma, DTypeEnum
from ..locale import Lang
from ..traffic import EV, EVDict, RoadNet, DetectFiles
from .misc import VehicleType, random_diff, _TripInner, _EVInner, _xmlSaver
from .poly import PolygonMan

DictPDF = Dict[int, Union[PDDiscrete[int], None]]

TAZ_TYPE_LIST = ("Home", "Work", "Relax", "Other")

class RoutingCacheMode(Enum):
    """Routing cache mode"""
    NONE = 0  # No cache
    RUNTIME = 1 # Cache during runtime
    STATIC = 2 # Static cache in generation time

    def __str__(self):
        return ("None", "Runtime", "Static")[self.value]

    def __repr__(self):
        return self.value
    
class TripsGenMode(Enum):
    """Generation mode"""

    AUTO = "Auto"  # Automatic
    TYPE = "TYPE"  # Type-based
    POLY = "Poly"  # Polygon-based

    def __str__(self):
        return self.value

    def __repr__(self):
        return self.value

class EVsGenerator:    
    """Class to generate trips"""
    def __init__(self, CROOT: str, PNAME: str, seed,
            mode: TripsGenMode = TripsGenMode.AUTO,
            route_cache: RoutingCacheMode = RoutingCacheMode.NONE):
        """
        Initialization
            CROOT: Trip parameter folder
            PNAME: SUMO configuration folder
            seed: Random seed
        """
        _fn = DetectFiles(PNAME)
        random.seed(seed)
        self.vTypes = [VehicleType(**x) for x in ReadOnlyTable(CROOT + "/ev_types.csv",dtype=DTypeEnum.FLOAT32).to_list_of_dict()]
        # Define various functional area types
        self._route_cache_mode = route_cache
        self.net = RoadNet.load(_fn["net"])
        if mode == TripsGenMode.AUTO:
            if _fn.node_type: mode = TripsGenMode.TYPE
            elif _fn.poly and _fn.net and _fn.fcs: mode = TripsGenMode.POLY
            else: raise RuntimeError(Lang.ERROR_NO_TAZ_OR_POLY)
        if mode == TripsGenMode.TYPE:
            assert _fn.node_type, Lang.ERROR_NO_TAZ_OR_POLY
            self._mode = "type"
            self.dic_nodetype = {}
            with open(_fn.node_type, "r") as fp:
                for ln in fp.readlines():
                    name, lst = ln.split(":")
                    self.dic_nodetype[name.strip()] = [x.strip() for x in lst.split(",")]
        elif mode == TripsGenMode.POLY:
            assert _fn.poly and _fn.net and _fn.fcs, Lang.ERROR_NO_TAZ_OR_POLY
            self._mode = "poly"
            polys = PolygonMan(_fn.poly)
            self.dic_nodetype = {k:[] for k in TAZ_TYPE_LIST}
            for poly in polys:
                poly_type = poly.getConvertedType()
                if poly_type:
                    dist, node = self.net.find_nearest_node_with_distance(*poly.center())
                    # Ensure the edge is in the largest strongly connected component
                    if dist < 200 and self.net.is_node_in_largest_scc(node.id):
                        self.dic_nodetype[poly_type].append(node.id)
        else:
            raise RuntimeError(Lang.ERROR_NO_TAZ_OR_POLY)
        
        # Start time of first trip
        self.pdf_start_weekday = PDGamma(6.63, 65.76, 114.54)
        self.pdf_start_weekend = PDGamma(3.45, 84.37, 197.53)
        
        # Spatial transfer probability of weekday and weekend. 
        # key1 = from_type, key2 = time (0~95, each unit = 15min), value = CDF of (to_type1, to_type2, to_type3, to_type4)
        self.PSweekday:Dict[str, DictPDF] = {}
        self.PSweekend:Dict[str, DictPDF] = {}
        # Parking duration CDF of weekday and weekend.
        self.park_cdf_wd:Dict[str, CDDiscrete[int]] = {} 
        self.park_cdf_we:Dict[str, CDDiscrete[int]] = {}

        def read_trans_pdfs(path:str) -> DictPDF:
            tbwd = ReadOnlyTable(path, dtype=DTypeEnum.FLOAT32)
            times = [int(x) for x in tbwd.head[1:]]
            values = list(map(int, tbwd.col(0)))
            ret:DictPDF = {}
            for i in range(1, len(times)+1):
                weights = list(map(float, tbwd.col(i)))
                assert len(values) == len(weights)
                try:
                    ret[i] = PDDiscrete(values, weights)
                except ZeroDivisionError:
                    ret[i] = None
            return ret
        
        for dtype in TAZ_TYPE_LIST:
            self.PSweekday[dtype] = read_trans_pdfs(f"{CROOT}/space_transfer_probability/{dtype[0]}_spr_weekday.csv")
            self.PSweekend[dtype] = read_trans_pdfs(f"{CROOT}/space_transfer_probability/{dtype[0]}_spr_weekend.csv")
            self.park_cdf_wd[dtype] = CDDiscrete(f"{CROOT}/duration_of_parking/{dtype[0]}_spr_weekday.csv", True, int)
            self.park_cdf_we[dtype] = CDDiscrete(f"{CROOT}/duration_of_parking/{dtype[0]}_spr_weekend.csv", True, int)

        self.soc_pdf = PDDiscrete.fromCSVFileI(f"{CROOT}/soc_dist.csv", True)
    
    def __getPs(self, is_weekday: bool, dtype: str, time_index:int):
        return self.PSweekday[dtype].get(time_index, None) if is_weekday else self.PSweekend[dtype].get(time_index, None)
    
    def __getDest1(self, pfr: str, weekday: bool = True):
        """
        Get the destination of the trip secondary to the first trip
            pfr: Departure functional area type, such as "Home"
            weekday: Whether it is weekday or weekend
        Returns: 
            First trip: First departure time, arrival destination functional area type, such as "Work"
        """
        pdf = None
        while pdf is None:
            init_time = self.pdf_start_weekday.sample() if weekday else self.pdf_start_weekend.sample()
            if init_time >= 86400:
                continue
            # Time index (0~95, each unit = 15min)
            init_time_i = int(init_time / 15)
            pdf = self.__getPs(weekday, pfr, init_time_i)
        next_place = TAZ_TYPE_LIST[pdf.sample()]
        return int(init_time), next_place

    def __getDestA(self, from_type:str, init_time_i:int, weekday: bool):
        """
        Get the destination of the next trip for non-first trips
            from_type: Departure type, such as "Home"
            init_time_i: Time index (0~95, each unit = 15min)
            weekday: Whether it is weekday or weekend
        Returns:
            Destination type
        """
        cdf = self.__getPs(weekday, from_type, init_time_i)
        return "Home" if cdf is None else TAZ_TYPE_LIST[cdf.sample()]

    def __getNextNode(self, from_node:str, next_place_type:str) -> str:
        return random_diff(self.dic_nodetype[next_place_type], from_node)        
    
    def __genFirstTrip1(self, trip_id, weekday: bool = True):
        """
        Generate the first trip of the first day
            trip_id: Trip ID
            weekday: hether it is weekday or weekend
        Return a InnerTrip instance
        """
        from_Type = "Home"
        from_node = random.choice(self.dic_nodetype[from_Type])
        # Get departure time and destination area type
        depart_time_min, to_Type = self.__getDest1(from_Type, weekday)  # Minimum departure time is 0
        to_node = self.__getNextNode(from_node, to_Type)
        return _TripInner(trip_id, depart_time_min * 60, from_node, from_Type, to_node, to_Type)

    cdf_dict = {}

    def __genStopTimeIdx(self, from_type:str, weekday: bool):
        cdf = self.park_cdf_wd[from_type] if weekday else self.park_cdf_we[from_type]
        return int(cdf.sample() + 1)

    def __genTripA(
        self, trip_id:str, from_type:str, from_node:str, start_time:int, weekday: bool = True
    ) -> _TripInner:
        """
        Generate the second trip
            trip_id: Trip ID
            from_type: Departure area type, such as "Home"
            from_node: Departure roadside, such as "T6"
            start_time: Departure time of the first trip, in seconds since midnight
            weekday: Whether it is weekday or weekend
        """
        depart_time_min = 1440
        cnt = 0
        while depart_time_min >= 1440:  # If the departure time is after midnight, regenerate
            stop_time_idx = self.__genStopTimeIdx(from_type, weekday)
            depart_time_min = start_time // 60 + stop_time_idx * 15 + 20
            cnt += 1
            if cnt > 10:
                depart_time_min = start_time // 60 + 1
                break
        next_place_type = self.__getDestA(from_type, stop_time_idx, weekday)
        to_node = self.__getNextNode(from_node, next_place_type)
        return _TripInner(trip_id, depart_time_min * 60, from_node, from_type, to_node, next_place_type)

    def __genTripF(
        self, trip_id:str, from_type:str, from_node:str,
        start_time:int, first_node:str, weekday: bool = True,
    ):
        """
        Generate the third trip
            trip_id: Trip ID
            from_TAZ: Departure area TAZ type, such as "TAZ1"
            from_type: Departure area type, such as "Home"
            from_EDGE: Departure roadside, such as "gnE29"
            start_time: Departure time of the first trip, in seconds since midnight
            first_TAZ: First trip's destination area TAZ type, such as "TAZ2"
            first_EDGE: First trip's destination roadside, such as "gnE2"
            weekday: Whether it is weekday or weekend
        """
        if first_node == from_node:
            return None
        depart_time_min = 1440
        cnt = 0
        while depart_time_min >= 1440:  # If the departure time is after midnight, regenerate
            stop_time_idx = self.__genStopTimeIdx(from_type, weekday)
            depart_time_min = start_time // 60 + stop_time_idx * 15 + 20
            cnt += 1
            if cnt > 10:
                return None
        return _TripInner(trip_id, depart_time_min * 60, from_node, from_type, first_node, "Home")

    def __genTripsChain1(self, ev:_EVInner):  # vehicle_trip
        """
        Generate a full day of trips on the first day
            ev: vehicle instance
        """
        daynum = 0
        weekday = True
        trip_1 = self.__genFirstTrip1("trip0_1", weekday)
        trip_2 = self.__genTripA("trip0_2",trip_1.toT,trip_1.toN,trip_1.DPTT,weekday)
        trip_3 = self.__genTripF("trip0_3",trip_2.toT,trip_2.toN,trip_2.DPTT,trip_1.frN,weekday)
        
        ev._add_trip(daynum, trip_1)
        ev._add_trip(daynum, trip_2)
        if trip_3: # Trip3: if O==D, don't generate trip 3
            if trip_3.DPTT < 86400:  # If the departure time is after midnight, it is not valid
                ev._add_trip(daynum, trip_3)
            else:
                trip_2.toN = trip_1.frN

    def __genFirstTripA(self, trip_id:str, ev: _EVInner, weekday: bool = True):
        """
        Generate the first trip of a non-first day
            trip_id: Trip ID
            vehicle_node: Vehicle node, such as rootNode.getElementsByTagName("vehicle")[0]
            weekday: Whether it is weekday or weekend
        """
        trip_last = ev.trips[-1]
        from_node = trip_last.toN
        # Get departure time and destination area type
        from_Type = "Home"
        depart_time_min, to_Type = self.__getDest1(from_Type, weekday)
        to_node = self.__getNextNode(from_node, from_Type)
        return _TripInner(trip_id, depart_time_min * 60, from_node, from_Type, to_node, to_Type)

    def __genTripsChainA(self, ev: _EVInner, daynum: int = 1):  # vehicle_trip
        """
        Generate a full day of trips on a non-first day
        """
        weekday = (daynum - 1) % 7 + 1 in [1, 2, 3, 4, 5]
        trip2_1 = self.__genFirstTripA(f"trip{daynum}_1", ev, weekday)
        trip2_2 = self.__genTripA(f"trip{daynum}_2", trip2_1.toT,trip2_1.toN,trip2_1.DPTT,weekday)
        trip2_3 = self.__genTripF(f"trip{daynum}_3", trip2_2.toT,trip2_2.toN, trip2_2.DPTT,trip2_1.frN,weekday)
                    
        ev._add_trip(daynum, trip2_1)
        ev._add_trip(daynum, trip2_2)
        if trip2_3:
            if trip2_3.DPTT < 86400:  # If the departure time is after midnight, it is not valid
                ev._add_trip(daynum, trip2_3)
            else:
                trip2_2.toN = trip2_1.frN
        
    def __genEV(self, veh_id: str, day_count:int, **kwargs) -> _EVInner:
        '''
        Generate a full week of trips for a vehicle as an inner instance
        '''
        ev = _EVInner(veh_id, random.choice(self.vTypes), self.soc_pdf.sample()/100.0, **kwargs)
        self.__genTripsChain1(ev)
        for j in range(1, day_count + 1):
            self.__genTripsChainA(ev, j)
        return ev

    def genEV(self, veh_id: str, **kwargs) -> EV:
        """
        Generate a full week of trips for a vehicle.
        The generated vehicle is returned as an EV instance, and will not be held in the buffer.
            veh_id: ID of the vehicle
            v2g_prop: Proportion of users willing to participate in V2G
            omega: PDFunc | None = None,
            krel: PDFunc | None = None,
            ksc: PDFunc | None = None,
            kfc: PDFunc | None = None,
            kv2g: PDFunc | None = None
        """
        return self.__genEV(veh_id, **kwargs).toEV()

    def genEVs(
        self, N: int, fname: Optional[str] = None, day_count: int = 7, silent: bool = False, **kwargs
    ) -> EVDict:
        """
        Generate EV and trips of N vehicles.
        The generated vehicles are returned as an EVDict instance, and will be saved to the file if fname is provided.
        The vehicles will not be held in the buffer.
            N: Number of vehicles
            fname: Saved file name (if None, not saved)
            day_count: Number of days
            silent: Whether silent mode
            v2g_prop: Proportion of users willing to participate in V2G
            omega: PDFunc | None = None,
            krel: PDFunc | None = None,
            ksc: PDFunc | None = None,
            kfc: PDFunc | None = None,
            kv2g: PDFunc | None = None
        """
        st_time = time.time()
        last_print_time = 0
        saver = _xmlSaver(fname) if fname else None
        ret = EVDict()
        for i in range(0, N):
            ev = self.__genEV("v" + str(i), day_count,
                cache_route = self._route_cache_mode != RoutingCacheMode.NONE, **kwargs)
            ret.add(ev.toEV())
            if saver:
                saver.write(ev)
            if not silent and time.time()-last_print_time>1:
                print(f"\r{i+1}/{N}, {(i+1)/N*100:.2f}%", end="")
                last_print_time=time.time()
        if not silent:
            print(f"\r{N}/{N}, 100.00%")
            print(Lang.INFO_DONE_WITH_SECOND.format(round(time.time() - st_time, 1)))
        if saver:
            saver.close()
        return ret

class ManualEVsGenerator:
    """Class to manually add EVs to the buffer"""
    def __init__(self):
        self.__evs:Dict[str, _EVInner] = {}
    
    def addEV(self, vid:str, bcap_kWh:float, range_km:float, efc_rate_kW:float, 
            esc_rate_kW:float, max_V2G_kW:float, soc:float, omega:float, 
            krel:float, ksc:float, kfc:float, kv2g:float, cache_route:bool = False) -> _EVInner:
        """
        Add an EV to the generator's buffer
            vid: Vehicle ID
            vtype: Vehicle type
            soc: State of charge (0.0~1.0)
        """
        if vid in self.__evs:
            raise ValueError(f"Vehicle ID {vid} already exists.")
        ev = _EVInner(vid, VehicleType(id=-1, bcap_kWh=bcap_kWh, range_km=range_km,
            efc_rate_kW=efc_rate_kW, esc_rate_kW=esc_rate_kW, max_V2G_kW=max_V2G_kW), 
            soc, omega=omega, krel=krel, ksc=ksc, kfc=kfc, kv2g=kv2g, cache_route=cache_route)
        self.__evs[vid] = ev
        return ev
    
    def dumpEVs(self, fname: str):
        """
        Dump all EVs in the buffer to a file
            fname: File name
        """
        saver = _xmlSaver(fname)
        for ev in self.__evs.values():
            saver.write(ev)
        saver.close()
    
    def getEVs(self) -> EVDict:
        """
        Get all EVs in the buffer as an EVDict instance
        """
        ret = EVDict()
        for ev in self.__evs.values():
            ret.add(ev.toEV())
        return ret