from matplotlib import pyplot as plt
from matplotlib.path import Path
from matplotlib.patches import PathPatch

import numpy
import geopandas as gpd
import shapely
from shapely import LineString
from shapely.wkt import loads
import cartagen as c4

# A simple road segment (LineString) as initiator
initial = LineString([(0, 0), (10, 0), (20, 1)])
final = LineString([(0, 1), (10, 1.5), (20, 2)])

crs = 3857
initiator_initial_gdf = gpd.GeoDataFrame([{'name': 'init', 'geometry': initial}], crs=crs)
initiator_final_gdf = gpd.GeoDataFrame([{'name': 'init', 'geometry': final}], crs=crs)

# Movable objects (e.g., a nearby boundary line)
movable_geom = LineString([(0, 3), (10, 3), (20, 4)])
objects = gpd.GeoDataFrame([{'name': 'movable_line', 'geometry': movable_geom}], crs=crs)

# Frozen objects (None for this simple case)
frozen = gpd.GeoDataFrame(geometry=[])

# --- Algorithm Inputs ---
# Distance of propagation (SizePZ in the paper) [cite: 132]
distance = 5.0 
# Interval for initial vector computation [cite: 111]
interval = 2.0 
vectors_gdf, max_dep = c4.interpolate_displacement_vectors(initial, final, interval, crs=crs)

# propagate the displacement to the movable objects
propagated_gdf = c4.propagation_crow_flies(objects, initial, final, frozen, distance, interval)

fig = plt.figure(1, (10, 10))

#############################################################

sub1 = fig.add_subplot(111)
sub1.set_aspect('equal')
sub1.axes.get_xaxis().set_visible(False)
sub1.axes.get_yaxis().set_visible(False)

for line in objects.geometry:
    path1 = Path(numpy.asarray(line.coords)[:, :2])
    sub1.add_patch(PathPatch(path1, facecolor="none", edgecolor='blue', linewidth=1))

for line in propagated_gdf.geometry:
    path1 = Path(numpy.asarray(line.coords)[:, :2])
    sub1.add_patch(PathPatch(path1, facecolor="none", edgecolor='red', linewidth=1))

for line in initiator_initial_gdf.geometry:
    path1 = Path(numpy.asarray(line.coords)[:, :2])
    sub1.add_patch(PathPatch(path1, facecolor="none", edgecolor='blue', linewidth=1))

for line in initiator_final_gdf.geometry:
    path1 = Path(numpy.asarray(line.coords)[:, :2])
    sub1.add_patch(PathPatch(path1, facecolor="none", edgecolor='red', linewidth=1))

for v_idx, vector_row in vectors_gdf.iterrows():
    vect_geom = vector_row['geometry']
    if vect_geom.length > 0:
        sub1.annotate('', xy=(vect_geom.coords[1][0], vect_geom.coords[1][1]), 
                        xytext=(vect_geom.coords[0][0], vect_geom.coords[0][1]),
                        arrowprops=dict(arrowstyle='->', color='red'))

sub1.autoscale_view()
plt.show()