#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Figure 1: the five steps of the global basin delineation.

No input data. Output: ../Figures/Fig01_Diagram.png and .pdf.

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

from pathlib import Path

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import font_manager as fm
from matplotlib.patches import FancyBboxPatch, Rectangle

HERE = Path(__file__).resolve().parent
OUT_DIR = HERE.parent / "Figures"

# --- fonts --------------------------------------------------------------------
for d in (Path("/Applications/Microsoft Word.app/Contents/Resources/DFonts"),
          Path("/Applications/Microsoft PowerPoint.app/Contents/Resources/DFonts")):
    for f in ("Consola.ttf", "Consolab.ttf", "Consolai.ttf"):
        if (d / f).exists():
            fm.fontManager.addfont(str(d / f))
_names = {f.name for f in fm.fontManager.ttflist}
SANS = "Arial" if "Arial" in _names else "DejaVu Sans"
MONO = "Consolas" if "Consolas" in _names else ("Menlo" if "Menlo" in _names else "DejaVu Sans Mono")
mpl.rcParams["pdf.fonttype"] = 42

# --- colours of the pptx --------------------------------------------------------
INK, BROWN, RED = "#1A1A1A", "#6B4F1F", "#C00000"
HEAD_FILL, HEAD_EDGE = "#2F5C7E", "#1F3F58"
BODY_FILL, BODY_EDGE = "#EAF0F5", "#5A7A95"
NOTE_FILL, NOTE_EDGE = "#F5EFE0", "#B89968"
BAND_FILL = "#F0EDE3"
ARROW = "#404040"

W, H = 6.77, 7.02                  # inches; drawing units are inches, y downward
X0 = 0.05                          # left edge of the content (pptx x = 1.49 in)


def style(font="sans", bold=False, italic=False, size=9.0, color=INK):
    return dict(family=SANS if font == "sans" else MONO, weight="bold" if bold else "normal",
                style="italic" if italic else "normal", size=size, color=color)


class Canvas:
    def __init__(self):
        self.fig = plt.figure(figsize=(W, H), facecolor="white")
        self.ax = self.fig.add_axes([0, 0, 1, 1])
        self.ax.set_xlim(0, W)
        self.ax.set_ylim(H, 0)
        self.ax.axis("off")
        self.r = self.fig.canvas.get_renderer()

    def runs(self, x, y, runs, ha="center"):
        """One line of text made of runs of different styles, centred (or left-aligned) at x, middle at y."""
        texts, widths = [], []
        for s, st in runs:
            t = self.ax.text(0, y, s, va="center", ha="left", fontfamily=st["family"],
                             fontweight=st["weight"], fontstyle=st["style"], fontsize=st["size"],
                             color=st["color"], zorder=5)
            texts.append(t)
            widths.append(t.get_window_extent(self.r).width / self.fig.dpi)
        total = sum(widths)
        cur = x - total / 2.0 if ha == "center" else x
        for t, w in zip(texts, widths):
            t.set_x(cur)
            cur += w

    def box(self, x, y, w, h, fill, edge, lw=0.75):
        self.ax.add_patch(Rectangle((x, y), w, h, fc=fill, ec=edge, lw=lw, zorder=2))

    def lines(self, x, y, w, h, rows, lead=1.3):
        """Rows of runs centred in a box; lead is the line height in units of 9 pt."""
        step = lead * 9 / 72.0
        y0 = y + h / 2.0 - step * (len(rows) - 1) / 2.0
        for i, row in enumerate(rows):
            self.runs(x + w / 2.0, y0 + i * step, row)

    def header_body(self, x, y, w, head, rows, h_head=0.19, h_body=0.56, lead=1.3):
        self.box(x, y, w, h_head, HEAD_FILL, HEAD_EDGE)
        self.runs(x + w / 2.0, y + h_head / 2.0, [(head, style(bold=True, color="white"))])
        self.box(x, y + h_head, w, h_body, BODY_FILL, BODY_EDGE)
        self.lines(x, y + h_head, w, h_body, rows, lead)

    def arrow(self, x0, y0, x1, y1):
        self.ax.annotate("", xy=(x1, y1), xytext=(x0, y0), zorder=3,
                         arrowprops=dict(arrowstyle="->", lw=0.75, color=ARROW, shrinkA=0, shrinkB=0,
                                         mutation_scale=9))

    def title(self, y, bold_text, note):
        self.runs(X0, y, [(bold_text, style(bold=True, size=10.5)), (note, style(italic=True, color=BROWN))],
                  ha="left")


def mono(s, **k):
    return (s, style("mono", **k))


def sans(s, **k):
    return (s, style("sans", **k))


def bid():
    return mono("Basin_ID − 1", bold=True, color=BROWN)


def main():
    c = Canvas()
    COL2, COL3, BW = X0 + 2.36, X0 + 4.72, 1.94       # three columns of Steps 1, 4, 5

    # --- index convention banner ---------------------------------------------
    c.box(X0, 0.05, 6.67, 0.36, NOTE_FILL, NOTE_EDGE)
    c.runs(X0 + 0.12, 0.23, [sans("Index convention:  ", bold=True, color=BROWN),
                             mono("Basin_ID = Table_Index + 1", bold=True, size=10.5, color=RED),
                             sans("     →   one-step translation between raster pixels and table rows",
                                  italic=True, size=8, color=BROWN)], ha="left")

    # --- Step 1 ----------------------------------------------------------------
    c.title(0.57, "Step 1.  Plant outlet seeds  ", "(BAT → BSN, once before any pass)")
    y = 0.75
    c.header_body(X0, y, BW, "Attribute Table (BAT)",
                  [[mono("BAT[Table_Index] =")], [mono("{Basin_ID, outlet_row,")], [mono("outlet_col, …}")]])
    c.header_body(COL3, y, BW, "Global 2D Raster (BSN)",
                  [[mono("BSN[outlet_row,")], [mono("outlet_col] = Basin_ID")], [mono("others = NoData")]])
    ym = y + 0.19 + 0.28
    c.arrow(X0 + BW, ym, COL3, ym)
    c.runs((X0 + BW + COL3) / 2.0, ym - 0.10, [sans("for each row in BAT", italic=True, size=8, color=BROWN)])

    # --- Step 2 ----------------------------------------------------------------
    c.title(1.70, "Step 2.  Sliding-window sweep schedule  ", "(6 passes; 36° windows; Pass 2–5 step 18°)")
    y, hh, hb, pw = 1.88, 0.17, 0.48, 1.01
    xs = [X0, X0 + 1.12, X0 + 2.25, X0 + 3.37, X0 + 4.50, X0 + 5.62]
    c.ax.add_patch(FancyBboxPatch((xs[1] - 0.05, y - 0.03), xs[5] + pw - xs[1] + 0.10, hh + hb + 0.06,
                                  boxstyle="round,pad=0,rounding_size=0.06", fc=BAND_FILL, ec=NOTE_EDGE,
                                  lw=0.5, ls="--", zorder=1))
    bodies = [["centered on", "30 largest", "basins"], ["baseline", "(0°, 0°)"], ["lon shift", "(+18°, 0°)"],
              ["lat shift", "(0°, +18°)"], ["diagonal", "(+18°, +18°)"], ["anti-meridian", "wrap-around"]]
    for k, (x, rows) in enumerate(zip(xs, bodies)):
        c.box(x, y, pw, hh, HEAD_FILL, HEAD_EDGE)
        c.runs(x + pw / 2.0, y + hh / 2.0, [sans(f"Pass {k + 1}", bold=True, color="white")])
        c.box(x, y + hh, pw, hb, BODY_FILL, BODY_EDGE)
        c.lines(x, y + hh, pw, hb, [[sans(r, size=8)] for r in rows], lead=1.2)
        if k:
            c.arrow(xs[k - 1] + pw, y + hh + hb / 2.0, x, y + hh + hb / 2.0)

    # --- Step 3 ----------------------------------------------------------------
    c.title(2.70, "Step 3.  Mark complete basins as “closed”  ", "(at the start of each pass)")
    y, h = 2.88, 0.44
    xo, wo, xe, we = X0, 2.85, X0 + 3.22, 3.44
    c.box(xo, y, wo, h, BODY_FILL, BODY_EDGE)
    c.lines(xo, y, wo, h, [[sans("observed", bold=True), sans(" (labels currently in BSN)", italic=True, color=BROWN)],
                           [mono("labeled_count["), bid(), mono("]")]], lead=1.45)
    c.box(xe, y, we, h, BODY_FILL, BODY_EDGE)
    c.lines(xe, y, we, h, [[sans("expected", bold=True),
                            sans(" (UPG, downstream-routed pixel count)", italic=True, color=BROWN)],
                           [mono("BAT["), bid(), mono("].basin_cell_count")]], lead=1.45)
    yc = 3.48
    c.arrow(xo + 1.61, y + h, xo + 1.61, yc)
    c.arrow(X0 + 5.05, y + h, X0 + 5.05, yc)
    c.box(X0, yc, 6.67, 0.30, NOTE_FILL, NOTE_EDGE)
    c.runs(X0 + 6.67 / 2.0, yc + 0.15, [sans("Compare:  ", bold=True),
                                        mono("labeled_count[Basin_ID − 1]", bold=True, size=10, color=RED),
                                        sans("  =  ", bold=True, size=10, color=RED),
                                        mono("BAT[Basin_ID − 1].basin_cell_count", bold=True, size=10, color=RED),
                                        sans("   ?", bold=True, size=10)])
    ym = 3.92
    c.arrow(X0 + 3.33, yc + 0.30, X0 + 3.33, ym)
    c.box(X0, ym, 6.67, 0.30, BODY_FILL, BODY_EDGE)
    c.runs(X0 + 6.67 / 2.0, ym + 0.15, [mono("MSK = 2 (closed)", bold=True),
                                        sans("   mark basin and skip it in subsequent passes")])

    # --- Step 4 ----------------------------------------------------------------
    c.title(4.42, "Step 4.  Per-tile upstream tracing  ", "(inside each pass, every tile)")
    y, hh, hb = 4.60, 0.20, 0.56
    c.header_body(X0, y, BW, "Global 2D Input",
                  [[mono("MSK"), sans("  (read-only)", italic=True, size=8, color=BROWN)],
                   [mono("DIR"), sans("  (read-only)", italic=True, size=8, color=BROWN)],
                   [mono("BSN"), sans("  (read)", italic=True, size=8, color=BROWN)]], hh, hb)
    c.header_body(COL2, y, BW, "Local Tile (in-memory)",
                  [[mono("mskTile = MSK[window]")], [mono("dirTile = DIR[window]")], [mono("bsnTile = BSN[window]")]],
                  hh, hb)
    c.header_body(COL3, y, BW, "Global 2D Write-back (BSN)",
                  [[mono("BSN[window] = bsnTile")], [sans("(updated labels only)", italic=True, size=8, color=BROWN)]],
                  hh, hb)
    ymid = y + hh + hb / 2.0
    c.arrow(X0 + BW, ymid, COL2, ymid)
    c.arrow(COL2 + BW, ymid, COL3, ymid)
    yi = 5.54
    c.arrow(COL2 + BW / 2.0, y + hh + hb, COL2 + BW / 2.0, yi)
    xi, wi = X0 + 1.77, 3.92
    c.box(xi, yi, wi, 0.32, NOTE_FILL, NOTE_EDGE)
    c.runs(xi + wi / 2.0, yi + 0.16, [mono("if mskTile == 1 (open):", color=BROWN),
                                      sans("   trace and label upstream pixels", italic=True, color=BROWN)])

    # --- Step 5 ----------------------------------------------------------------
    c.title(6.03, "Step 5.  Record bounding boxes  ", "(BSN → BAT, once after all basins delineated)")
    y = 6.21
    c.header_body(X0, y, BW, "Attribute Table (BAT)",
                  [[mono("BAT["), bid(), mono("].{")], [mono("row_min..max,")], [mono("col_min..max}")]], 0.19, 0.56)
    c.header_body(COL3, y, BW, "Global 2D Raster (BSN)", [[mono("BSN[row, col] = Basin_ID")]], 0.19, 0.56)
    ym = y + 0.19 + 0.28
    c.arrow(COL3, ym, X0 + BW, ym)
    c.runs((X0 + BW + COL3) / 2.0, ym - 0.10, [sans("for each labeled pixel", italic=True, size=8, color=BROWN)])

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    for ext in ("png", "pdf"):
        p = OUT_DIR / f"Fig01_Diagram.{ext}"
        c.fig.savefig(p, dpi=340, facecolor="white")
        print(f"wrote {p}")
    print(f"fonts: {SANS}, {MONO}")


if __name__ == "__main__":
    main()
