#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Figure 7: the boundary walk on the example basin 8595.

Input: ../Data/ex8595_*.tif.
Output: ../Figures/Fig07_BoundaryWalk_v11.png and .pdf.

Run from this folder; ../Data holds the inputs and the figure is written to ../Figures.
"""

import os

import numpy as np
import rasterio

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.patches import Patch


# ============================================================================
# Style (ESSD / AGU – Helvetica, cm mathtext, 9 pt)
# ============================================================================
mpl.rcParams.update({
    "font.family":         "Helvetica",
    "font.size":            9,
    "axes.labelsize":       9,
    "axes.titlesize":      10,
    "axes.titleweight":    "bold",
    "xtick.labelsize":      8,
    "ytick.labelsize":      8,
    "mathtext.fontset":   "cm",
    "axes.linewidth":     0.6,
    "axes.unicode_minus": False,
})


# ============================================================================
# Paths
# ============================================================================
HERE     = os.path.dirname(os.path.abspath(__file__))
ROOT     = os.path.abspath(os.path.join(HERE, "..", ".."))
DATA     = os.path.join(HERE, "..", "Data")
OUT_DIR  = os.path.join(HERE, "..", "Figures")
os.makedirs(OUT_DIR, exist_ok=True)

FOCUS_BSN_ID = 8595      # global_basin_id (the v1.1 basin raster holds global ids)
SKEL_KM2     = 10.0      # UPA threshold for the basin skeleton (noise filter)


# ============================================================================
# D8 code conventions
# ============================================================================
# Direction (dr, dc) -> the D8 code that *makes a pixel point that way*.
# Used to find a pixel's *upstream* neighbours: a neighbour at (dr, dc)
# is upstream iff its D8 code is the inverse direction.
INV = {(0, 1): 16, (1, 1): 32, (1, 0): 64, (1, -1): 128,
       (0, -1): 1, (-1, -1): 2, (-1, 0): 4, (-1, 1): 8}

# Direct (dr, dc) lookup for a given D8 code (downstream offset).
D8_OFF = {1: (0, 1), 2: (1, 1), 4: (1, 0), 8: (1, -1),
          16: (0, -1), 32: (-1, -1), 64: (-1, 0), 128: (-1, 1)}


# ============================================================================
# Read rasters + tight crop around focus basin
# ============================================================================
def _read(name):
    with rasterio.open(os.path.join(DATA, f"ex8595_{name}.tif")) as s:
        return s.read(1)

dir_arr  = _read("dir")
upa      = _read("upa")
dfs_in   = _read("dfs_in")
dfs_out  = _read("dfs_out")
bsn_full = _read("bsn")
bmask_full = (bsn_full == FOCUS_BSN_ID)
H_full, W_full = bmask_full.shape

rows = np.where(bmask_full.any(axis=1))[0]
cols = np.where(bmask_full.any(axis=0))[0]
PAD  = 6
r0 = max(0, rows[0] - PAD); r1 = min(H_full, rows[-1] + 1 + PAD)
c0 = max(0, cols[0] - PAD); c1 = min(W_full, cols[-1] + 1 + PAD)
sl = (slice(r0, r1), slice(c0, c1))

dir_c   = dir_arr[sl]
upa_c   = upa[sl]
dfin_c  = dfs_in[sl]
dfout_c = dfs_out[sl]
bmask   = bmask_full[sl]
H, W    = bmask.shape


# ============================================================================
# Helpers
# ============================================================================
def upstream_neighbours(i, j, mask, dir_local):
    out = []
    for (di, dj), code in INV.items():
        ni, nj = i + di, j + dj
        if (0 <= ni < mask.shape[0] and 0 <= nj < mask.shape[1]
                and mask[ni, nj] and dir_local[ni, nj] == code):
            out.append((ni, nj))
    return out


def walk_max_upa(start, mask, upa_local, dir_local):
    """Greedy upstream walk: at each step take the upstream neighbour with
    the largest UPA.  Used for tracing the main stem and other skeleton
    branches."""
    chain = [start]
    cur   = start
    while True:
        ups = upstream_neighbours(*cur, mask, dir_local)
        if not ups:
            break
        cur = max(ups, key=lambda p: upa_local[p[0], p[1]])
        chain.append(cur)
    return chain


def ups_set(T):
    """Upstream catchment of pixel T: dfs_in(T) <= dfs_in <= dfs_out(T)."""
    a, b = dfin_c[T], dfout_c[T]
    return bmask & (dfin_c >= a) & (dfin_c <= b)


def boundary_8(mask):
    """Boundary pixels by the eight-neighbor test: pixels of the mask with at
    least one of their eight neighbors outside it, as in the production code
    (main_merit90m_flowmorph_step3_boundary_walk.c) and the Fig. 7 caption.
    (Until 1 Nov 2026 this was a 4-neighbor test, boundary_4.)"""
    out = np.zeros_like(mask)
    out[1:,  :] |= mask[1:,  :] & ~mask[:-1, :]
    out[:-1, :] |= mask[:-1, :] & ~mask[1:,  :]
    out[:, 1:]  |= mask[:, 1:]  & ~mask[:, :-1]
    out[:, :-1] |= mask[:, :-1] & ~mask[:, 1:]
    out[1:,  1:]  |= mask[1:,  1:]  & ~mask[:-1, :-1]
    out[1:,  :-1] |= mask[1:,  :-1] & ~mask[:-1, 1:]
    out[:-1, 1:]  |= mask[:-1, 1:]  & ~mask[1:,  :-1]
    out[:-1, :-1] |= mask[:-1, :-1] & ~mask[1:,  1:]
    out[0,  :]  |= mask[0,  :];  out[-1, :] |= mask[-1, :]
    out[:,  0]  |= mask[:, 0];   out[:, -1]  |= mask[:, -1]
    return out


# ============================================================================
# Outlet (max-UPA pixel within focus mask)
# ============================================================================
upa_in = np.where(bmask, upa_c, -1)
oi, oj = np.unravel_index(np.argmax(upa_in), upa_in.shape)


# Sanity check: ups_set(outlet) should equal bmask.  If not, fall back
# to ups_set(outlet) so border + interior always line up.
ups_out = ups_set((oi, oj))
diff = ups_out ^ bmask
print(f"|bmask|={int(bmask.sum())}  |ups_set(outlet)|={int(ups_out.sum())}  "
      f"diff={int(diff.sum())} px")
assert not diff.any(), "the outlet's interval does not hold exactly the basin"


# ============================================================================
# Main stem and two targets T1 (upstream) / T2 (downstream)
# ============================================================================
main_stem = walk_max_upa((oi, oj), bmask, upa_c, dir_c)
N_main    = len(main_stem)
print(f"main stem N = {N_main}")

# main_stem[0] = outlet, main_stem[N-1] = source.  T1 closer to source
# (upstream), T2 midstream (downstream of T1).
T1 = main_stem[(3 * N_main) // 4]
T2 = main_stem[N_main // 2]
ups_T1   = ups_set(T1)
ups_T2   = ups_set(T2)
ring_T1  = boundary_8(ups_T1)
ring_T2  = boundary_8(ups_T2)
ring_common  = ring_T1 &  ring_T2
ring_deleted = ring_T1 & ~ring_T2
ring_added   = ring_T2 & ~ring_T1
print(f"|ring(T1)|={int(ring_T1.sum())}  |ring(T2)|={int(ring_T2.sum())}  "
      f"common={int(ring_common.sum())}  deleted={int(ring_deleted.sum())}  "
      f"added={int(ring_added.sum())}")


# ============================================================================
# Skeleton extraction (UPA ≥ SKEL_KM2): greedy chain decomposition
# ============================================================================
skel_mask = bmask & (upa_c >= SKEL_KM2)
remaining = skel_mask.copy()
chains    = []
while remaining.any():
    upa_rem  = np.where(remaining, upa_c, -1)
    si, sj   = np.unravel_index(np.argmax(upa_rem), upa_rem.shape)
    chain    = walk_max_upa((si, sj), remaining, upa_c, dir_c)
    for (i, j) in chain:
        remaining[i, j] = False
    chains.append(chain)

main_idx = max(range(len(chains)), key=lambda k: upa_c[chains[k][0]])
main_chain = chains[main_idx]
ms_idx = {p: k for k, p in enumerate(reversed(main_chain))}  # source=0


def junction_to_main(chain):
    """Walk downstream from a chain's downstream end until reaching a
    main-stem pixel; return its index along the main stem (source=0)."""
    cur, seen = chain[0], set()
    for _ in range(dir_c.size):
        if cur in ms_idx: return ms_idx[cur]
        if cur in seen:   return -1
        seen.add(cur)
        d = int(dir_c[cur])
        if d not in D8_OFF: return -1
        di, dj = D8_OFF[d]
        ni, nj = cur[0] + di, cur[1] + dj
        if not (0 <= ni < dir_c.shape[0] and 0 <= nj < dir_c.shape[1]):
            return -1
        cur = (ni, nj)
    return -1


others = [c for k, c in enumerate(chains) if k != main_idx]
others_with_j = [(junction_to_main(c), c) for c in others]
others_with_j.sort(key=lambda t: t[0])
chains = [c for _, c in others_with_j] + [main_chain]
N_branches = len(chains)
print(f"N branches (skeleton) = {N_branches}")

# main-stem split for panel (d): 33 % and 67 % along source -> outlet
ms_src_to_out = list(reversed(chains[-1]))
L = len(ms_src_to_out)
cut_a = ms_src_to_out[L // 3]
cut_b = ms_src_to_out[(2 * L) // 3]
print(f"main stem split: cut@33%={cut_a}, cut@67%={cut_b}")


# ============================================================================
# Plot setup
# ============================================================================
fig, axes = plt.subplots(2, 2, figsize=(8.0, 3.8))

DFIN_DISP = np.where(bmask & (dfin_c >= 0), dfin_c, np.nan).astype(np.float64)
DFIN_VMIN = float(np.nanpercentile(DFIN_DISP, 1))
DFIN_VMAX = float(np.nanpercentile(DFIN_DISP, 99))
BASIN_BND = boundary_8(bmask)
BB_R, BB_C = np.where(BASIN_BND)


def base(ax, draw_skel=True):
    """All panels share: dfs_in viridis fill + skeleton lines + basin
    border (low-zorder so panel-specific overlays sit on top) + outlet ▼."""
    ax.set_facecolor("white")
    ax.imshow(DFIN_DISP, cmap="viridis", vmin=DFIN_VMIN, vmax=DFIN_VMAX,
              origin="upper", interpolation="nearest", alpha=0.9, zorder=1)
    if draw_skel:
        for ch in chains:
            rr = np.array([p[0] for p in ch])
            cc = np.array([p[1] for p in ch])
            ax.plot(cc, rr, "-", color="white",
                    linewidth=0.6, alpha=0.95, zorder=3)
    ax.scatter(BB_C, BB_R, s=3, c="black", marker="s",
               linewidths=0, alpha=1.0, zorder=5)
    ax.plot(oj, oi, marker="v", color="#E53935", markersize=6,
            markeredgecolor="black", markeredgewidth=0.6,
            linestyle="none", zorder=12)
    ax.set_xlim(-0.5, W - 0.5); ax.set_ylim(H - 0.5, -0.5)
    ax.set_aspect("equal")
    ax.set_xticks([]); ax.set_yticks([])
    for sp in ax.spines.values():
        sp.set_linewidth(0.5); sp.set_color("#888888")


def scatter_pixels(ax, mask, color, size=3.5, marker="s",
                   alpha=0.9, zorder=5):
    rr, cc = np.where(mask)
    ax.scatter(cc, rr, s=size, c=color, marker=marker,
               linewidths=0, alpha=alpha, zorder=zorder)


LABEL_BOX = dict(boxstyle="round,pad=0.16", fc="white", ec="#888888",
                 lw=0.4, alpha=0.95)


# ----------------------------------------------------------------------------
# (a) ring(T₁) — split into two parts per manuscript caption:
#   gray = segment on the basin perimeter (BASIN_BND ∩ ring_T1)
#   red  = new closing edge interior to the basin (ring_T1 \ BASIN_BND)
# ----------------------------------------------------------------------------
ring_T1_on_perim    = ring_T1 &  BASIN_BND
ring_T1_closing_edge = ring_T1 & ~BASIN_BND
ax = axes[0, 0]
base(ax)
scatter_pixels(ax, ring_T1_on_perim,    "#9aa0a6", size=3.5, alpha=1.0, zorder=7)
scatter_pixels(ax, ring_T1_closing_edge, "#dc2626", size=4.5, alpha=1.0, zorder=9)
ax.plot(T1[1], T1[0], marker="o", color="black", markersize=5,
        linestyle="none", zorder=13)
ax.text(T1[1] + 6, T1[0] + 6, r"$T_1$", color="black",
        fontsize=8.5, weight="bold",
        ha="left", va="top", bbox=LABEL_BOX, zorder=14)
ax.legend(handles=[
    Patch(facecolor="#9aa0a6", alpha=0.85, label="on basin perimeter"),
    Patch(facecolor="#dc2626", alpha=0.95, label="closing edge"),
], loc="lower right", fontsize=7.5, frameon=True,
   edgecolor="#888", framealpha=0.95)
ax.set_title(r"(a) Boundary of upstream$(T_1)$", loc="left", pad=3)


# ----------------------------------------------------------------------------
# (b) ring(T₁) ↔ ring(T₂) symmetric difference
# ----------------------------------------------------------------------------
ax = axes[0, 1]
base(ax)
scatter_pixels(ax, ring_common,  "#9aa0a6", size=3.5, alpha=0.95, zorder=7)
scatter_pixels(ax, ring_deleted, "#dc2626", size=4.0, alpha=0.95, zorder=8)
scatter_pixels(ax, ring_added,   "#16a34a", size=4.0, alpha=0.95, zorder=8)
for tgt, lab in [(T1, r"$T_1$"), (T2, r"$T_2$")]:
    ax.plot(tgt[1], tgt[0], marker="o", color="black",
            markersize=5, linestyle="none", zorder=13)
    ax.text(tgt[1] + 6, tgt[0] + 6, lab, color="black",
            fontsize=8.5, weight="bold",
            ha="left", va="top", bbox=LABEL_BOX, zorder=14)
ax.legend(handles=[
    Patch(facecolor="#9aa0a6", alpha=0.85, label="common"),
    Patch(facecolor="#dc2626", alpha=0.95, label="deleted"),
    Patch(facecolor="#16a34a", alpha=0.95, label="added"),
], loc="lower right", fontsize=7.5, frameon=True,
   edgecolor="#888", framealpha=0.95)
ax.set_title(r"(b) Move target to $T_2$", loc="left", pad=3)


def downstream_arrow(ax, pos, color="black", length=10.0, lw=1.4, zorder=11):
    d = int(dir_c[pos])
    if d not in D8_OFF: return
    di, dj = D8_OFF[d]
    norm = (di * di + dj * dj) ** 0.5
    tip  = (pos[1] + dj * length / norm, pos[0] + di * length / norm)
    base_pt = (pos[1], pos[0])
    ax.annotate("", xy=tip, xytext=base_pt,
                arrowprops=dict(arrowstyle="-|>", color=color,
                                lw=lw, mutation_scale=14,
                                shrinkA=2, shrinkB=0),
                zorder=zorder)


# ----------------------------------------------------------------------------
# (c) Branch tips
# ----------------------------------------------------------------------------
ax = axes[1, 0]
base(ax, draw_skel=True)
LABEL_FS = 8.5
for k, ch in enumerate(chains, 1):
    head = ch[-1]
    downstream_arrow(ax, head, color="black", length=10.0, lw=1.4)
    ax.plot(head[1], head[0], marker="o", color="black",
            markersize=4, linestyle="none", zorder=12)
    ax.text(head[1] + 2, head[0] - 2, f"S{k}", color="black",
            fontsize=LABEL_FS, weight="bold",
            ha="left", va="bottom", zorder=13)
ax.set_title(rf"(c) {N_branches} branch tips (S$_1$…S$_{{{N_branches}}}$)",
             loc="left", pad=3)


# ----------------------------------------------------------------------------
# (d) Main stem split into S_N, S_{N+1}, S_{N+2}
# ----------------------------------------------------------------------------
ax = axes[1, 1]
base(ax, draw_skel=True)
for k, ch in enumerate(chains, 1):
    head = ch[-1]
    downstream_arrow(ax, head, color="black", length=10.0, lw=1.4)
    ax.plot(head[1], head[0], marker="o", color="black",
            markersize=4, linestyle="none", zorder=12)
    ax.text(head[1] + 2, head[0] - 2, f"S{k}", color="black",
            fontsize=LABEL_FS, weight="bold",
            ha="left", va="bottom", zorder=13)

SPLIT_BLUE = "#2563EB"
for k, cut in enumerate([cut_a, cut_b], 1):
    downstream_arrow(ax, cut, color=SPLIT_BLUE, length=10.0,
                     lw=1.6, zorder=14)
    ax.plot(cut[1], cut[0], marker="o", color=SPLIT_BLUE,
            markersize=5, linestyle="none", zorder=15)
    ax.text(cut[1] + 3, cut[0] - 3, f"S{N_branches + k}",
            color=SPLIT_BLUE, fontsize=LABEL_FS + 0.5, weight="bold",
            ha="left", va="bottom", zorder=16)
ax.set_title(rf"(d) Split main stem S$_{{{N_branches}}}$ into S$_{{{N_branches}}}$, "
             rf"S$_{{{N_branches+1}}}$, S$_{{{N_branches+2}}}$",
             loc="left", pad=3)


# ============================================================================
# Save
# ============================================================================
fig.subplots_adjust(left=0.015, right=0.985, top=0.93, bottom=0.02,
                    hspace=0.18, wspace=0.04)

for ext in ("png", "pdf"):
    out = os.path.join(OUT_DIR, f"Fig07_BoundaryWalk_v11.{ext}")
    fig.savefig(out, dpi=300, bbox_inches="tight", facecolor="white")
    print(f"Saved: {out}")
plt.close(fig)