#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ornl_bridges_to_fast_format_20260807.py

Created by: Andy Carter, PE
2026.08.07 -- Assitance from Claude Opus 4.8

Combined pipeline built from four notebooks (run in order):

    part1  -> build bridge lines + opening area / slope / transitions + xsec JSON
              writes: 01_bridge_ornl_20260628_ln_4326.fgb
                      02_snbi_txdot_pnt_4326.fgb
                      03_streams_ornl_ln_4326.fgb
    part2  -> flip cross sections so station 0 is consistent looking downstream
              writes: 04_bridge_flipped_on_ornl_ln_4326.fgb
    part3  -> assign NBI number from nearest SNBI point where missing (<= 40 ft)
              writes: 05_bridge_lines_snbi_filled_4326.fgb
    part4  -> build final bridge data tables / rating curves / per-bridge JSON
              writes: 06_bridge_lines_for_accounting_4326.fgb
                      07_for_FAST_database/07_01_t_bridge_rating_curve.parquet
                      07_for_FAST_database/07_02_s_bridge_pnt_4326.fgb
                      07_for_FAST_database/07_03_json_per_bridge/<uuid>.json

Each stage reads the output of the previous stage from disk, so the four
run_* functions run sequentially at the bottom.
"""

import os
import ast
import json
import uuid
from io import BytesIO

import numpy as np
import pandas as pd
import geopandas as gpd
import requests

from scipy.interpolate import interp1d
from scipy.integrate import trapezoid

from shapely.geometry import LineString
from shapely.geometry import LineString as ShapelyLine
from tqdm import tqdm

import matplotlib.pyplot as plt  # used only by the optional plotting helper


# ===========================================================================
# Config (consolidated from all four notebooks)
# ===========================================================================

# Per-district state-plane CRS (US survey foot zones).
list_epsg_txdot_dist = {
    '1_PAR':  'EPSG:2276',  # Texas North Central
    '2_FTW':  'EPSG:2276',  # Texas North Central
    '3_WFS':  'EPSG:2276',  # Texas North Central
    '4_AMA':  'EPSG:2275',  # Texas North
    '5_LBB':  'EPSG:2276',  # Texas North Central
    '6_ODA':  'EPSG:2277',  # Texas Central
    '7_SJT':  'EPSG:2277',  # Texas Central
    '8_ABL':  'EPSG:2276',  # Texas North Central
    '9_WAC':  'EPSG:2277',  # Texas Central
    '10_TYL': 'EPSG:2277',  # Texas Central
    '11_LFK': 'EPSG:2277',  # Texas Central
    '12_HOU': 'EPSG:2278',  # Texas South Central
    '13_YKM': 'EPSG:2278',  # Texas South Central
    '14_AUS': 'EPSG:2277',  # Texas Central
    '15_SAT': 'EPSG:2278',  # Texas South Central
    '16_CRP': 'EPSG:2279',  # Texas South
    '17_BRY': 'EPSG:2278',  # Texas South Central *
    '18_DAL': 'EPSG:2276',  # Texas North Central
    '19_ATL': 'EPSG:2276',  # Texas North Central
    '20_BMT': 'EPSG:2278',  # Texas South Central
    '21_PHR': 'EPSG:2279',  # Texas South
    '22_LRD': 'EPSG:2279',  # Texas South *
    '23_BWD': 'EPSG:2277',  # Texas Central
    '24_ELP': 'EPSG:2277',  # Texas Central
    '25_CHS': 'EPSG:2276',  # Texas North Central *
}

# Which districts to process. The original notebooks looped over the full
# dict above (the ['2_FTW'] "interim testing" lists were commented out).
# To restrict to a subset, e.g. only Fort Worth, use:
DICT_DISTRICTS_TO_PROCESS = {'12_HOU': 'EPSG:2278'}

#DICT_DISTRICTS_TO_PROCESS = list_epsg_txdot_dist

# --- Folders ---------------------------------------------------------------
STR_INPUT_FILEPATH = r"E:\bridge_conversions_20260805\input"
STR_OUTPUT_FOLDER_PATH = r"E:\bridge_conversions_20260805\output_20260807"

# --- Part 1 web source -----------------------------------------------------
# It must be the base URL/path such that (for each district)
#   f"{STR_WEB_ROOT}{str_dist}/bridge_profiles_bareearth_snbi_osm_nwm_matching.gpkg"

STR_WEB_ROOT = 'https://web.corral.tacc.utexas.edu/nfiedata/sf3/alpha10/'

# --- Part 1 input files ----------------------------------------------------
str_bridge_thickness_filepath = os.path.join(
    STR_INPUT_FILEPATH, 'nbi_bridges_texas_4326.parquet'
)
str_ornl_stream_filepath = os.path.join(
    STR_INPUT_FILEPATH,
    'demDerived_reaches_split_filtered_addedAttributes_crosswalked_3081.fgb'
)

# --- Part 3 SNBI points ----------------------------------------------------
# SNBI downloaded 2026.08.06 from
# https://gis-txdot.opendata.arcgis.com/datasets/TXDOT::txdot-bridges-snbi/about
STR_SNBI_POINTS = os.path.join(
    STR_INPUT_FILEPATH, 'snbi_from_TxDOT_20260806_pnt_4326.fgb'
)
# bridge-number column in the SNBI POINTS file -- verify this is correct
STR_SNBI_ID_FIELD = 'ID01_BRDG_NBR'
# distance to find an SNBI point near a bridge line that currently has no NBI
FLT_SEARCH_DIST_FT = 40

# --- Part 4 input files ----------------------------------------------------
    
STR_FILEPATH_TO_HAND_RATING_CURVES = os.path.join(
    STR_INPUT_FILEPATH, 'hydroTable_rp_bf_lmtdischarge_cda.parquet'
)

STR_T_NEXTGEN_TO_NWM_FILEPATH = os.path.join(
    STR_INPUT_FILEPATH, 't_nextgen_to_nwm.parquet'
)
 
# (Alternate values from the first assignment, for reference:)
#   os.path.join(STR_INPUT_FILEPATH, 'hydroTable_rp_bf_lmtdischarge_cda.parquet')
#   os.path.join(STR_INPUT_FILEPATH, 't_nextgen_to_nwm.parquet')

# --- Constants -------------------------------------------------------------
FLT_MAX_GRADE = 0.11               # 11% maximum grade
M_TO_FT = 3.28084
CMS_TO_CFS = 35.3147
FLT_MIN_OPENING_SQ_FT = 100        # minimum area under bridge allowed

FLT_DEFAULT_BRIDGE_THICK_STEP_4 = 3.14 # default bridge thickness in feet if one isn't provided
FLT_MAX_BRIDGE_THICK_STEP_4 = 7.0 # max allowed bridge thickness in feet

# US survey foot -> meter (EPSG:2275-2279 are ftUS zones)
FLT_USFT_TO_M = 1200.0 / 3937.0


# ===========================================================================
# Part 1 -- helpers
# ===========================================================================

def fn_convert_z_to_feet(geom):
    new_coords = [(x, y, z * M_TO_FT) for x, y, z in geom.coords]
    return LineString(new_coords)


def _profile(coords):
    xy = np.asarray([(c[0], c[1]) for c in coords], float)
    d = np.concatenate([[0.0], np.cumsum(np.hypot(np.diff(xy[:, 0]), np.diff(xy[:, 1])))])
    z = np.asarray([c[2] for c in coords], float)
    return d, z


def _intersect(dist, za, zb):
    inter = ShapelyLine(zip(dist, za)).intersection(ShapelyLine(zip(dist, zb)))
    if inter.is_empty:
        return [], []
    if inter.geom_type == "Point":
        return [inter.x], [inter.y]
    if inter.geom_type in ("MultiPoint", "GeometryCollection"):
        pts = [g for g in inter.geoms if g.geom_type == "Point"]
        return [g.x for g in pts], [g.y for g in pts]
    return [], []


def _r(a, nd=3):
    return [round(float(v), nd) for v in a]


def build_xsec(row_be, gdf_road, df_areas):
    bridge_id = row_be["bridge_id"]
    matches = gdf_road[gdf_road["bridge_id"] == bridge_id]
    if len(matches) == 0:
        return None
    row_rd = matches.iloc[0]
    bridge_thi = float(row_be["bridge_thickness"])

    dist_be, z_be = _profile(list(row_be.geometry.coords))
    dist_rd, z_rd = _profile(list(row_rd.geometry.coords))
    z_low_chord = z_rd - bridge_thi

    lo = max(dist_be[0], dist_rd[0]); hi = min(dist_be[-1], dist_rd[-1])
    dist_common = np.union1d(dist_be, dist_rd)
    dist_common = dist_common[(dist_common >= lo) & (dist_common <= hi)]
    if dist_common.size < 2:
        return None

    z_be_i = interp1d(dist_be, z_be, bounds_error=False, fill_value="extrapolate")(dist_common)
    z_lc_i = interp1d(dist_rd, z_low_chord, bounds_error=False, fill_value="extrapolate")(dist_common)
    z_rd_i = interp1d(dist_rd, z_rd, bounds_error=False, fill_value="extrapolate")(dist_common)

    # low-chord (combo) line: max(be, low_chord), with be x lowchord intersections inserted
    z_combo = np.maximum(z_be_i, z_lc_i)
    ix, iy = _intersect(dist_common, z_be_i, z_lc_i)
    combo = sorted(list(zip(dist_common, z_combo)) + list(zip(ix, iy)), key=lambda p: p[0])
    dist_combo = np.array([p[0] for p in combo]); z_combo_final = np.array([p[1] for p in combo])

    # road-surface (max envelope) line, with be x road intersections inserted
    z_max = np.maximum(z_be_i, z_rd_i)
    ix2, iy2 = _intersect(dist_common, z_be_i, z_rd_i)
    mx = sorted(list(zip(dist_common, z_max)) + list(zip(ix2, iy2)), key=lambda p: p[0])
    dist_max = np.array([p[0] for p in mx]); z_max_final = np.array([p[1] for p in mx])

    # gray bridge-structure band: road surface interpolated onto dist_combo vs the low-chord line
    z_rd_shade = interp1d(dist_rd, z_rd, bounds_error=False, fill_value="extrapolate")(dist_combo)

    oa = df_areas.loc[df_areas["bridge_id"] == bridge_id, "opening_area_sqft"]
    opening_area = float(oa.values[0]) if len(oa) else None

    return {
        "bridge_id": str(bridge_id),
        "brdg_nbr": str(row_be.get("ID01_BRDG_NBR", "") or ""),
        "bridge_thi": round(bridge_thi, 3),
        "opening_area": opening_area,
        "ground":   {"d": _r(dist_be),    "z": _r(z_be)},
        "lowchord": {"d": _r(dist_combo), "z": _r(z_combo_final)},
        "roadsurf": {"d": _r(dist_max),   "z": _r(z_max_final)},
        "band":     {"d": _r(dist_combo), "z_top": _r(z_rd_shade), "z_bot": _r(z_combo_final)},
        "trans":    {"x": _r(ix),         "y": _r(iy)},
    }


# ===========================================================================
# Part 1 -- run
# ===========================================================================

def run_part1():
    """Build bridge lines, opening area / slope / transition metrics, xsec JSON."""
    int_total_bridge = 0  # noqa: F841  (kept from original)

    print('  -- Reading ORNL stream lines')

    # Load Texas hydrofabric (Yan's Texas 3 meter) streamlines
    gdf_streams = gpd.read_file(str_ornl_stream_filepath)

    # (loaded in the original; kept for parity though unused below)
    df_nbi_thickness_lookup = pd.read_parquet(str_bridge_thickness_filepath)  # noqa: F841

    # --- Load the thickness table and build a lookup on ID01_BRDG_NBR ---
    df_thickness = pd.read_parquet(str_bridge_thickness_filepath)

    # Keep one thickness per bridge number (guard against dupes) and drop nulls
    df_thick_lookup = (
        df_thickness[['ID01_BRDG_NBR', 'bridge_thickness']]
        .dropna(subset=['ID01_BRDG_NBR', 'bridge_thickness'])
        .drop_duplicates(subset='ID01_BRDG_NBR', keep='first')
        .set_index('ID01_BRDG_NBR')['bridge_thickness']
    )

    for str_dist, STR_DIST_PROJECTION_EPSG in DICT_DISTRICTS_TO_PROCESS.items():
        print('----------------')

        # ##############################
        # Step 00
        print('  -- Step 0: Loading web data')
        str_bare_earth_path = f"{STR_WEB_ROOT}{str_dist}/bridge_profiles_bareearth_snbi_osm_nwm_matching.gpkg"
        str_road_profile_path = f"{STR_WEB_ROOT}{str_dist}/bridge_profiles_surface_snbi_osm_nwm_matching.gpkg"

        str_bridge_attr_path = f"{STR_WEB_ROOT}{str_dist}/bridge_profiles_attr_snbi_osm_nwm_matching.json"

        gdf_bare_earth = gpd.read_file(str_bare_earth_path)
        gdf_road_profile = gpd.read_file(str_road_profile_path)

        data = requests.get(str_bridge_attr_path).json()

        int_current_bridge_count = len(gdf_bare_earth)
        print(f'    -- District {str_dist}: Bridge lines found: {int_current_bridge_count}')

        # ##############################
        # Step 01 -- determine streams that cross ORNL streams
        print('  -- Step 1: Determine bridges that cross ORNL streams')
        records = []
        for item in data:
            snbi = item.get("attr_snbi", {})
            if snbi.get("OVER_WATER") == 1:
                record = {"bridge_id": item["bridge_id"]}
                record.update(snbi)
                records.append(record)
        df_snbi = pd.DataFrame(records)
        gdf_conflates_snbi = gdf_bare_earth.merge(
            df_snbi[["bridge_id", "ID01_BRDG_NBR", "distance2nearest_snbi"]],
            on="bridge_id",
            how="left"
        )
        # Append a 'has_snbi'
        gdf_conflates_snbi["has_snbi"] = gdf_conflates_snbi["distance2nearest_snbi"].notna()

        # Find the bridges that intersect Yan's 3 meter hydrofabric
        gdf_intersects = gpd.sjoin(
            gdf_bare_earth,
            gdf_streams[["HydroID", "feature_id", "order_", "geometry"]],
            how="inner",
            predicate="intersects"
        )

        # If a bridge line intersects multiple streams (Yan's 3 meter),
        # keep the one with the highest order_ (first found if tied)
        gdf_intersects = (
            gdf_intersects
            .sort_values("order_", ascending=False)
            .groupby(level=0)
            .first()
            .reset_index()
        )

        # Merge the intersecting stream's HydroID onto the conflation table
        gdf_conflates_snbi = gdf_conflates_snbi.merge(
            gdf_intersects[["bridge_id", "HydroID"]],
            on="bridge_id",
            how="left"
        )

        # Keep HydroID as a nullable integer (merge upcasts to float due to NaNs for non-matches)
        gdf_conflates_snbi["HydroID"] = gdf_conflates_snbi["HydroID"].astype("Int64")

        # Append boolean flag: a bridge intersects an ORNL stream iff it got a HydroID
        gdf_conflates_snbi["intersects_ORNL_stream"] = gdf_conflates_snbi["HydroID"].notna()
        print(f"    -- Bare earth profiles intersecting a ORNL streamline: "
              f"{len(gdf_intersects)} / {len(gdf_bare_earth)}")

        # ##############################
        # Step 02 -- Compute Bridge Opening area
        print('  -- Step 2: Compute bridge opening area')

        gdf_conflates_snbi_stateplane = gdf_conflates_snbi.to_crs(STR_DIST_PROJECTION_EPSG)
        gdf_conflates_snbi_stateplane["geometry"] = \
            gdf_conflates_snbi_stateplane["geometry"].apply(fn_convert_z_to_feet)

        # Filter gdf_road_profile to matching bridge_ids
        matching_ids = gdf_conflates_snbi_stateplane["bridge_id"]
        gdf_road_matched = gdf_road_profile[gdf_road_profile["bridge_id"].isin(matching_ids)].copy()

        # Reproject x,y to district CRS (feet)
        gdf_road_matched_stateplane = gdf_road_matched.set_crs("EPSG:3081").to_crs(STR_DIST_PROJECTION_EPSG)

        # Convert Z from meters to feet
        gdf_road_matched_stateplane["geometry"] = \
            gdf_road_matched_stateplane["geometry"].apply(fn_convert_z_to_feet)

        # gdf_conflates_snbi is the bare earth profile
        gdf_conflates_snbi_stateplane['bridge_thickness'] = FLT_DEFAULT_BRIDGE_THICK_STEP_4

        # --- Map thickness onto the conflated bridges by ID01_BRDG_NBR ---
        mapped_thickness = gdf_conflates_snbi_stateplane['ID01_BRDG_NBR'].map(df_thick_lookup)

        # Fallback to default where there is no match, then cap at the max
        gdf_conflates_snbi_stateplane['bridge_thickness'] = (
            mapped_thickness
            .fillna(FLT_DEFAULT_BRIDGE_THICK_STEP_4)
            .clip(upper=FLT_MAX_BRIDGE_THICK_STEP_4)
        )

        areas = []

        for idx in range(len(gdf_conflates_snbi_stateplane)):

            row_be = gdf_conflates_snbi_stateplane.iloc[idx]
            bridge_id = row_be["bridge_id"]

            # Find matching road profile
            row_rd_matches = gdf_road_matched_stateplane[
                gdf_road_matched_stateplane["bridge_id"] == bridge_id
            ]

            if row_rd_matches.empty:
                areas.append({
                    "idx": idx,
                    "bridge_id": bridge_id,
                    "opening_area_sqft": None
                })
                continue

            row_rd = row_rd_matches.iloc[0]
            bridge_thi = row_be["bridge_thickness"]

            # Ground profile (bare earth)
            coords_be = list(row_be.geometry.coords)
            dist_be = np.insert(
                np.cumsum(
                    np.sqrt(
                        np.diff([c[0] for c in coords_be])**2 +
                        np.diff([c[1] for c in coords_be])**2
                    )
                ),
                0,
                0
            )
            z_be = np.array([c[2] for c in coords_be])

            # Road/deck profile
            coords_rd = list(row_rd.geometry.coords)
            dist_rd = np.insert(
                np.cumsum(
                    np.sqrt(
                        np.diff([c[0] for c in coords_rd])**2 +
                        np.diff([c[1] for c in coords_rd])**2
                    )
                ),
                0,
                0
            )
            z_rd = np.array([c[2] for c in coords_rd])

            # Low chord elevation
            z_low_chord = z_rd - bridge_thi

            # Common distance domain
            dist_common = np.union1d(dist_be, dist_rd)
            dist_common = dist_common[
                (dist_common >= max(dist_be[0], dist_rd[0])) &
                (dist_common <= min(dist_be[-1], dist_rd[-1]))
            ]

            if len(dist_common) < 2:
                areas.append({
                    "idx": idx,
                    "bridge_id": bridge_id,
                    "opening_area_sqft": 0
                })
                continue

            # Interpolate profiles
            z_be_interp = interp1d(
                dist_be, z_be, bounds_error=False, fill_value="extrapolate"
            )(dist_common)

            z_lc_interp = interp1d(
                dist_rd, z_low_chord, bounds_error=False, fill_value="extrapolate"
            )(dist_common)

            # =================================================
            # SIGN CONVENTION
            #   Low chord below ground  -> negative
            #   Ground below low chord  -> positive
            # =================================================
            gap = z_lc_interp - z_be_interp

            # Signed area
            opening_area = trapezoid(gap, dist_common)

            areas.append({
                "idx": idx,
                "bridge_id": bridge_id,
                "opening_area_sqft": opening_area
            })

        # Create DataFrame
        df_areas = pd.DataFrame(areas)

        df_areas["opening_area_sqft"] = (
            pd.to_numeric(df_areas["opening_area_sqft"], errors="coerce")
            .round(0)
            .astype("Int64")
        )

        # Filter small / negative openings
        df_small_openings = df_areas[
            df_areas["opening_area_sqft"] < FLT_MIN_OPENING_SQ_FT
        ].copy()

        print(f"    -- Bridge opening less than {FLT_MIN_OPENING_SQ_FT} sqft: "
              f"{len(df_small_openings)}")

        # Convert areas list to DataFrame if needed
        df_open_area = pd.DataFrame(areas)

        # Keep only required fields
        df_open_area = df_open_area[["bridge_id", "opening_area_sqft"]]

        # Merge area back onto bridge GeoDataFrame
        gdf_conflates_snbi_stateplane = gdf_conflates_snbi_stateplane.merge(
            df_open_area, on="bridge_id", how="left"
        )

        # Create boolean flag
        gdf_conflates_snbi_stateplane["has_min_area"] = (
            gdf_conflates_snbi_stateplane["opening_area_sqft"] > FLT_MIN_OPENING_SQ_FT
        )

        # Optional: fill missing areas
        gdf_conflates_snbi_stateplane["opening_area_sqft"] = (
            gdf_conflates_snbi_stateplane["opening_area_sqft"].fillna(0).astype(int)
        )

        int_count_ORNL_and_min_opening = (
            (gdf_conflates_snbi_stateplane['intersects_ORNL_stream'] == True) &
            (gdf_conflates_snbi_stateplane['has_min_area'] == True)
        ).sum()

        print(f"    -- Intersects ORNL streamline with min area: "
              f"{int_count_ORNL_and_min_opening} / {len(gdf_bare_earth)}")

        # ##############################
        # Step 03 -- maximum road slope per bridge
        print('  -- Step 3: Determine maximum road slope per bridge')

        slope_counts = []
        for idx in range(len(gdf_conflates_snbi_stateplane)):
            row_be = gdf_conflates_snbi_stateplane.iloc[idx]
            bridge_id = row_be["bridge_id"]

            row_rd_matches = gdf_road_matched_stateplane[
                gdf_road_matched_stateplane["bridge_id"] == bridge_id
            ]
            if len(row_rd_matches) == 0:
                slope_counts.append({"idx": idx, "bridge_id": bridge_id,
                                     "n_steep_segments": None, "max_grade": None})
                continue

            row_rd = row_rd_matches.iloc[0]
            coords_rd = list(row_rd.geometry.coords)

            dist_rd = np.cumsum(np.sqrt(
                np.diff([c[0] for c in coords_rd])**2 +
                np.diff([c[1] for c in coords_rd])**2
            ))
            dist_rd = np.insert(dist_rd, 0, 0)
            z_rd = np.array([c[2] for c in coords_rd])

            delta_z = np.diff(z_rd)
            delta_d = np.diff(dist_rd)
            slopes = np.abs(delta_z / delta_d)

            n_steep = int((slopes > FLT_MAX_GRADE).sum())
            max_grade = round(float(slopes.max()), 4)

            slope_counts.append({"idx": idx, "bridge_id": bridge_id,
                                 "n_steep_segments": n_steep, "max_grade": max_grade})

        df_slope_counts = pd.DataFrame(slope_counts)
        df_slope_counts = df_slope_counts.sort_values(
            "n_steep_segments", ascending=False
        ).reset_index(drop=True)

        gdf_conflates_snbi_stateplane = gdf_conflates_snbi_stateplane.merge(
            df_slope_counts[["bridge_id", "max_grade"]],
            on="bridge_id",
            how="left",
        )

        gdf_conflates_snbi_stateplane["is_flat_road"] = ~(
            gdf_conflates_snbi_stateplane["max_grade"] > FLT_MAX_GRADE
        )

        gdf_conflates_snbi_stateplane["max_grade"] = \
            gdf_conflates_snbi_stateplane["max_grade"].round(3)

        int_count_ORNL_min_area_flat = (
            (gdf_conflates_snbi_stateplane['intersects_ORNL_stream'] == True) &
            (gdf_conflates_snbi_stateplane['has_min_area'] == True) &
            (gdf_conflates_snbi_stateplane['is_flat_road'] == True)
        ).sum()

        print(f"    -- ORNL-min area-flat road: "
              f"{int_count_ORNL_min_area_flat} / {len(gdf_bare_earth)}")

        # ##############################
        # Step 04 -- abutment / transition points
        print('  -- Step 4: Determine abutment points')

        transition_counts = []
        for idx in range(len(gdf_conflates_snbi_stateplane)):
            row_be = gdf_conflates_snbi_stateplane.iloc[idx]
            bridge_id = row_be["bridge_id"]

            row_rd_matches = gdf_road_matched_stateplane[
                gdf_road_matched_stateplane["bridge_id"] == bridge_id
            ]
            if row_rd_matches.empty:
                transition_counts.append({"idx": idx, "bridge_id": bridge_id,
                                          "n_transition_pts": None})
                continue

            row_rd = row_rd_matches.iloc[0]
            bridge_thi = row_be["bridge_thickness"]

            # Bridge (ground) profile
            coords_be = list(row_be.geometry.coords)
            dist_be = np.insert(
                np.cumsum(np.sqrt(
                    np.diff([c[0] for c in coords_be])**2 +
                    np.diff([c[1] for c in coords_be])**2
                )),
                0, 0
            )
            z_be = np.array([c[2] for c in coords_be])

            # Road profile
            coords_rd = list(row_rd.geometry.coords)
            dist_rd = np.insert(
                np.cumsum(np.sqrt(
                    np.diff([c[0] for c in coords_rd])**2 +
                    np.diff([c[1] for c in coords_rd])**2
                )),
                0, 0
            )
            z_rd = np.array([c[2] for c in coords_rd])

            # Low chord elevation
            z_low_chord = z_rd - bridge_thi

            # Common distance domain
            dist_common = np.union1d(dist_be, dist_rd)
            dist_common = dist_common[
                (dist_common >= max(dist_be[0], dist_rd[0])) &
                (dist_common <= min(dist_be[-1], dist_rd[-1]))
            ]

            # Interpolate onto common domain
            z_be_interp = interp1d(
                dist_be, z_be, bounds_error=False, fill_value="extrapolate"
            )(dist_common)

            z_lc_interp = interp1d(
                dist_rd, z_low_chord, bounds_error=False, fill_value="extrapolate"
            )(dist_common)

            # Compute gap and count sign changes (transitions)
            gap = z_lc_interp - z_be_interp
            sign = np.sign(gap)

            # Avoid double counting exact zeros
            sign[sign == 0] = 1

            n_pts = np.sum(sign[:-1] * sign[1:] < 0)

            transition_counts.append({"idx": idx, "bridge_id": bridge_id,
                                      "n_transition_pts": int(n_pts)})

        df_transition_counts = pd.DataFrame(transition_counts)
        df_transition_counts["n_transition_pts"] = (
            pd.to_numeric(df_transition_counts["n_transition_pts"], errors="coerce")
            .astype("Int64")
        )

        gdf_conflates_snbi_stateplane = gdf_conflates_snbi_stateplane.merge(
            df_transition_counts[["bridge_id", "n_transition_pts"]],
            on="bridge_id",
            how="left",
        )

        gdf_conflates_snbi_stateplane["n_transition_pts"] = (
            gdf_conflates_snbi_stateplane["n_transition_pts"].fillna(0).astype(int)
        )

        gdf_conflates_snbi_stateplane["has_two_transition"] = (
            gdf_conflates_snbi_stateplane["n_transition_pts"] == 2
        )

        int_count_ORNL_min_area_flat_two_trans = (
            (gdf_conflates_snbi_stateplane['intersects_ORNL_stream'] == True) &
            (gdf_conflates_snbi_stateplane['has_min_area'] == True) &
            (gdf_conflates_snbi_stateplane['is_flat_road'] == True) &
            (gdf_conflates_snbi_stateplane['has_two_transition'] == True)
        ).sum()

        print(f"    -- ORNL-min area-flat road-two trans: "
              f"{int_count_ORNL_min_area_flat_two_trans} / {len(gdf_bare_earth)}")

        # ##############################
        # Step 05 -- create bridge lines output
        print('  -- Step 5: Creating district ORNL Bridge line FGB')

        xsec_col = []
        for idx in range(len(gdf_conflates_snbi_stateplane)):
            row = gdf_conflates_snbi_stateplane.iloc[idx]
            try:
                d = build_xsec(row, gdf_road_matched_stateplane, df_areas)
            except Exception as e:
                print(f"[{idx}] xsec failed: {e}")
                d = None
            xsec_col.append(json.dumps(d, separators=(",", ":")) if d is not None else None)

        gdf_out = gdf_conflates_snbi_stateplane.copy()
        gdf_out["xsec"] = xsec_col

        # Reproject to WGS84 and export
        gdf_out = gdf_out.to_crs(epsg=4326)

        out_folder = os.path.join(STR_OUTPUT_FOLDER_PATH, str_dist)
        os.makedirs(out_folder, exist_ok=True)

        str_filename = '01_bridge_ornl_20260628_ln_4326.fgb'
        out_path = os.path.join(out_folder, str_filename)
        gdf_out.to_file(out_path, driver="FlatGeobuf")

        n_ok = sum(v is not None for v in xsec_col)
        print(f"    -- Wrote {out_path}; {n_ok} with cross-section data.")

        # ##############################
        # Step 06 -- create snbi point output
        print('  -- Step 6: Copying SNBI data')

        str_snbi_path = f"{STR_WEB_ROOT}{str_dist}/snbi.gpkg"
        out_path_snbi = os.path.join(out_folder, "02_snbi_txdot_pnt_4326.fgb")

        with requests.get(str_snbi_path, stream=True, timeout=60) as r:
            r.raise_for_status()
            data_snbi = BytesIO(r.content)

        gdf_snbi = gpd.read_file(data_snbi)
        gdf_snbi.to_file(out_path_snbi, driver="FlatGeobuf")

        print(f"    -- Wrote {out_path_snbi}")

        # ##############################
        # Step 07 -- Clip ORNL stream lines to district boundary
        print('  -- Step 7: Saving districts ORNL stream lines')

        str_district_boundary_path = f"{STR_WEB_ROOT}{str_dist}/District.gpkg"
        gdf_district_boundary = gpd.read_file(str_district_boundary_path)

        # Union + buffer once into a single geometry
        buffer_geom = gdf_district_boundary.geometry.unary_union.buffer(5000)

        # Spatial-index-backed filter: keep whole stream lines that touch the buffer
        gdf_ORNL_streams_4326 = gdf_streams[gdf_streams.intersects(buffer_geom)].to_crs(epsg=4326)

        out_path_ornl_streams = os.path.join(out_folder, "03_streams_ornl_ln_4326.fgb")
        gdf_ORNL_streams_4326.to_file(out_path_ornl_streams, driver="FlatGeobuf")

        print(f"    -- Wrote {out_path_ornl_streams}")


# ===========================================================================
# Part 2 -- helpers (flip cross sections)
# ===========================================================================

def fn_get_segment_at_intersection(line, other_line):
    """Get the direction vector of the segment of `line` closest to `other_line`."""
    intersection_pt = line.intersection(other_line)
    if intersection_pt.is_empty:
        # fallback: nearest point on `line` to `other_line`
        intersection_pt = line.interpolate(line.project(other_line.centroid))

    # If the intersection is a multi-part geometry (e.g. crosses twice),
    # use its centroid so .distance() below behaves as a point comparison.
    if intersection_pt.geom_type not in ("Point",):
        intersection_pt = intersection_pt.centroid

    coords = list(line.coords)
    if len(coords) < 2:
        # degenerate line; no direction can be derived
        return np.array([0.0, 0.0])

    min_dist = float("inf")
    seg_start, seg_end = coords[0], coords[1]

    for i in range(len(coords) - 1):
        seg = LineString([coords[i], coords[i + 1]])
        dist = seg.distance(intersection_pt)
        if dist < min_dist:
            min_dist = dist
            seg_start, seg_end = coords[i], coords[i + 1]

    return np.array(seg_end[:2]) - np.array(seg_start[:2])


def fn_get_cross_product_sign(bridge_line, stream_line):
    """
    Returns +1 if bridge line goes left-to-right looking downstream,
    -1 if right-to-left, 0 if parallel.
    Uses only the local intersecting segments of each line.
    """
    stream_dir = fn_get_segment_at_intersection(stream_line, bridge_line)
    bridge_dir = fn_get_segment_at_intersection(bridge_line, stream_line)
    # 2D cross product (z-component)
    cross_z = stream_dir[0] * bridge_dir[1] - stream_dir[1] * bridge_dir[0]
    return int(np.sign(cross_z))


def to_int64_hydroid(s):
    """Normalize HydroID to a common int64 type.
    Handles strings, floats, '24710733.0', and pandas nullable Int64."""
    return s.astype("float64").astype("int64")


def fn_lookup_sign(row, hydro_geom_lookup):
    """Compute cross product sign, tolerant of missing keys."""
    hydro_id = row["HydroID"]
    if hydro_id not in hydro_geom_lookup.index:
        return np.nan
    return fn_get_cross_product_sign(row["geometry"], hydro_geom_lookup.loc[hydro_id])


def fn_flip_bridge_cross_section(data):
    """
    Flip a bridge cross section so station 0 starts at the opposite end.
    Mirrors every station about the section width L (d -> L - d), re-sorts
    ascending, and carries each paired elevation array along. Elevations unchanged.
    Accepts a dict or a JSON string.

    Robust to empty or missing station lists: L is computed only from
    profiles that actually contain stations, and empty profiles are left
    untouched. If no profile has any stations, the data is returned unchanged.
    """
    if isinstance(data, str):
        data = json.loads(data)

    station_key = {
        "ground":   "d",
        "lowchord": "d",
        "roadsurf": "d",
        "band":     "d",
        "trans":    "x",
    }

    # Only consider profiles that (a) are present, (b) have the station key,
    # and (c) have a non-empty station list.
    present = {
        p: k for p, k in station_key.items()
        if p in data
        and isinstance(data[p].get(k), (list, tuple))
        and len(data[p][k]) > 0
    }

    # Nothing to flip - return an untouched deep copy.
    if not present:
        return json.loads(json.dumps(data))

    L = max(max(data[p][k]) for p, k in present.items())

    flipped = json.loads(json.dumps(data))  # deep copy, leave original intact
    for profile, s_key in present.items():
        prof = flipped[profile]
        mirrored = [L - s for s in prof[s_key]]
        order = sorted(range(len(mirrored)), key=lambda i: mirrored[i])
        prof[s_key] = [mirrored[i] for i in order]
        for k in prof:
            if k == s_key:
                continue
            # Only reorder paired arrays that match the station length.
            if isinstance(prof[k], (list, tuple)) and len(prof[k]) == len(order):
                prof[k] = [prof[k][i] for i in order]
    return flipped


def fn_flip_row(row):
    """
    For a bridge line that runs right-to-left looking downstream
    (cross_product_sign == 1), reverse the LineString geometry AND
    flip the paired cross-section so station 0 lines up with the new
    start vertex. Returns (geometry, xsec) as a tuple.
    """
    geom = row["geometry"]
    xsec = row["xsec"]

    # Reverse the vertex order of the line
    flipped_geom = LineString(list(geom.coords)[::-1])

    # Flip the cross section (handles dict or JSON string)
    flipped_xsec = fn_flip_bridge_cross_section(xsec)

    # Preserve the original storage type of 'xsec' (str vs dict)
    if isinstance(xsec, str):
        flipped_xsec = json.dumps(flipped_xsec)

    return flipped_geom, flipped_xsec


def fn_plot_bridge_cross_section(data, ax=None, show_trans=True):
    """
    Plot a bridge cross section: ground, bridge deck (band between low chord
    and road surface), and abutment transition points.

    `data` may be a dict or a JSON string. Returns the matplotlib Axes.
    (Optional / exploratory helper -- not called by the pipeline.)
    """
    if isinstance(data, str):
        data = json.loads(data)

    if ax is None:
        _, ax = plt.subplots(figsize=(12, 6))

    # Bridge deck: fill between low chord (bottom) and road surface (top)
    b = data["band"]
    ax.fill_between(b["d"], b["z_bot"], b["z_top"],
                    color="0.55", alpha=0.7, label="Bridge deck", zorder=2)

    # Natural ground, filled down to a floor below the channel
    g = data["ground"]
    ax.fill_between(g["d"], g["z"], min(g["z"]) - 2,
                    color="#c2a878", alpha=0.35, zorder=1)
    ax.plot(g["d"], g["z"], color="#7a5a2e", lw=2.0, label="Ground", zorder=3)

    # Road surface and low chord
    ax.plot(data["roadsurf"]["d"], data["roadsurf"]["z"],
            color="black", lw=1.8, label="Road surface", zorder=4)
    ax.plot(data["lowchord"]["d"], data["lowchord"]["z"],
            color="firebrick", lw=1.6, ls="--", label="Low chord", zorder=4)

    # Transition points (abutment edges of the opening)
    if show_trans and "trans" in data:
        t = data["trans"]
        ax.scatter(t["x"], t["y"], color="blue", s=55, zorder=5,
                   label="Transition pts", edgecolor="white", linewidth=0.8)

    ax.set_xlabel("Station along cross section, d")
    ax.set_ylabel("Elevation, z")
    ax.set_title(f"Bridge {data.get('brdg_nbr', data.get('bridge_id'))} - cross section")
    ax.legend(loc="lower center", ncol=5, frameon=True, fontsize=9)
    ax.grid(True, ls=":", alpha=0.5)
    ax.margins(x=0.01)
    return ax


# ===========================================================================
# Part 2 -- run
# ===========================================================================

def run_part2():
    """Flip bridge cross sections so station 0 is consistent looking downstream."""
    for str_dist, STR_DIST_PROJECTION_EPSG in DICT_DISTRICTS_TO_PROCESS.items():
        str_bridge_lines = os.path.join(STR_OUTPUT_FOLDER_PATH, str_dist, '01_bridge_ornl_20260628_ln_4326.fgb')
        str_snbi_points = os.path.join(STR_OUTPUT_FOLDER_PATH, str_dist, '02_snbi_txdot_pnt_4326.fgb')
        str_stream_lines = os.path.join(STR_OUTPUT_FOLDER_PATH, str_dist, '03_streams_ornl_ln_4326.fgb')

        gdf_bridge_lines = gpd.read_file(str_bridge_lines)
        gdf_snbi_points = gpd.read_file(str_snbi_points)  # noqa: F841 (kept from original)
        gdf_stream_lines = gpd.read_file(str_stream_lines)

        tqdm.pandas()

        # Streams side: cast, then build the HydroID -> geometry lookup
        gdf_stream_lines = gdf_stream_lines.copy()
        gdf_stream_lines["HydroID"] = to_int64_hydroid(gdf_stream_lines["HydroID"])

        hydro_geom_lookup = gdf_stream_lines.set_index("HydroID")["geometry"]

        # Collapse duplicate HydroIDs so .loc returns a single geometry, not a Series
        if not hydro_geom_lookup.index.is_unique:
            print("Warning: duplicate HydroIDs in streams - keeping first of each.")
            hydro_geom_lookup = hydro_geom_lookup[~hydro_geom_lookup.index.duplicated(keep="first")]

        # --- Bridges side ---
        gdf_intersects = gdf_bridge_lines[gdf_bridge_lines["intersects_ORNL_stream"] == True].copy()

        # Drop rows with null HydroID before casting (can't cast <NA> to int64)
        n_null = gdf_intersects["HydroID"].isna().sum()
        if n_null:
            print(f"Warning: dropping {n_null} rows with null HydroID.")
            gdf_intersects = gdf_intersects[gdf_intersects["HydroID"].notna()]

        gdf_intersects["HydroID"] = to_int64_hydroid(gdf_intersects["HydroID"])

        # --- Report any HydroIDs still missing from the lookup ---
        missing = set(gdf_intersects["HydroID"]) - set(hydro_geom_lookup.index)
        if missing:
            print(f"Warning: {len(missing)} HydroIDs not found in streams lookup; "
                  f"those rows will get NaN. Examples: {list(missing)[:5]}")

        gdf_intersects["cross_product_sign"] = gdf_intersects.progress_apply(
            lambda row: fn_lookup_sign(row, hydro_geom_lookup), axis=1
        )

        print(gdf_intersects["cross_product_sign"].value_counts(dropna=False))

        # Mask of rows that need flipping (sign == 1). NaN and -1 are left alone.
        mask_flip = gdf_intersects["cross_product_sign"] == 1
        print(f"Flipping {mask_flip.sum()} of {len(gdf_intersects)} rows.")

        # Apply and unpack back into the two columns
        flipped = gdf_intersects.loc[mask_flip].progress_apply(
            fn_flip_row, axis=1, result_type="expand"
        )

        if not flipped.empty:
            gdf_intersects.loc[mask_flip, "geometry"] = flipped[0].values
            gdf_intersects.loc[mask_flip, "xsec"] = flipped[1].values

        # Sanity check: re-derive the sign and confirm the flips took
        gdf_intersects["cross_product_sign_after"] = gdf_intersects.progress_apply(
            lambda row: fn_lookup_sign(row, hydro_geom_lookup), axis=1
        )
        print(gdf_intersects["cross_product_sign_after"].value_counts(dropna=False))

        str_flipped_streams_on_ornl = os.path.join(
            STR_OUTPUT_FOLDER_PATH, str_dist, '04_bridge_flipped_on_ornl_ln_4326.fgb'
        )
        gdf_intersects.to_file(str_flipped_streams_on_ornl, driver="FlatGeobuf")

        print(f"    -- Wrote {str_flipped_streams_on_ornl}")

        # (Original notebook previewed one cross section here:)
        # fn_plot_bridge_cross_section(gdf_intersects.iloc[0]['xsec']); plt.show()


# ===========================================================================
# Part 3 -- helpers (assign NBI where missing)
# ===========================================================================

def _extract_brdg_nbr(x):
    """Pull 'brdg_nbr' out of the xsec JSON blob; returns None on bad/empty."""
    if pd.isna(x):
        return None
    try:
        return json.loads(x).get('brdg_nbr')
    except (json.JSONDecodeError, TypeError):
        return None


def fn_match_district(str_dist, str_epsg, gdf_snbi_points_4326):
    """
    For one district: fill ID01_BRDG_NBR on the no-NBI bridge lines from the
    nearest SNBI point within FLT_SEARCH_DIST_FT, and record that distance in
    meters. Returns the augmented full bridge-lines GeoDataFrame.
    """
    str_lines = os.path.join(
        STR_OUTPUT_FOLDER_PATH, str_dist, '04_bridge_flipped_on_ornl_ln_4326.fgb'
    )
    gdf_bridge_lines = gpd.read_file(str_lines).to_crs(str_epsg)

    # reproject a fresh copy of the points into this district's CRS
    gdf_snbi_points = gdf_snbi_points_4326.to_crs(str_epsg)

    # bridge_id must be unique or the left merge below fans out into dup rows
    assert gdf_bridge_lines['bridge_id'].is_unique, \
        f"[{str_dist}] bridge_id is not unique in gdf_bridge_lines"

    # pull brdg_nbr out of the xsec json (kept for downstream validation use)
    gdf_bridge_lines['xsec_brdg_nbr'] = gdf_bridge_lines['xsec'].apply(_extract_brdg_nbr)

    # ----- 1) crossings that currently have no NBI record --------------------
    gdf_no_nbi = gdf_bridge_lines[gdf_bridge_lines['ID01_BRDG_NBR'].isna()].copy()
    int_no_nbi = len(gdf_no_nbi)

    # ----- 2) nearest SNBI point within the search distance ------------------
    gdf_no_nbi = gdf_no_nbi.reset_index(drop=True)
    gdf_no_nbi['__line_id'] = gdf_no_nbi.index
    gdf_snbi_points = gdf_snbi_points.reset_index(drop=True)
    gdf_snbi_points['__pt_id'] = gdf_snbi_points.index

    gdf_buf = gdf_no_nbi[['__line_id', 'geometry']].copy()
    gdf_buf['geometry'] = gdf_buf.geometry.buffer(FLT_SEARCH_DIST_FT)

    gdf_hits = gpd.sjoin(
        gdf_buf,
        gdf_snbi_points[['__pt_id', 'geometry']],
        how='inner', predicate='intersects',
    )

    if len(gdf_hits):
        line_geom = gdf_no_nbi.geometry.loc[gdf_hits['__line_id']].values
        pt_geom = gdf_snbi_points.geometry.loc[gdf_hits['__pt_id']].values
        gdf_hits['distance2nearest_snbi_ft'] = [
            l.distance(p) for l, p in zip(line_geom, pt_geom)
        ]
        gdf_hits = (gdf_hits
                    .sort_values('distance2nearest_snbi_ft')
                    .drop_duplicates('__line_id', keep='first'))
    else:
        gdf_hits = gdf_hits.reindex(
            columns=['__line_id', '__pt_id', 'distance2nearest_snbi_ft']
        )

    gdf_near = gdf_no_nbi.merge(
        gdf_hits[['__line_id', '__pt_id', 'distance2nearest_snbi_ft']],
        on='__line_id', how='left',
    )

    # map the matched point's SNBI bridge number in via __pt_id
    ser_pt_to_id = gdf_snbi_points.set_index('__pt_id')[STR_SNBI_ID_FIELD]
    gdf_near['nearest_snbi_ID01_BRDG_NBR'] = gdf_near['__pt_id'].map(ser_pt_to_id)

    # ----- 3) fold matches back onto the FULL lines layer --------------------
    gdf_matched = gdf_near.loc[
        gdf_near['distance2nearest_snbi_ft'].notna(),
        ['bridge_id', 'nearest_snbi_ID01_BRDG_NBR', 'distance2nearest_snbi_ft'],
    ]
    gdf_bridge_lines = gdf_bridge_lines.merge(gdf_matched, on='bridge_id', how='left')

    # distance in meters (NaN where no match)
    gdf_bridge_lines['distance2nearest_snbi_m'] = (
        gdf_bridge_lines['distance2nearest_snbi_ft'] * FLT_USFT_TO_M
    )

    # flag the rows we're about to fill, before overwriting the NaNs
    gdf_bridge_lines['snbi_filled'] = \
        gdf_bridge_lines['nearest_snbi_ID01_BRDG_NBR'].notna()

    # fill only the previously-missing NBI numbers; existing ones are untouched
    gdf_bridge_lines['ID01_BRDG_NBR'] = gdf_bridge_lines['ID01_BRDG_NBR'].fillna(
        gdf_bridge_lines['nearest_snbi_ID01_BRDG_NBR']
    )

    # drop the helper columns
    gdf_bridge_lines = gdf_bridge_lines.drop(
        columns=['nearest_snbi_ID01_BRDG_NBR', 'distance2nearest_snbi_ft']
    )

    int_within = int(gdf_bridge_lines['snbi_filled'].sum())
    print(f"[{str_dist}] no-NBI: {int_no_nbi} | "
          f"filled within {FLT_SEARCH_DIST_FT:.0f} ft: {int_within} | "
          f"none: {int_no_nbi - int_within}")

    return gdf_bridge_lines


# ===========================================================================
# Part 3 -- run
# ===========================================================================

def run_part3():
    """Assign NBI number from nearest SNBI point where missing (<= FLT_SEARCH_DIST_FT)."""
    # read the SNBI points once, keep the pristine 4326 copy
    gdf_snbi_points_4326 = gpd.read_file(STR_SNBI_POINTS)

    dict_results = {}
    for str_dist, str_epsg in DICT_DISTRICTS_TO_PROCESS.items():
        gdf_bridge_lines = fn_match_district(str_dist, str_epsg, gdf_snbi_points_4326)
        dict_results[str_dist] = gdf_bridge_lines

        # write the augmented lines back out per district
        str_out = os.path.join(STR_OUTPUT_FOLDER_PATH, str_dist, '05_bridge_lines_snbi_filled_4326.fgb')
        gdf_bridge_lines.to_crs('EPSG:4326').to_file(str_out, driver='FlatGeobuf')

    return dict_results


# ===========================================================================
# Part 4 -- helpers (build final data tables)
# ===========================================================================

def fn_extract_min_low_chord_min_overtop_min_ground(xsec, tol=0.01):
    """
    Returns (min_low_chord, min_overtop, min_ground), each rounded to 2 dp.

    min_low_chord : lowest low-chord elevation across the opening, ignoring
                    points sitting on natural ground (abutment slopes) EXCEPT
                    the two `trans` abutment-face points, which always count.
    min_overtop   : min_low_chord + bridge_thi - top-of-deck elevation at the
                    lowest opening point; the stage at which water overtops.
    min_ground    : lowest ground elevation (channel thalweg).

    Returns (nan, nan, min_ground) if no low-chord points qualify;
    (nan, nan, nan) if ground is empty too.
    """
    if isinstance(xsec, str):
        xsec = json.loads(xsec)

    lc_d = np.asarray(xsec["lowchord"]["d"], float)
    lc_z = np.asarray(xsec["lowchord"]["z"], float)
    g_d = np.asarray(xsec["ground"]["d"], float)
    g_z = np.asarray(xsec["ground"]["z"], float)

    min_ground = round(float(g_z.min()), 2) if g_z.size else np.nan

    order = np.argsort(g_d)
    g_at_lc = np.interp(lc_d, g_d[order], g_z[order])
    off_ground = np.abs(lc_z - g_at_lc) > tol

    trans_x = np.asarray(xsec.get("trans", {}).get("x", []), float)
    if trans_x.size:
        is_trans = np.isclose(lc_d[:, None], trans_x[None, :], atol=tol).any(axis=1)
    else:
        is_trans = np.zeros_like(lc_d, dtype=bool)

    keep = off_ground | is_trans
    if not keep.any():
        return np.nan, np.nan, min_ground

    min_low_ch = round(float(lc_z[keep].min()), 2)
    thi = float(xsec.get("bridge_thi", np.nan))
    min_overtop = round(min_low_ch + thi, 2)
    return min_low_ch, min_overtop, min_ground


def fn_compute_hand_r(row, df_hydro):
    hydro_id = int(row['HydroID'])
    df_h = df_hydro[df_hydro['FATSGTID'] == hydro_id].sort_values('stage')
    if df_h.empty:
        return str([])
    tuples = [
        (round(r['discharge_cms'] * CMS_TO_CFS, 1), round(r['stage'] * M_TO_FT, 1))
        for _, r in df_h.iterrows()
    ]
    return str(tuples)


def fn_build_rating_curve_for_t_bridge_rating_curve(row):
    hand_r = row["hand_r"]
    if isinstance(hand_r, str):
        hand_r = ast.literal_eval(hand_r)

    if len(hand_r) < 2:
        return None, "[]"

    # drop first tuple, add min_ground to elevation (second value)
    min_ground = row["min_ground"]
    trimmed = [(int(round(q)), round(elv + min_ground, 2)) for q, elv in hand_r[1:]]

    # min_flow is second tuple's first value (first after drop)
    min_flow = trimmed[0][0] if trimmed else None

    return min_flow, str(trimmed)


def fn_compute_zone_limits(row):
    ground_elv = row["ground_elv"]

    # if stored as a "[...]" string, parse it back into a list of floats
    if isinstance(ground_elv, str):
        ground_elv = [float(x) for x in ast.literal_eval(ground_elv)]

    flt_min_low_ch = row["min_low_ch"]
    flt_min_ground = row["min_ground"]

    flt_buffer_ground = 1.0
    list_zones = [flt_buffer_ground * -1, 0.5, 2.0, 5.0]  # THIS IS HOW THE WARNING ZONES ARE DEFINED
    # [-1.0, 0.5, 2.0, 5.0]

    list_depth_from_min = [x - flt_min_ground for x in ground_elv]
    flt_dist_to_low_ch = flt_min_low_ch - flt_min_ground

    list_zone_limits = []

    # Zone top: max depth from min ground
    flt_zone_top_depth = max(list_depth_from_min)
    list_zone_limits.append(flt_zone_top_depth)

    # Zone 1 bottom
    flt_bottom_depth = flt_dist_to_low_ch - list_zones[1]
    list_zone_limits.append(flt_bottom_depth)

    # Zone 2
    if flt_dist_to_low_ch > list_zones[1]:
        if flt_dist_to_low_ch > list_zones[2]:
            flt_bottom_depth = flt_dist_to_low_ch - list_zones[2]
        else:
            flt_bottom_depth = list_zones[0]
        list_zone_limits.append(flt_bottom_depth)

    # Zone 3
    if flt_dist_to_low_ch > list_zones[2]:
        if flt_dist_to_low_ch > list_zones[3]:
            flt_bottom_depth = flt_dist_to_low_ch - list_zones[3]
        else:
            flt_bottom_depth = list_zones[0]
        list_zone_limits.append(flt_bottom_depth)

    # Zone 4
    if flt_dist_to_low_ch > list_zones[3]:
        list_zone_limits.append(list_zones[0])

    list_zone_limits = [round(x, 2) for x in list_zone_limits]
    return str(list_zone_limits)


def fn_parse_xsec(xsec):
    """
    Parse a bridge cross-section dict into four equal-length lists,
    all sampled at the ground survey stations.
    Returns a dict with:
        sta          -> ground['d']            (station / offset)
        ground_elv   -> ground['z']            (natural ground elevation)
        deck_elev    -> roadsurf['z']  interpolated onto sta
        low_ch_elev  -> band['z_bot']  interpolated onto sta
    The deck and low-chord series are surveyed on their own station
    grids, so they're linearly interpolated onto the ground stations
    to guarantee equal length and one-to-one alignment. All values
    are rounded to two decimal places.
    """
    sta = np.asarray(xsec["ground"]["d"], dtype=float)
    ground_elv = np.asarray(xsec["ground"]["z"], dtype=float)
    # roadsurf -> deck
    deck_d = np.asarray(xsec["roadsurf"]["d"], dtype=float)
    deck_z = np.asarray(xsec["roadsurf"]["z"], dtype=float)
    # band -> low chord (underside)
    lc_d = np.asarray(xsec["band"]["d"], dtype=float)
    lc_z = np.asarray(xsec["band"]["z_bot"], dtype=float)
    deck_elev = np.interp(sta, deck_d, deck_z)
    low_ch_elev = np.interp(sta, lc_d, lc_z)
    return {
        "sta": np.round(sta, 2).tolist(),
        "ground_elv": np.round(ground_elv, 2).tolist(),
        "deck_elev": np.round(deck_elev, 2).tolist(),
        "low_ch_elev": np.round(low_ch_elev, 2).tolist(),
    }


def fn_sync_xsec_id(row):
    x = row['xsec']

    # Handle the case where xsec is stored as a JSON string rather than a dict
    if isinstance(x, str):
        x = json.loads(x)

    if isinstance(x, dict):
        x = dict(x)                     # copy so we don't mutate a shared reference
        x['brdg_nbr'] = row['ID01_BRDG_NBR']
        # If you also want the internal id to match, uncomment:
        # x['bridge_id'] = row['ID01_BRDG_NBR']

    return x


def fn_row_to_bridge_json(row):
    def csv(seq):
        # -> "a, b, c"  (no brackets). Handles real lists, bracketed strings,
        #    and already-joined comma strings.
        if isinstance(seq, str):
            s = seq.strip()
            if s.startswith("[") and s.endswith("]"):
                seq = ast.literal_eval(s)
            else:
                return s
        return ", ".join(str(v) for v in seq)

    def as_str(v):
        # keep brackets: -> "[...]"
        return v if isinstance(v, str) else str(v)

    return {
        "sta":           csv(row["sta"]),
        "ground_elv":    csv(row["ground_elv"]),
        "deck_elev":     csv(row["deck_elev"]),
        "uuid":          row["uuid"],
        "low_ch_elv":    csv(row["low_ch_elev"]),   # source key: low_ch_elev
        "min_low_ch":    float(row["min_low_ch"]),
        "min_ground":    float(row["min_ground"]),
        "hand_r":        as_str(row["hand_r"]),      # keeps brackets
        "anno_xs_title": row["anno_xs_title"],
        "anno_latlong":  row["anno_latlong"],
        "anno_nbi":      row["anno_nbi"],
        "anno_comid":    row["anno_comid"],
        "zone_limits":   as_str(row["zone_limits"]),  # keeps brackets
    }


# ===========================================================================
# Part 4 -- run
# ===========================================================================

def run_part4():
    """Build final bridge data tables, rating curves and per-bridge JSON."""
    # read the lookup for NWM to ORNL stream lines
    df_nextgen_to_nwm_lookup = pd.read_parquet(STR_T_NEXTGEN_TO_NWM_FILEPATH, engine="pyarrow")

    # Build the lookup: nextgen_id -> feature_id
    lookup = df_nextgen_to_nwm_lookup.set_index("nextgen_id")["feature_id"]

    # Normalize both keys to the SAME type so the match actually lands.
    # Strings are the safe common denominator (handles int-vs-str, and stray floats like 123.0).
    lookup.index = lookup.index.astype(str).str.strip()

    for str_dist, str_epsg in DICT_DISTRICTS_TO_PROCESS.items():
        str_xs_for_forecasting = os.path.join(
            STR_OUTPUT_FOLDER_PATH, str_dist, '05_bridge_lines_snbi_filled_4326.fgb'
        )
        gdf_bridge_lines = gpd.read_file(str_xs_for_forecasting)

        # convert JSON string to dict
        gdf_bridge_lines['xsec'] = gdf_bridge_lines['xsec'].apply(
            lambda s: json.loads(s) if isinstance(s, str) else s
        )

        gdf_bridge_lines_for_s_bridge_pnt = gdf_bridge_lines.copy()

        keep_cols = [
            "ID01_BRDG_NBR", "HydroID", "osm_names", "river_names",
            "bridge_thickness", "opening_area_sqft", "max_grade", "n_transition_pts",
            "intersects_ORNL_stream", "has_min_area", "is_flat_road", "has_two_transition",
            "xsec", "geometry",
        ]

        gdf_bridge_lines_for_s_bridge_pnt = gdf_bridge_lines_for_s_bridge_pnt[keep_cols].copy()

        gdf_bridge_lines_for_s_bridge_pnt[["min_low_ch", "min_overtop", "min_ground"]] = (
            gdf_bridge_lines_for_s_bridge_pnt["xsec"].apply(
                lambda x: pd.Series(fn_extract_min_low_chord_min_overtop_min_ground(x))
            )
        )

        # Compute bridge line centroid in original CRS
        centroids = gdf_bridge_lines_for_s_bridge_pnt.geometry.centroid

        # Convert bridge line centroids to EPSG:4326
        centroids_4326 = centroids.to_crs(epsg=4326)

        # Create Lat / Long columns (rounded to 4 decimals)
        gdf_bridge_lines_for_s_bridge_pnt["Long"] = centroids_4326.x.round(4)
        gdf_bridge_lines_for_s_bridge_pnt["Lat"] = centroids_4326.y.round(4)

        # Extract the HAND rating curve for each bridge line
        df_hydro = pd.read_parquet(STR_FILEPATH_TO_HAND_RATING_CURVES)
        gdf_bridge_lines_for_s_bridge_pnt["hand_r"] = gdf_bridge_lines_for_s_bridge_pnt.apply(
            lambda row: fn_compute_hand_r(row, df_hydro), axis=1
        )

        # Create a uuid per bridge
        gdf_bridge_lines_for_s_bridge_pnt["uuid"] = gdf_bridge_lines_for_s_bridge_pnt.apply(
            lambda _: str(uuid.uuid4()), axis=1
        )

        # for each item...from 'HydroID' (nextgen_id in df_nextgen_to_nwm_lookup)
        # determine the 'feature_id' and write it to that row
        hydro_key = gdf_bridge_lines_for_s_bridge_pnt["HydroID"].astype(str).str.strip()
        gdf_bridge_lines_for_s_bridge_pnt["feature_id"] = hydro_key.map(lookup)

        # Create the cross section annotation text
        gdf_xs = gdf_bridge_lines_for_s_bridge_pnt.copy()

        osm = gdf_xs["osm_names"].fillna("").astype(str).str.strip()
        riv = gdf_xs["river_names"].fillna("").astype(str).str.strip()

        osm_t = osm.str.title()
        both = (osm != "") & (riv != "")

        gdf_xs["anno_xs_title"] = np.where(
            both,
            osm_t + " @ " + riv,
            np.where(osm != "", osm_t, riv),   # else whichever exists; "" if neither
        )

        gdf_xs["anno_latlong"] = "Lat/Long: (" + gdf_xs["Lat"].astype(str) + "," + gdf_xs["Long"].astype(str) + ")"
        gdf_xs["anno_nbi"] = "NBI: " + gdf_xs["ID01_BRDG_NBR"].astype(str)
        gdf_xs["anno_comid"] = "NWM COMID: " + gdf_xs["feature_id"].astype(str)

        # Parse once per row into a dict of the four lists
        parsed = gdf_xs["xsec"].apply(fn_parse_xsec)

        # Explode each key into its own column, storing the list as a string
        for key in ("sta", "ground_elv", "deck_elev", "low_ch_elev"):
            gdf_xs[key] = parsed.apply(lambda d, k=key: str(d[k]))

        # Generate Warning Zone Limits
        gdf_xs["zone_limits"] = gdf_xs.apply(fn_compute_zone_limits, axis=1)

        # FIX xsec json to include and match NBI in 'xsec' json
        gdf_xs['xsec'] = gdf_xs.apply(fn_sync_xsec_id, axis=1)

        gdf_xs['dist'] = str_dist

        # Write the output (all lines that cross ORNL FAST streams)
        str_out_bridges = os.path.join(
            STR_OUTPUT_FOLDER_PATH, str_dist, '06_bridge_lines_for_accounting_4326.fgb'
        )
        gdf_xs.to_crs('EPSG:4326').to_file(str_out_bridges, driver='FlatGeobuf')

        # These are the only points where we are going to make a forecast
        gdf_xs_filtered = gdf_xs[
            gdf_xs["intersects_ORNL_stream"]
            & gdf_xs["has_min_area"]
            & gdf_xs["has_two_transition"]
            & gdf_xs["is_flat_road"]
        ].copy()

        # building the t_bridge_rating_curve
        # get the list_rating curve (removing 0,0) and min flow
        results = gdf_xs_filtered.apply(fn_build_rating_curve_for_t_bridge_rating_curve, axis=1)

        df_for_t_bridge_rating_curve = pd.DataFrame({
            "uuid_bridge":       gdf_xs_filtered["uuid"].astype(str),
            "nextgen_id":        gdf_xs_filtered["HydroID"].astype(str),
            "min_flow":          results.apply(lambda x: x[0]),
            "list_rating_curve": results.apply(lambda x: x[1]),
        })

        gdf_s_bridge_pnt = gpd.GeoDataFrame({
            "uuid_bridge": gdf_xs_filtered["uuid"].astype(str),
            "BRDG_ID":     gdf_xs_filtered["ID01_BRDG_NBR"].astype(str),
            "min_low_ch":  gdf_xs_filtered["min_low_ch"].astype(float),
            "min_ground":  gdf_xs_filtered["min_ground"].astype(float),
            "min_overtop": gdf_xs_filtered["min_overtop"].astype(float),
            "name":        gdf_xs_filtered["osm_names"].astype(str),
            "ref":         "",
            "nhd_name":    gdf_xs_filtered["river_names"].astype(str),
            "geometry":    gdf_xs_filtered.geometry.centroid,
        }, geometry="geometry", crs=gdf_xs_filtered.crs)

        records_json_per_bridge = [fn_row_to_bridge_json(r) for _, r in gdf_xs_filtered.iterrows()]

        # Write the output
        str_out_bridges_dir = os.path.join(STR_OUTPUT_FOLDER_PATH, str_dist, '07_for_FAST_database')
        os.makedirs(str_out_bridges_dir, exist_ok=True)

        str_out_t_bridge_rating_curve = os.path.join(str_out_bridges_dir, '07_01_t_bridge_rating_curve.parquet')
        df_for_t_bridge_rating_curve.to_parquet(str_out_t_bridge_rating_curve)

        str_out_s_bridge_pnt = os.path.join(str_out_bridges_dir, '07_02_s_bridge_pnt_4326.fgb')
        gdf_s_bridge_pnt.to_crs('EPSG:4326').to_file(str_out_s_bridge_pnt, driver='FlatGeobuf')

        str_out_json_per_bridge = os.path.join(str_out_bridges_dir, '07_03_json_per_bridge')
        os.makedirs(str_out_json_per_bridge, exist_ok=True)

        for record in records_json_per_bridge:
            out_path = os.path.join(str_out_json_per_bridge, f"{record['uuid']}.json")
            with open(out_path, 'w') as f:
                json.dump(record, f, indent=4)


# ===========================================================================
# Orchestration
# ===========================================================================

def main():
    print("========== Part 1: build bridge lines / metrics / xsec ==========")
    run_part1()

    print("========== Part 2: flip cross sections ==========")
    run_part2()

    print("========== Part 3: assign NBI where missing ==========")
    run_part3()

    print("========== Part 4: build final bridge data tables ==========")
    run_part4()

    print("Done.")


if __name__ == "__main__":
    main()