import heapq
from time import perf_counter
from src.entity.network_entity import *
from src.entity.utils import Logger

from src.io_utils.input_data_utils import read_mandl_data


def transit_assignment_gp(hypergraph: Graph, demand_list, enable_logging=False):
    """
    Gradient Projection (GP) method for solving the static transit assignment problem.
    Reference:
    Xu, Z., Xie, J., Liu, X., & Nie, Y. (Marco). (2020).
    Hyperpath-based algorithms for the transit equilibrium assignment problem.
    Transportation Research Part E: Logistics and Transportation Review, 143, 102102.
    https://doi.org/10.1016/j.tre.2020.102102
    :param hypergraph:
    :param demand_list:
    :param enable_logging:
    """
    mylogger = Logger(__name__)
    if not enable_logging:
        mylogger.mute()


    update_cost(hypergraph)     # update link cost in the hypergraph
    demand_dict = OD_demand.destination_based_OD(demand_list)   # rearrange OD by destinations
    path_pool = {w: [] for w in demand_list}  # setup path pool

    # clear link flow / node waiting time
    for a in hypergraph.link_dict.values():
        a.flow = 0
    for node in hypergraph.node_dict.values():
        node.demand_wait = {d: 0 for d in demand_dict.keys()}

    """
    Step1: Initialization. All-or-nothing assignment
    """
    t1 = perf_counter()
    mylogger.logger.info("start initial assignment:")

    for w_list in demand_dict.values():
        # compute the shortest path for each OD pair, and assign all flows to it
        init_path_dict = label_setting_shortest_path(hypergraph, w_list)
        for w, path in init_path_dict.items():
            path.path_flow = w.demand
            path_pool[w].append(path)
            # update link flow
            for a in path.link_dict.values():
                a.flow += path.path_flow * path.link_prob[a.link_id]
            for node in path.node_dict.values():
                node.demand_wait[w.des_node] = node.demand_wait.get(w.des_node, 0) + path.path_flow * path.unit_node_wait[node.node_id]

    update_cost(hypergraph)
    mylogger.logger.info(f"Initialization time: {perf_counter() - t1:.2f} secs")

    """
    Step2: Main loop
    """
    base_time = perf_counter()
    time_record, gap_record = [0], [1.0]
    path_num_record = [sum([len(pool) for pool in path_pool.values()])]
    iter_step = 0
    epsilon = 5e-4
    time_limit = 5
    mylogger.logger.info(f"in iter: {iter_step}\tElapsed time: {time_record[-1]}\tRelative Gap: {gap_record[-1]}")

    while gap_record[-1] >= epsilon and time_record[-1] <= time_limit:
        iter_step += 1
        for w_list in demand_dict.values():
            for w in w_list:
                # (i) update path cost
                for path in path_pool[w]:
                    path.cal_path_cost()
                # (ii) compute current shortest path
                new_path = label_setting_shortest_path(hypergraph, [w])
                new_path = list(new_path.values())[0]
                if new_path not in path_pool[w]:
                    path_pool[w].append(new_path)
                # (iii) sort the path by cost
                path_pool[w].sort(key=lambda x: x.cost)
                bsc_path = path_pool[w][0]

                # (iv) updating process
                #   (iv.i) update nonbasic path flows
                for non_bsc_path in path_pool[w][1:]:
                    first_drv = non_bsc_path.cost - bsc_path.cost  # first-order derivative
                    second_drv = 0  # second-order derivative
                    for link in set(non_bsc_path.link_dict.values()) | set(bsc_path.link_dict.values()):
                        term1 = non_bsc_path.get_link_prob(link.link_id) - bsc_path.get_link_prob(link.link_id)
                        term2 = sum([value * (non_bsc_path.get_link_prob(l.link_id) - bsc_path.get_link_prob(l.link_id))
                                     for l, value in link.partial_dev.items()])
                        second_drv += term1 * term2
                    record = non_bsc_path.path_flow
                    non_bsc_path.path_flow = max(0, record - first_drv / second_drv)
                    diff = non_bsc_path.path_flow - record

                    for link in set(non_bsc_path.link_dict.values()):
                        link.flow += non_bsc_path.link_prob[link.link_id] * diff
                    for link in set(bsc_path.link_dict.values()):
                        link.flow += bsc_path.link_prob[link.link_id] * (-diff)

                    for node in non_bsc_path.node_dict.values():
                        node.demand_wait[w.des_node] = node.demand_wait.get(w.des_node, 0) + non_bsc_path.unit_node_wait[
                            node.node_id] * diff
                    for node in bsc_path.node_dict.values():
                        node.demand_wait[w.des_node] = node.demand_wait.get(w.des_node, 0) - bsc_path.unit_node_wait[node.node_id] * diff

                #   (iv.ii) update basic path flows
                bsc_path.path_flow = w.demand - sum([p.path_flow for p in path_pool[w][1:]])

                #   (iv.iii) delete path with too little flow
                path_pool[w] = list(filter(lambda x: x.path_flow >= 1e-7, path_pool[w]))

                #   (iv.iv) update link cost
                update_cost(hypergraph)

        # compute node waiting time
        for i in hypergraph.node_dict.values():
            i.wait = sum(i.demand_wait.values())

        time_record.append(perf_counter() - base_time)  # 璁板綍鏃堕棿鍜岃绠楃簿搴?
        path_num_record.append(sum([len(pool) for pool in path_pool.values()]))
        gap_record.append(calculate_relative_gap(path_pool, demand_list))
        mylogger.logger.info(
            f"in iter: {iter_step}\tElapsed time: {time_record[-1]:.2f} secs\tPool size: {path_num_record[-1]}\tRelative gap: {gap_record[-1]}")

    # v_a -> v_a^w
    for a in hypergraph.link_dict.values():
        a.flow_demand = {}
    for d, paths in path_pool.items():
        for path in paths:
            for a in path.link_dict.values():
                f = path.path_flow * path.link_prob[a.link_id]
                a.flow_demand[d] = f if d not in a.flow_demand.keys() else a.flow_demand[d] + f

    return time_record, gap_record


def update_cost(hypergraph: Graph):
    for link in hypergraph.link_dict.values():

        # update link cost
        link.travel_cost = link.cal_cost()
        link.cost = link.travel_cost + link.fare_cost

        # update derivative
        link.partial_dev.update(link.cal_partial_dev())



def calculate_relative_gap(path_pool, demand_list):
    best_path_cost = sum([r.demand * path_pool[r][0].cost for r in demand_list])
    total_cost = sum([path.cost * path.path_flow for od_pool in path_pool.values() for path in od_pool])
    return 1 - best_path_cost / total_cost



def label_setting_shortest_path(hypergraph: Graph, w_list):
    """
    A label-setting algorithm for solving the (multi-origins-one-destination) shortest hyperpath problem
    in the hypergraph.
    """

    class CustomPriorityQueue:
        """
        This class defines a priority queue structure for labels realized by Python heapq.
        It has the following methods:
        qsize(): return the queue size
        pop(): pop the first element (with the highest priority)
        put(): add an element to the queue by priority ascending order
        """

        def __init__(self):
            self.queue = []
            self.index = 0

        def qsize(self):
            return len(self.queue)

        def pop(self):
            return heapq.heappop(self.queue)[-1]

        def put(self, item, priority):
            heapq.heappush(self.queue, (priority, self.index, item))
            self.index += 1

    # noinspection PyShadowingNames
    def compute_attractive_set(expect_cost):
        """
        A greedy method to compute the optimal attractive set `R_opt`.
        Note that the vehicle headway should follow a negative exponential distribution;
        otherwise, the optimal attractive set can only be obtained by enumeration.
        :param expect_cost: Expected cost from each link to the destination node.
        """
        R_opt = []  # Optimal attractive set
        u_opt = float("inf")  # Minimum generalized cost
        L = sorted(expect_cost, key=lambda x: expect_cost[x])
        for link in L:
            R_temp = R_opt + [link]
            temp = max(sum([l.freq for l in R_temp]), 1e-8)
            pi_l = {l: l.freq / temp for l in R_temp}
            u = 1 / temp + sum([expect_cost[l] * pi_l[l] for l in R_temp])
            if u < u_opt - 1e-6:
                R_opt.append(link)
                u_opt = u
            else:
                break

        return R_opt, u_opt

    # Initialization
    des_node = w_list[0].des_node
    ori_nodes = [w.ori_node for w in w_list]

    u_i = {node: float("inf") for node in hypergraph.node_dict.values()}  # u_i = inf if i != d, otherwise 0
    u_i[des_node] = 0
    Q = CustomPriorityQueue()
    Q.put(des_node, u_i[des_node])  # Use a priority queue to store nodes to be processed
    L_i_opt = {node: set() for node in hypergraph.node_dict.values()}
    # Maintain an extra dictionary to record whether a node has been visited.
    # This is valid in the label-setting algorithm because each node is visited at most once.
    # Therefore, there is no need to dynamically adjust the priority of existing queue entries.
    visited = {i: False for i in hypergraph.node_dict.values()}
    unsolved_num = len(ori_nodes)  # The algorithm can terminate early once all origins are settled
    while Q.qsize() > 0 and unsolved_num > 0:
        j = Q.pop()  # Current node being processed
        if not visited[j]:
            for inLink in j.in_links:
                if inLink.link_type == "board":  # case1: boarding link
                    i = inLink.from_node
                    expect_cost = {l: l.cost + u_i[l.to_node] for l in i.out_links}  # s_jp = c_i_jp + u_jp
                    R_opt_i, u_opt_i = compute_attractive_set(expect_cost)  # Compute the optimal attractive set and its cost
                    if u_opt_i < u_i[i]:
                        u_i[i] = u_opt_i
                        L_i_opt[i] = R_opt_i
                        Q.put(i, u_i[i])

                else:  # case2: other links
                    i = inLink.from_node
                    if inLink.cost + u_i[j] < u_i[i]:
                        u_i[i] = inLink.cost + u_i[j]  # Update the label of node i
                        L_i_opt[i] = [inLink]
                        Q.put(i, u_i[i])

            visited[j] = True  # Mark the current node as processed
            if j in ori_nodes:
                unsolved_num -= 1

    # Build the shortest hyperpath from each origin to the destination
    path_opt_dict = {}
    for w in w_list:
        ori_node = w.ori_node
        path_opt = Hyperpath()

        # Use DFS to build node_dict and link_dict, while recording node depth
        def cal_path_DFS(node, depth, node_dict, link_dict, depth_dict):
            d = max(node_depth.get(node, 0), depth)
            depth_dict[node] = d
            node_dict[node.node_id] = node
            link_dict.update({l.link_id: l for l in L_i_opt[node]})
            next_nodes = [l.to_node for l in L_i_opt[node]]
            for sub_n in next_nodes:
                cal_path_DFS(sub_n, d + 1, node_dict, link_dict, depth_dict)

        node_depth = {}
        path_opt.node_dict = {}
        path_opt.link_dict = {}
        cal_path_DFS(ori_node, 0, path_opt.node_dict, path_opt.link_dict, node_depth)

        # Recompute path probabilities
        unsolved_node = sorted(node_depth, key=lambda x: node_depth[x])
        e_ij = {i: {l: l.freq / sum([ls.freq for ls in L_i_opt[i]]) for l in L_i_opt[i]} for i in unsolved_node}
        for cur_n in unsolved_node:
            # Compute the probability of the current node
            if cur_n == ori_node or cur_n == des_node:
                path_opt.node_prob[cur_n.node_id] = 1.0
            else:
                in_links = set(cur_n.in_links) & set(path_opt.link_dict.values())
                path_opt.node_prob[cur_n.node_id] = sum(
                    [path_opt.node_prob[l.from_node.node_id] * e_ij[l.from_node][l] for l in in_links])
            out_links = L_i_opt[cur_n]
            # Compute the probability of each outgoing link from the current node
            for l in out_links:
                path_opt.link_prob[l.link_id] = path_opt.node_prob[cur_n.node_id] * e_ij[cur_n][l]

        # Compute path cost
        path_opt.cost = u_i[ori_node]
        # Compute the unit waiting time at each node along the path
        path_opt.load_unit_node_waiting()
        # Store the result
        path_opt_dict[w] = path_opt

    return path_opt_dict


if __name__ == '__main__':
    print("transit assignment test:")
    config = ModelConfig()
    _, hypergraph, route_list, line_list, demand_list = read_mandl_data(config.candidate_routes, config)
    transit_assignment_gp(hypergraph, demand_list, enable_logging=True)
