from __future__ import annotations

from itertools import combinations
from pathlib import Path

import pandas as pd

from src.busopt.config import ModelConfig
from src.entity.network_entity import Graph, Hyperlink, Line, Link, Node, OD_demand, Route


REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_DATA_DIR = REPO_ROOT / "data" / "Mandl"


def _resolve_path(path_like: str | Path) -> Path:
    path = Path(path_like)
    return path if path.is_absolute() else REPO_ROOT / path


def read_mandl_data(candidate_route_file: str | Path, config: ModelConfig):
    stop_file = DEFAULT_DATA_DIR / "Mandl_stop.txt"
    base_file = DEFAULT_DATA_DIR / "Mandl_base.txt"
    trip_file = DEFAULT_DATA_DIR / "Mandl_trip.txt"
    candidate_route_file = _resolve_path(candidate_route_file)

    basegraph = Graph(config)
    hypergraph = Graph(config)
    demand_list = []

    for _, row in pd.read_csv(stop_file).iterrows():
        node_id = str(row.stopid)
        basegraph.node_dict[node_id] = Node(node_id)
        hypergraph.node_dict[node_id + "_s"] = Node(node_id + "_s", node_type="stop")

    for _, row in pd.read_csv(base_file).iterrows():
        from_node_id = str(int(row.fromstop))
        to_node_id = str(int(row.tostop))
        cost = row.fft / 60
        length = cost * 30
        new_link = Link(
            link_id=from_node_id + ", " + to_node_id,
            from_node=basegraph.node_dict[from_node_id],
            to_node=basegraph.node_dict[to_node_id],
            config=config,
            static_cost=cost,
            length=length,
        )
        basegraph.link_dict[new_link.link_id] = new_link
        basegraph.node_dict[from_node_id].out_links.append(new_link)
        basegraph.node_dict[to_node_id].in_links.append(new_link)

    for _, row in pd.read_csv(trip_file).iterrows():
        origin_id = str(int(row.origin)) + "_s"
        destination_id = str(int(row.destination)) + "_s"
        demand_list.append(
            OD_demand(hypergraph.node_dict[origin_id], hypergraph.node_dict[destination_id], float(row.demand))
        )

    demand_list = sorted(demand_list, key=lambda item: item.demand, reverse=True)
    demand_list = demand_list[0::2][: config.od_num]
    for demand in demand_list:
        demand.demand = demand.demand * config.demand_coef / 24

    route_list = []
    line_list = []
    with candidate_route_file.open("r", encoding="utf-8") as file_obj:
        next(file_obj)
        for line in file_obj:
            route_id, org_freq, *stops = line.rstrip("\n").split(",")
            route = Route(
                node_seq=[basegraph.node_dict[stop_id] for stop_id in stops],
                org_freq=float(org_freq),
                freq=float(org_freq),
                route_id=route_id,
            )
            route.cal_length(basegraph)
            route.cal_time(basegraph)
            route_list.append(route)

    route_list = route_list[: 2 * config.route_num]
    for index in range(0, len(route_list), 2):
        upstream_route = route_list[index]
        downstream_route = route_list[index + 1]
        line_list.append(
            Line(
                line_id=upstream_route.route_id.split("_")[0],
                org_freq=upstream_route.freq,
                length=upstream_route.length,
                time=upstream_route.time,
                routes={"upstream": upstream_route, "downstream": downstream_route},
            )
        )

    load_route_on_hypergraph(basegraph, hypergraph, route_list, config)
    load_flow_dependent_links(hypergraph)
    enumerate_hyperlinks(hypergraph)
    return basegraph, hypergraph, route_list, line_list, demand_list


def load_route_on_hypergraph(basegraph, hypergraph, route_list, config: ModelConfig):
    for route in route_list:
        for node in route.node_seq:
            node_id = node.node_id + "_" + route.route_id
            hypergraph.node_dict[node_id] = Node(node_id, node_type="transit")

        for base_i, base_j in zip(route.node_seq[:-1], route.node_seq[1:]):
            hyper_i = hypergraph.node_dict[base_i.node_id + "_" + route.route_id]
            hyper_j = hypergraph.node_dict[base_j.node_id + "_" + route.route_id]
            transit_link = Link(
                link_id=hyper_i.node_id + ", " + hyper_j.node_id,
                from_node=hyper_i,
                to_node=hyper_j,
                config=config,
                static_cost=basegraph.link_dict[base_i.node_id + ", " + base_j.node_id].static_cost,
                length=basegraph.link_dict[base_i.node_id + ", " + base_j.node_id].length,
                in_route=route,
                link_type="transit",
                link_attr=route.attr,
            )
            transit_link.fare_cost = transit_link.length * config.tau / config.vot
            transit_link.cost = transit_link.travel_cost + transit_link.fare_cost
            hypergraph.link_dict[transit_link.link_id] = transit_link
            route.hyperlinks[transit_link.link_id] = transit_link
            hyper_i.out_links.append(transit_link)
            hyper_j.in_links.append(transit_link)

        for index, base_node in enumerate(route.node_seq):
            stop_node = hypergraph.node_dict[base_node.node_id + "_s"]
            transit_node = hypergraph.node_dict[base_node.node_id + "_" + route.route_id]
            if index != len(route.node_seq) - 1:
                boarding_link = Link(
                    link_id=stop_node.node_id + ", " + transit_node.node_id,
                    from_node=stop_node,
                    to_node=transit_node,
                    config=config,
                    freq=route.freq,
                    in_route=route,
                    link_type="board",
                    link_attr=route.attr,
                )
                hypergraph.link_dict[boarding_link.link_id] = boarding_link
                route.hyperlinks[boarding_link.link_id] = boarding_link
                stop_node.out_links.append(boarding_link)
                transit_node.in_links.append(boarding_link)
            if index != 0:
                alighting_link = Link(
                    link_id=transit_node.node_id + ", " + stop_node.node_id,
                    from_node=transit_node,
                    to_node=stop_node,
                    config=config,
                    in_route=route,
                    link_type="alight",
                    link_attr=route.attr,
                )
                hypergraph.link_dict[alighting_link.link_id] = alighting_link
                route.hyperlinks[alighting_link.link_id] = alighting_link
                transit_node.out_links.append(alighting_link)
                stop_node.in_links.append(alighting_link)

        route.transit_links = [link for link in route.hyperlinks.values() if link.link_type == "transit"]
        route.board_links = [link for link in route.hyperlinks.values() if link.link_type == "board"]
        route.alight_links = [link for link in route.hyperlinks.values() if link.link_type == "alight"]


def load_flow_dependent_links(hypergraph):
    for link in hypergraph.link_dict.values():
        if link.link_type == "transit":
            transit_node_id = link.from_node.node_id
            last_idx = transit_node_id.find("_")
            transfer_node_id = transit_node_id[:last_idx] + "_s"
            link.dept_link = hypergraph.link_dict[transfer_node_id + ", " + transit_node_id]
        elif link.link_type == "board":
            for outbound_link in link.to_node.out_links:
                if outbound_link.link_type == "transit":
                    link.dept_link = outbound_link
        elif link.link_type in {"walk", "alight"}:
            continue
        else:
            raise ValueError(f"Unsupported link type: {link.link_type}")


def enumerate_hyperlinks(hypergraph):
    for node in [item for item in hypergraph.node_dict.values() if item.node_type == "stop"]:
        out_links_board = [link for link in node.out_links if link.link_type == "board"]
        all_combinations = [
            list(comb)
            for size in range(1, len(out_links_board) + 1)
            for comb in combinations(out_links_board, size)
        ]
        for comb in all_combinations:
            hyperlink = Hyperlink(node, comb)
            hyperlink.has_waiting = True
            hypergraph.hyperlink_dict[hyperlink.id] = hyperlink
            node.out_hyperlinks.append(hyperlink)
            for to_node in hyperlink.to_node_list:
                to_node.in_hyperlinks.append(hyperlink)

        out_links_walk = [link for link in node.out_links if link.link_type == "walk"]
        for link in out_links_walk:
            hyperlink = Hyperlink(node, [link])
            hypergraph.hyperlink_dict[hyperlink.id] = hyperlink
            node.out_hyperlinks.append(hyperlink)
            link.to_node.in_hyperlinks.append(hyperlink)

    for node in [item for item in hypergraph.node_dict.values() if item.node_type == "transit"]:
        for link in node.out_links:
            hyperlink = Hyperlink(node, [link])
            hypergraph.hyperlink_dict[hyperlink.id] = hyperlink
            node.out_hyperlinks.append(hyperlink)
            link.to_node.in_hyperlinks.append(hyperlink)
