import geopandas as gpd
import fiona
from pybind11_geobuf import geojson

output = geojson.FeatureCollection()
gpkg_path = 'data.gpkg'

layers = fiona.listlayers(gpkg_path)
for layer in layers:
    if 'style' in layer:
        continue
    gdf = gpd.read_file(gpkg_path, layer=layer)
    gdf.set_crs(epsg=3857, inplace=True)
    gdf_wgs84 = gdf.to_crs(epsg=4326)
    opath = f'output_{layer}.json'
    gdf_wgs84.to_file(opath, driven='GeoJSON')
    fc = geojson.FeatureCollection().load(opath)
    output.extend(fc)
output.dump('output.json', indent=True)

import numpy as np
import fiona
import geopandas as gpd
from pybind11_geobuf import geojson, rapidjson

output = geojson.FeatureCollection()
gpkg_path = 'data.gpkg'

gdf = gpd.read_file(gpkg_path, layer='centerline')
gdf.set_crs(epsg=3857, inplace=True)
gdf_wgs84 = gdf.to_crs(epsg=4326)
centerline = geojson.FeatureCollection().from_rapidjson(rapidjson().loads(gdf_wgs84.to_json()))[0]
# print(centerline.as_numpy())

gdf = gpd.read_file(gpkg_path, layer='transects')
gdf.set_crs(epsg=3857, inplace=True)
gdf_wgs84 = gdf.to_crs(epsg=4326)
transects = geojson.FeatureCollection().from_rapidjson(rapidjson().loads(gdf_wgs84.to_json()))

from polyline_ruler import PolylineRuler
river = PolylineRuler(centerline.as_numpy(), is_wgs84=True)
centers = []
offsets = []
for t in transects:
    lla = t.as_numpy().mean(axis=0)
    hit, seg_idx, t = river.pointOnLine(lla)
    centers.append(hit)
    offsets.append(river.range(segment_index=seg_idx, t=t))
index = np.argsort(offsets)

polygon_original = []
polygon_updated1 = []
for idx in index:
    polygon_original.extend(transects[idx].as_numpy())
    # use river lateral direction (via scanline)
    # see https://github.com/cubao/index/blob/master/docs/notebooks/polyline-ruler.ipynb
    # right, left = river.scanline(offsets[idx], min=-250.0, max=250.0)

    # or use original transect direction
    from pybind11_geobuf import tf
    enu_dir = tf.lla2enu(transects[idx].as_numpy())[1]
    enu_dir /= np.linalg.norm(enu_dir)
    # I checked, transect is right to left direction:
    #       forward_dir = river.arrow(offsets[idx])[1]
    #       # (x, y) --ccw-rot90--> -y, x
    #       left_dir = [-forward_dir[1], forward_dir[0], 0.0]
    #       enu_dir @ left_dir > 0
    enus = enu_dir * np.array([-250.0, 250.0])[:, np.newaxis]
    right, left = tf.enu2lla(enus, anchor_lla=centers[idx])
    polygon_updated1.extend([right, left])

    # update gpkg
    from shapely.geometry import LineString
    gdf_wgs84.at[idx, 'geometry'] = LineString([right, left])

polygon_original = np.array([
    *polygon_original[::2],
    *polygon_original[1::2][::-1],
    polygon_original[0],
])
polygon_updated1 = np.array([
    *polygon_updated1[::2],
    *polygon_updated1[1::2][::-1],
    polygon_updated1[0],
])
f0 = geojson.Feature().geometry(geojson.Polygon(polygon_original)).properties({'stroke': '#ffff00', 'title': 'original'})
# print(f0.to_rapidjson().dumps()) # paste to geojson.io
f1 = geojson.Feature().geometry(geojson.Polygon(polygon_updated1)).properties({'stroke': '#00ff00', 'title': 'updated1'})
# print(f1.to_rapidjson().dumps())

output = geojson.FeatureCollection()
output.extend([f0,f1])
# print(output.to_rapidjson().dumps())
output.dump('river_polygons.json', indent=True)

gdf_updated = gdf_wgs84.to_crs(epsg=3857)
gdf_updated.to_file('updated_transects.gpkg', layer='transects', driver='GPKG')