"""
review_bridges.py

End-to-end bridge cross-section review tool:
  1. Walks a root folder for all files named 06_bridge_lines_for_accounting_4326.fgb
  2. Filters each to matched records (stream intersect, min area, two transitions,
     not a flat road) and combines them into one GeoDataFrame.
  3. Opens an interactive window to review each cross section by keystroke
     (3 = approve, 1 = disapprove).
  4. Writes out the combined-matched set and the approved-only set as .fgb.

Run from a terminal with your geopandas env active:
    python review_bridges.py
Requires a GUI matplotlib backend (TkAgg / Qt) — a plain `python` run gives you one.
"""

import os
import json
import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages

# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #
ROOT = r"E:\bridge_conversions_20260805\output_20260807"
TARGET_NAME = "06_bridge_lines_for_accounting_4326.fgb"

MATCHED_GDF_PATH = os.path.join(ROOT, "matched_bridges_4326.fgb")
APPROVED_GDF_PATH = os.path.join(ROOT, "approved_bridges_4326.fgb")
REVIEW_CSV = os.path.join(ROOT, "cross_section_reviews.csv")
PDF_PATH = os.path.join(ROOT, "matched_bridge_cross_sections.pdf")

# Column holding the cross-section JSON/dict. None = auto-detect.
DATA_COL = None
# Bridge id column for stable review keys (survives re-runs / re-ordering).
# Set to e.g. "brdg_nbr" if you have it; None keys off row position.
ID_COL = None

APPROVE_KEY = "3"
DISAPPROVE_KEY = "1"

# Set True to also render the 10x10 contact-sheet PDF of all matches.
WRITE_CONTACT_SHEET = False
NROWS, NCOLS = 10, 10
PER_SHEET = NROWS * NCOLS


# --------------------------------------------------------------------------- #
# Plotting
# --------------------------------------------------------------------------- #
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.
    """
    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


# --------------------------------------------------------------------------- #
# Loading / matching
# --------------------------------------------------------------------------- #
def match_mask(gdf):
    return (
        (gdf["intersects_ORNL_stream"] == True) &
        (gdf["has_min_area"] == True) &
        (gdf["has_two_transition"] == True) &
        (gdf["is_flat_road"] == False)
    )


def detect_data_col(gdf):
    """Find an object column whose values parse to the cross-section dict."""
    geom_name = gdf.geometry.name
    for col in gdf.columns:
        if col == geom_name:
            continue
        s = gdf[col].dropna()
        if s.empty:
            continue
        val = s.iloc[0]
        try:
            d = json.loads(val) if isinstance(val, str) else val
        except (json.JSONDecodeError, TypeError):
            continue
        if isinstance(d, dict) and {"band", "ground", "roadsurf", "lowchord"} <= set(d):
            return col
    return None


def find_files(root, target_name):
    files = []
    for dirpath, _, filenames in os.walk(root):
        for fn in filenames:
            if fn == target_name:
                files.append(os.path.join(dirpath, fn))
    return files


def collect_matched_bridges(root, target_name, data_col=None):
    """
    Walk `root`, filter each target file to matched rows, and combine into one
    GeoDataFrame. Returns (matched_bridges, plot_col).
    """
    files = find_files(root, target_name)
    print(f"Found {len(files)} matching file(s).")

    matched_gdfs = []
    detected_cols = []
    for fp in files:
        gdf = gpd.read_file(fp)
        col = data_col or detect_data_col(gdf)
        sub = gdf[match_mask(gdf)].copy()
        sub["source_file"] = fp
        matched_gdfs.append(sub)
        detected_cols.append(col)
        note = f"data col='{col}'" if col else "no cross-section data column found"
        print(f"  {fp}: {len(sub)} matches ({note})")

    if not matched_gdfs:
        raise SystemExit("No matching files found.")

    matched = gpd.GeoDataFrame(
        pd.concat(matched_gdfs, ignore_index=True),
        crs=matched_gdfs[0].crs,
    )
    plot_col = next((c for c in detected_cols if c is not None), None)
    print(f"\nTotal matched bridges (combined): {len(matched)}")
    return matched, plot_col


# --------------------------------------------------------------------------- #
# Persistence helper (guards the JSON/dict column against FlatGeobuf)
# --------------------------------------------------------------------------- #
def write_fgb(gdf, path, data_col=None):
    out = gdf.copy()
    if data_col and data_col in out.columns and len(out):
        first = out[data_col].dropna()
        if not first.empty and isinstance(first.iloc[0], dict):
            out[data_col] = out[data_col].apply(
                lambda v: json.dumps(v) if isinstance(v, dict) else v
            )
    out.to_file(path, driver="FlatGeobuf")
    print(f"Wrote {len(out)} rows: {path}")


# --------------------------------------------------------------------------- #
# Optional contact-sheet PDF (10x10 per page)
# --------------------------------------------------------------------------- #
def write_contact_sheet(records, output_pdf):
    if not records:
        print("Nothing to plot for contact sheet.")
        return
    with PdfPages(output_pdf) as pdf:
        n_sheets = (len(records) + PER_SHEET - 1) // PER_SHEET
        for start in range(0, len(records), PER_SHEET):
            chunk = records[start:start + PER_SHEET]
            fig, axes = plt.subplots(NROWS, NCOLS, figsize=(NCOLS * 3, NROWS * 2.2))
            axes = axes.ravel()
            handles, labels = None, None
            for ax, data in zip(axes, chunk):
                try:
                    fn_plot_bridge_cross_section(data, ax=ax)
                except Exception as e:
                    ax.text(0.5, 0.5, f"plot error:\n{e}", ha="center",
                            va="center", transform=ax.transAxes,
                            fontsize=6, color="red")
                if handles is None:
                    h, l = ax.get_legend_handles_labels()
                    if h:
                        handles, labels = h, l
                leg = ax.get_legend()
                if leg is not None:
                    leg.remove()
                ax.set_xlabel(""); ax.set_ylabel("")
                ax.title.set_fontsize(7)
                ax.tick_params(labelsize=5)
            for ax in axes[len(chunk):]:
                ax.axis("off")
            if handles:
                fig.legend(handles, labels, loc="lower center",
                           ncol=len(labels), fontsize=8)
            fig.tight_layout(rect=[0, 0.03, 1, 1])
            pdf.savefig(fig)
            plt.close(fig)
            print(f"  wrote sheet {start // PER_SHEET + 1}/{n_sheets} "
                  f"({len(chunk)} plots)")
    print(f"Contact sheet: {output_pdf}")


# --------------------------------------------------------------------------- #
# Interactive reviewer
# --------------------------------------------------------------------------- #
def review_cross_sections(gdf, data_col, id_col=None,
                          review_csv=REVIEW_CSV, resume=True):
    """
    Manually review each cross section. Keys:
      3 -> approve, 1 -> disapprove
      left/right -> move without changing a decision
      s -> save & quit,  q/esc -> quit (progress saved after each keystroke)
    Returns the GeoDataFrame of approved (==3) rows.
    """
    gdf = gdf.reset_index(drop=True)

    if id_col and id_col in gdf.columns:
        keys = gdf[id_col].astype(str).tolist()
    else:
        keys = [str(i) for i in range(len(gdf))]

    decisions = {}
    if resume and os.path.exists(review_csv):
        prior = pd.read_csv(review_csv, dtype={"key": str})
        decisions = dict(zip(prior["key"], prior["decision"].astype(int)))
        print(f"Resuming: {len(decisions)} prior decisions loaded.")

    def save():
        pd.DataFrame(
            {"key": list(decisions.keys()),
             "decision": list(decisions.values())}
        ).to_csv(review_csv, index=False)

    state = {"i": 0}

    # Jump to first undecided when resuming
    if resume and decisions:
        for idx, k in enumerate(keys):
            if k not in decisions:
                state["i"] = idx
                break

    fig, ax = plt.subplots(figsize=(12, 6))

    def draw():
        ax.clear()
        i = state["i"]
        k = keys[i]
        data = gdf.iloc[i][data_col]
        try:
            fn_plot_bridge_cross_section(data, ax=ax)
        except Exception as e:
            ax.text(0.5, 0.5, f"plot error:\n{e}", ha="center", va="center",
                    transform=ax.transAxes, color="red")
        prev = decisions.get(k)
        status = {3: "APPROVED", 1: "disapproved"}.get(prev, "— undecided —")
        ax.set_title(
            f"[{i + 1}/{len(gdf)}]  id={k}   current: {status}\n"
            f"3=approve  1=disapprove  <-/-> move  s=save&quit  q=quit   "
            f"({len(decisions)} decided)",
            fontsize=10,
        )
        fig.canvas.draw_idle()

    def advance():
        if state["i"] < len(gdf) - 1:
            state["i"] += 1
        else:
            print("Reached the last cross section.")
        draw()

    def on_key(event):
        k = keys[state["i"]]
        if event.key == APPROVE_KEY:
            decisions[k] = 3; save(); advance()
        elif event.key == DISAPPROVE_KEY:
            decisions[k] = 1; save(); advance()
        elif event.key == "right":
            advance()
        elif event.key == "left":
            if state["i"] > 0:
                state["i"] -= 1; draw()
        elif event.key == "s":
            save(); plt.close(fig)
        elif event.key in ("q", "escape"):
            plt.close(fig)

    fig.canvas.mpl_connect("key_press_event", on_key)
    draw()
    plt.show()  # blocks until window closed
    save()

    approved_keys = {k for k, v in decisions.items() if int(v) == 3}
    mask = [k in approved_keys for k in keys]
    approved = gdf[mask].copy()

    n_app = sum(1 for v in decisions.values() if int(v) == 3)
    n_dis = sum(1 for v in decisions.values() if int(v) == 1)
    print(f"\nReviewed {len(decisions)}/{len(gdf)}  "
          f"| approved: {n_app}  disapproved: {n_dis}")
    return approved


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #
def main():
    matched_bridges, plot_col = collect_matched_bridges(
        ROOT, TARGET_NAME, data_col=DATA_COL
    )

    # Save the combined matched set
    write_fgb(matched_bridges, MATCHED_GDF_PATH, data_col=plot_col)

    if plot_col is None:
        raise SystemExit(
            "No cross-section data column detected — cannot plot for review. "
            "Set DATA_COL explicitly if you know the column name."
        )

    if WRITE_CONTACT_SHEET:
        records = matched_bridges[plot_col].dropna().tolist()
        write_contact_sheet(records, PDF_PATH)

    # Interactive review
    print("\nOpening reviewer window — press 3 to approve, 1 to disapprove.")
    approved_bridges = review_cross_sections(
        matched_bridges, plot_col, id_col=ID_COL
    )
    print(f"Approved bridges: {len(approved_bridges)}")

    # Save approved-only set
    if len(approved_bridges):
        write_fgb(approved_bridges, APPROVED_GDF_PATH, data_col=plot_col)
    else:
        print("No approved rows — nothing written.")

    return approved_bridges


if __name__ == "__main__":
    approved_bridges = main()