from __future__ import annotations

from typing import Dict, List

import numpy as np

from src.busopt.config import ModelConfig


class Node:
    def __init__(self, node_id: str, node_type: str = "default"):
        self.node_id = node_id
        self.node_type = node_type
        self.in_links = []
        self.out_links = []
        self.out_hyperlinks = []
        self.in_hyperlinks = []
        self.wait = 0
        self.demand_wait = {}

    def __str__(self):
        return f"node_id: {self.node_id}"


class Link:
    def __init__(
        self,
        link_id: str,
        from_node: Node,
        to_node: Node,
        config: ModelConfig,
        in_route=None,
        link_attr=None,
        static_cost=0,
        length=0,
        freq=1e6,
        link_type="default",
    ):
        self.link_id = link_id
        self.from_node = from_node
        self.to_node = to_node
        self.config = config
        self.length = length
        self.in_route = in_route
        self.static_cost = static_cost
        self.cost = static_cost
        self.travel_cost = static_cost
        self.fare_cost = 0
        self.freq = freq
        self.flow = 0
        self.flow_demand = {}
        self.link_type = link_type
        self.dept_link = None
        self.partial_dev = {}
        self.link_attr = link_attr

    def __str__(self):
        return (
            f"link_id:{self.link_id},"
            f"travel_cost:{self.travel_cost},"
            f"fare_cost:{self.fare_cost},"
            f"freq:{self.freq},"
            f"flow:{self.flow}"
        )

    def link_cost_func_sym(self):
        config = self.config
        if self.link_type == "board":
            if self.freq == 0:
                return lambda v: 1e9
            return lambda v: config.beta_board * (v / (config.capacity * self.freq)) ** config.exponent_n
        if self.link_type == "transit":
            return lambda v: self.static_cost + config.beta_transit * (v / config.capacity) ** config.exponent_n
        if self.link_type in {"alight", "walk"}:
            return lambda v: self.static_cost
        raise ValueError(f"Unsupported link type: {self.link_type}")

    def link_cost_func_asym(self):
        config = self.config
        if self.link_type == "board":
            if self.freq == 0:
                return lambda v_b, v_t: 1e9
            return lambda v_b, v_t: config.beta_board * (
                (v_b + config.gamma_board * v_t - config.gamma_board * v_b) / (config.capacity * self.freq)
            ) ** config.exponent_n
        if self.link_type == "transit":
            freq = self.dept_link.freq
            if freq == 0:
                return lambda v_b, v_t: 1e9
            return lambda v_t, v_b: self.static_cost + config.beta_transit * (
                (v_t - v_b + config.gamma_transit * v_b) / (config.capacity * freq)
            ) ** config.exponent_n
        if self.link_type in {"alight", "walk"}:
            return lambda v_t, v_b: self.static_cost
        raise ValueError(f"Unsupported link type: {self.link_type}")

    def link_cost_func_asym_with_fvalue(self, fre_b=None, fre_dept=None):
        config = self.config
        if self.link_type == "board":
            return lambda v_b, v_t: config.beta_board * (
                (v_b + config.gamma_board * v_t - config.gamma_board * v_b) / (config.capacity * fre_b)
            ) ** config.exponent_n
        if self.link_type == "transit":
            return lambda v_t, v_b: self.static_cost + config.beta_transit * (
                (v_t - v_b + config.gamma_transit * v_b) / (config.capacity * fre_dept)
            ) ** config.exponent_n
        if self.link_type in {"alight", "walk"}:
            return lambda v_t, v_b: self.static_cost
        raise ValueError(f"Unsupported link type: {self.link_type}")

    def second_func(self):
        return lambda v_b, v_t: self.link_cost_func_asym()(v_b, v_t) * v_b

    def cal_cost(self):
        return self.link_cost_func_asym()(self.flow, self.dept_link.flow if self.dept_link is not None else None)

    def cal_partial_dev(self, v_point=None, func_case="asymmetric"):
        config = self.config

        if self.link_type == "board":
            related_transit = self.dept_link
            v_b = v_point if v_point is not None else self.flow
            v_t = related_transit.flow
            if self.freq == 0:
                return {self: 1e9, related_transit: 1e9}
            common = (
                config.exponent_n
                * config.beta_board
                * (v_b + config.gamma_board * v_t - config.gamma_board * v_b) ** (config.exponent_n - 1)
                / (config.capacity * self.freq) ** config.exponent_n
            )
            return {
                self: common * (1 - config.gamma_board),
                related_transit: common * config.gamma_board,
            }
        if self.link_type == "transit":
            related_board = self.dept_link
            v_t = v_point if v_point is not None else self.flow
            v_b = related_board.flow
            freq = related_board.freq
            if freq == 0:
                return {self: 1e9, related_board: 1e9}
            common = (
                config.exponent_n
                * config.beta_transit
                * (v_t - v_b + config.gamma_transit * v_b) ** (config.exponent_n - 1)
                / (config.capacity * freq) ** config.exponent_n
            )
            return {
                self: common,
                related_board: common * (config.gamma_transit - 1),
            }
        if self.link_type in {"alight", "walk"}:
            return {self: 0}
        raise ValueError(f"Unsupported link type: {self.link_type}")



class Graph:
    def __init__(self, config: ModelConfig):
        self.config = config
        self.node_dict = {}
        self.link_dict = {}
        self.hyperlink_dict = {}

    def build_graph(self, adj_matrix):
        node_num = adj_matrix.shape[0]
        for i in range(node_num):
            new_node = Node(str(i))
            self.node_dict[new_node.node_id] = new_node
        for i in range(node_num):
            for j in range(node_num):
                if not np.isnan(adj_matrix[i, j]):
                    from_node = self.node_dict.get(str(i))
                    to_node = self.node_dict.get(str(j))
                    new_link = Link(
                        link_id=from_node.node_id + ", " + to_node.node_id,
                        from_node=from_node,
                        to_node=to_node,
                        config=self.config,
                        static_cost=adj_matrix[i, j],
                    )
                    self.link_dict[new_link.link_id] = new_link
                    from_node.out_links.append(new_link)
                    to_node.in_links.append(new_link)

    def clear_info(self):
        for node in self.node_dict.values():
            node.wait = 0
        for link in self.link_dict.values():
            link.flow = 0

    def clear_isolate_node(self):
        self.node_dict = dict(filter(lambda item: item[1].out_links or item[1].in_links, self.node_dict.items()))

    def get_info(self):
        flow = {link_id: link.flow for link_id, link in self.link_dict.items()}
        waiting = {node_id: node.wait for node_id, node in self.node_dict.items()}
        link_cost = {link_id: link.cost for link_id, link in self.link_dict.items()}
        return flow, waiting, link_cost


class Hyperlink:
    def __init__(self, from_node, link_list):
        self.from_node: Node = from_node
        self.link_list: List[Link] = link_list
        self.has_waiting = False
        self.to_node_list = [link.to_node for link in link_list]
        self.id = f"{self.from_node.node_id}, " + ", ".join([node.node_id for node in self.to_node_list])
        self.link_dict = {link.to_node.node_id: link for link in self.link_list}
        self.link_prob: Dict[Link, float] = {}
        self.update_link_prob()

    def update_link_prob(self):
        freq_sum = sum([link.freq for link in self.link_list])
        self.link_prob = {link: link.freq / freq_sum for link in self.link_list}

    def __str__(self):
        return self.id


class Hyperpath:
    def __init__(self):
        self.node_dict = {}
        self.link_dict = {}
        self.node_prob = {}
        self.link_prob = {}
        self.unit_node_wait = {}
        self.path_flow = 0
        self.cost = float("inf")
        self.travel_and_fare_cost = float("inf")
        self.waiting_cost = float("inf")

    def cal_path_cost(self):
        self.travel_and_fare_cost = 0
        self.waiting_cost = 0
        for link in self.link_dict.values():
            self.travel_and_fare_cost += self.link_prob[link.link_id] * link.cost
        for node in self.node_dict.values():
            if node.node_id.endswith("s"):
                link_set = set(node.out_links) & set(self.link_dict.values())
                self.waiting_cost += 0 if not link_set else self.node_prob[node.node_id] / sum([link.freq for link in link_set])
        self.cost = self.travel_and_fare_cost + self.waiting_cost
        return self.cost

    def load_unit_node_waiting(self):
        for node in self.node_dict.values():
            if node.node_id.endswith("s"):
                link_set = set(node.out_links) & set(self.link_dict.values())
                self.unit_node_wait[node.node_id] = 0 if not link_set else self.node_prob[node.node_id] / sum(
                    [link.freq for link in link_set]
                )
            else:
                self.unit_node_wait[node.node_id] = 0

    def get_link_prob(self, link_id):
        return self.link_prob.get(link_id, 0)

    def __hash__(self) -> int:
        return super().__hash__()

    def __eq__(self, other):
        return self.link_dict == other.link_dict

    def __str__(self):
        return f"cost = {self.cost:.2f}, path_flow = {self.path_flow}"


class OD_demand:
    def __init__(self, ori_node, des_node, demand):
        self.ori_node = ori_node
        self.des_node = des_node
        self.demand = demand
        self.demand_id = f"{ori_node.node_id}, {des_node.node_id}"

    def __str__(self):
        return f"{self.ori_node.node_id} -> {self.des_node.node_id} : {self.demand}"

    @staticmethod
    def destination_based_OD(demand_list):
        demand_dict = {}
        for demand in demand_list:
            demand_dict.setdefault(demand.des_node, []).append(demand)
        return demand_dict


class Route:
    def __init__(self, node_seq, org_freq, freq=None, route_id=None):
        self.node_seq = node_seq
        self.org_freq = org_freq
        self.freq = freq
        self.route_id = route_id
        self.hypernodes = {}
        self.hyperlinks = {}
        self.length = None
        self.time = None
        self.attr = None

    def cal_length(self, basegraph):
        self.length = 0
        for ni, nj in zip(self.node_seq[:-1], self.node_seq[1:]):
            self.length += basegraph.link_dict.get(ni.node_id + ", " + nj.node_id).length
        return self.length

    def cal_time(self, basegraph):
        self.time = 0
        for ni, nj in zip(self.node_seq[:-1], self.node_seq[1:]):
            self.time += basegraph.link_dict.get(ni.node_id + ", " + nj.node_id).static_cost
        return self.time

    def __eq__(self, other):
        return self.node_seq == other.node_seq and self.freq == other.freq


class Line:
    def __init__(self, line_id, org_freq, length, time, routes):
        self.line_id = line_id
        self.org_freq = org_freq
        self.length = length
        self.time = time
        self.routes = routes
