import osmnx as ox

# Загрузка данных о дорогах Москвы
G = ox.graph.graph_from_place("Moscow", network_type="drive")

# Отображение дорог на карте
moscow_gdf = ox.geocoder.geocode_to_gdf("Moscow")
fig, ax = ox.plot.plot_graph(G, show=False, close=False, bgcolor="#111111", edge_color="#ffcb00", edge_linewidth=0.3, node_size=0)
moscow_gdf.plot(ax=ax, fc="#444444", ec=None, lw=1, alpha=1, zorder=-1)

# Настройка границ карты
margin = 0.02
west, south, east, north = moscow_gdf.union_all().bounds
margin_ns = (north - south) * margin
margin_ew = (east - west) * margin
ax.set_ylim((south - margin_ns, north + margin_ns))
ax.set_xlim((west - margin_ew, east + margin_ew))
plt.show()

with open('Moscow.geojson', 'w') as file:
    file.write(moscow_gdf.to_json())

nodes = G.nodes(data=True)
with open('nodes.csv', 'a') as file:
    file.write("id,lat,lonn")
    for (node, data) in nodes:
        file.write("%d,%f,%fn" % (node, data.get("y"), data.get("x")))

edges = G.edges(data=True)
def decode_maxspeed(maxspeed):
    match maxspeed:
        case str():
            match maxspeed.lower():
                case "ru:urban": return 60
                case "ru:rural": return 90
                case "ru:living_street": return 20
                case "ru:motorway": return 110
                case _: return int(maxspeed)
        case list(): return min(list(map(decode_maxspeed, maxspeed)))
        case _: return maxspeed

with open('edges.csv', 'a') as file:
    file.write("src,dst,maxspeed,length,geometryn")
    for (src, dst, data) in edges:
        maxspeed = decode_maxspeed(data.get("maxspeed", 999))
        length = float(data.get("length"))
        geometry = shapely.wkt.dumps(data.get("geometry"))
        file.write("%d,%d,%d,%f,%sn" % (src, dst, maxspeed, length, geometry))

from pyspark.sql import SparkSession

spark = SparkSession.builder 
    .config("spark.jars.packages", "graphframes:graphframes:0.8.4-spark3.5-s_2.12") 
    .master("local[*]") 
    .appName("GraphFrames") 
    .getOrCreate()

nodes = spark.read.options(header=True).csv("nodes.csv")
edges = spark.read.options(header=True).csv("edges.csv")

# Вычисление времени прохождения рёбер
edgesT = edges.withColumn("time", edges["length"] / edges["maxspeed"])

# Построение графа
from graphframes import *

g = GraphFrame(nodes, edgesT)

src = "257601812"
dst = "5840593081"

paths = g.shortestPaths(landmarks=[dst])
paths.filter(F.col("id") == src).show(truncate=False)