import heapq
import itertools
from dataclasses import dataclass
from math import sin, cos, acos, radians, isinf
from pathlib import Path
from typing import NewType

from tqdm import tqdm


StationId = NewType('StationId', int)
StationGroupId = NewType('StationGroupId', int)

@dataclass(eq=True, frozen=True)
class Station:
    id: StationId
    group_id: StationGroupId
    name: str
    longitude: float
    latitude: float

def compute_distance(station1: Station, station2: Station) -> float:
    R = 6378.137 # equatorial radius in km
    x1 = radians(station1.longitude)
    y1 = radians(station1.latitude)
    x2 = radians(station2.longitude)
    y2 = radians(station2.latitude)
    dist = R * acos(sin(y1) * sin(y2) + cos(y1) * cos(y2) * cos(x2 - x1))
    return dist


def list_route_distances(graph: dict[StationGroupId, list[tuple[StationGroupId, float]]], start: StationGroupId):
    distances = { node: float('inf') for node in graph }
    distances[start] = 0
    heap = [(0.0, start)]
    while heap:
        current_distance, current_node = heapq.heappop(heap)
        if current_distance > distances[current_node]:
            continue
        for neighbour, neighbour_distance in graph[current_node]:
            distance = current_distance + neighbour_distance
            if distance < distances[neighbour]:
                distances[neighbour] = distance
                heapq.heappush(heap, (distance, neighbour))
    return distances

def compute_route(graph: dict[StationGroupId, list[tuple[StationGroupId, float]]], start: StationGroupId, end: StationGroupId):
    distances = { node: float('inf') for node in graph }
    distances[start] = 0
    previous_nodes: dict[StationGroupId, StationGroupId | None] = { node: None for node in graph }
    heap = [(0.0, start)]
    while heap:
        current_distance, current_node = heapq.heappop(heap)
        if current_distance > distances[current_node]:
            continue
        for neighbour, neighbour_distance in graph[current_node]:
            distance = current_distance + neighbour_distance
            if distance < distances[neighbour]:
                distances[neighbour] = distance
                previous_nodes[neighbour] = current_node
                heapq.heappush(heap, (distance, neighbour))
    path: list[StationGroupId] = []
    current = end
    while current is not None:
        path.append(current)
        current = previous_nodes[current]
    path.reverse()
    return path

def main():

    stations: dict[StationGroupId, Station] = {}
    station_group_map: dict[StationId, StationGroupId] = {}

    for line in [*Path('.').glob('station*.csv')][0].read_text().strip().splitlines()[1:]:
        (
            station_cd,
            station_g_cd,
            station_name,
            _station_name_k,
            _station_name_r,
            _line_cd,
            _pref_cd,
            _post,
            _address,
            lon,
            lat,
            _open_ymd,
            _close_ymd,
            _e_status,
            _e_sort
        ) = line.split(',')

        station = Station(
            StationId(int(station_cd)),
            StationGroupId(int(station_g_cd)),
            station_name,
            float(lon),
            float(lat)
        )
        
        stations[StationGroupId(int(station_g_cd))] = station

        station_group_map[StationId(int(station_cd))] = StationGroupId(int(station_g_cd))

    graph: dict[StationGroupId, list[tuple[StationGroupId, float]]] = {}

    for line in [*Path('.').glob('join*.csv')][0].read_text().strip().splitlines()[1:]:

        (
            _line_cd,
            station_cd1,
            station_cd2
        ) = line.split(',')

        if len(station_cd1) == 6 or len(station_cd2) == 6:
            continue
        
        station1 = stations[station_group_map[StationId(int(station_cd1))]]
        station2 = stations[station_group_map[StationId(int(station_cd2))]]

        dist = compute_distance(station1, station2)

        if station1.group_id not in graph:
            graph[station1.group_id] = []
        
        graph[station1.group_id].append((station2.group_id, dist))

        if station2.group_id not in graph:
            graph[station2.group_id] = []
        
        graph[station2.group_id].append((station1.group_id, dist))
    
    straight_distances: dict[tuple[StationGroupId, StationGroupId], float] = {}

    for source_station_group_id, destination_station_group_id in tqdm(itertools.combinations(stations.keys(), r=2), total=len(stations) * (len(stations) - 1) // 2):
        source_station = stations[source_station_group_id]
        destination_station = stations[destination_station_group_id]
        dist = compute_distance(source_station, destination_station)
        straight_distances[(source_station_group_id, destination_station_group_id)] = dist
        straight_distances[(destination_station_group_id, source_station_group_id)] = dist

    max_score = 0.0
    max_pair = (StationGroupId(0), StationGroupId(0))
    max_distances = (0.0, 0.0)

    for source_station_group_id, _ in tqdm(stations.items()):
        if source_station_group_id not in graph:
            continue
        max_score_for_station = 0.0
        max_dest = StationGroupId(0)
        max_distances_for_station = (0.0, 0.0)
        for destination_station_group_id, route_distance in list_route_distances(graph, source_station_group_id).items():
            if source_station_group_id == destination_station_group_id:
                continue
            if isinf(route_distance):
                continue
            straight_distance = straight_distances[(destination_station_group_id, source_station_group_id)]
            score = route_distance / straight_distance
            if score > max_score_for_station:
                max_score_for_station = score
                max_dest = destination_station_group_id
                max_distances_for_station = (route_distance, straight_distance)
        if max_score_for_station > max_score:
            max_score = max_score_for_station
            max_pair = (source_station_group_id, max_dest)
            max_distances = max_distances_for_station
    
    print(stations[max_pair[0]].name, stations[max_pair[1]].name)
    print(max_distances)

    path = compute_route(graph, max_pair[0], max_pair[1])

    print(' -> '.join([ stations[station_group_id].name for station_group_id in path ]))
    
if __name__ == "__main__":
    main()