#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Figure A1: wall-clock time of the global basin delineation.

Input: ../Data/stats/time/.
Output: ../Figures/FigA1_Step12_Timing_v1.png and .pdf.

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

from __future__ import annotations

import re
from pathlib import Path
from datetime import datetime, timedelta

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt


# ============================================================================
# Style
# ============================================================================
mpl.rcParams.update({
    "font.family":      "Helvetica",
    "font.size":        10,
    "axes.labelsize":   11,
    "legend.fontsize":  9,
    "xtick.labelsize":  9.5,
    "ytick.labelsize":  9,
    "mathtext.fontset": "cm",
    "axes.linewidth":   0.8,
})

# ESSD-consistent appendix palette
C_IO    = "#BFBFBF"   # light gray for I/O bookends (DIR read, BSN write)
C_EDGE  = "black"
# Pass 1-6 use viridis(0.05 → 0.85) — deep (heavy work) to bright (light work)
_PASS_VIRIDIS = plt.get_cmap("viridis")(np.linspace(0.10, 0.85, 6))


# ============================================================================
# Paths
# ============================================================================
HERE      = Path(__file__).resolve().parent
STATS_DIR = HERE.parent / "Data" / "stats" / "time"
FIG_DIR   = HERE.parent / "Figures"
FIG_DIR.mkdir(parents=True, exist_ok=True)

LOG_RUN   = STATS_DIR / "run_delin_693.log"
LOG_SUM   = STATS_DIR / "Step12_BasinDelineation.txt"


# ============================================================================
# Log parsing
# ============================================================================
def _hms_to_seconds(hms: str) -> int:
    h, m, s = map(int, hms.split(":"))
    return h * 3600 + m * 60 + s


def parse_run_log(path: Path):
    """Return dict with per-stage wall-clock seconds.

    Stages: 'DIR_read', 'pass1'..'pass6', 'BSN_write'.
    Strategy:
      - start_of_run    = 'Start: YYYY-MM-DD HH:MM:SS'
      - first_pass_k_ts = first 'Time_HH:MM:SS' line after '===== Pass k ====='
      - end_of_passes   = last 'Time_HH:MM:SS' line inside Pass 6 block
      - end_of_run      = 'End: YYYY-MM-DD HH:MM:SS'
      - DIR_read = first_pass_1_ts - start_of_run
      - pass_k   = first_pass_{k+1}_ts - first_pass_k_ts    (k<6)
      - pass_6   = end_of_passes - first_pass_6_ts
      - BSN_write= end_of_run - end_of_passes
    """
    re_start  = re.compile(r"\[step12_delin\] Start:\s+\S+\s+(\d{2}:\d{2}:\d{2})")
    re_end    = re.compile(r"\[step12_delin\] End:\s+\S+\s+(\d{2}:\d{2}:\d{2})")
    re_pass   = re.compile(r"===== Pass (\d+):")
    re_time   = re.compile(r"Time_(\d{2}:\d{2}:\d{2})\s+\[pass(\d+)\]")

    start_sec = None
    end_sec   = None
    # first_ts_by_pass[k] = seconds of first Time_... line whose [passK] matches k
    first_ts_by_pass: dict[int, int] = {}
    last_ts_by_pass:  dict[int, int] = {}

    with open(path) as f:
        for line in f:
            m = re_start.search(line)
            if m and start_sec is None:
                start_sec = _hms_to_seconds(m.group(1))
                continue
            m = re_end.search(line)
            if m:
                end_sec = _hms_to_seconds(m.group(1))
                continue
            m = re_time.search(line)
            if m:
                ts = _hms_to_seconds(m.group(1))
                k  = int(m.group(2))
                if k not in first_ts_by_pass:
                    first_ts_by_pass[k] = ts
                last_ts_by_pass[k] = ts

    assert start_sec is not None, "missing Start: line"
    assert end_sec   is not None, "missing End: line"
    assert set(first_ts_by_pass.keys()) == {1, 2, 3, 4, 5, 6}, \
        f"unexpected passes: {sorted(first_ts_by_pass)}"

    # Handle day wrap (unlikely but safe)
    def _delta(a: int, b: int) -> int:
        d = b - a
        return d + 24 * 3600 if d < 0 else d

    stages = {}
    stages["DIR_read"] = _delta(start_sec, first_ts_by_pass[1])
    for k in range(1, 6):
        stages[f"pass{k}"] = _delta(first_ts_by_pass[k], first_ts_by_pass[k + 1])
    stages["pass6"]    = _delta(first_ts_by_pass[6], last_ts_by_pass[6])
    stages["BSN_write"] = _delta(last_ts_by_pass[6], end_sec)

    return stages


def parse_completion(path: Path):
    """Return two lists of length 6: msk_pct, bsn_pct (5-decimal precision,
    computed from raw pixel counts rather than the rounded 2-decimal % in
    the log).

    Handles both Step12 summary formats:
      pass1 form: '[pass1] MSK==2: <count> (...%)  BSN!=0: <count> (...%)'
      pass2+:     '[pass2] active=40/50  MSK==2: <count> (...) BSN!=0: <count> (...)'
    """
    re_total = re.compile(r"total_expected(?:\s+pixels)?:\s+(\d+)")
    re_pass  = re.compile(r"\[pass(\d+)\]")
    re_cnt   = re.compile(
        r"MSK==2:\s+(\d+)\s+\([\d.]+%\)\s+"
        r"BSN!=0:\s+(\d+)\s+\([\d.]+%\)"
    )
    total = None
    msk_cnt = [None] * 6
    bsn_cnt = [None] * 6
    with open(path) as f:
        for line in f:
            if total is None:
                m = re_total.search(line)
                if m:
                    total = int(m.group(1))
            m1 = re_pass.search(line)
            m2 = re_cnt.search(line)
            if m1 and m2:
                k = int(m1.group(1))
                if 1 <= k <= 6:
                    msk_cnt[k - 1] = int(m2.group(1))
                    bsn_cnt[k - 1] = int(m2.group(2))
    assert total is not None, "missing total_expected"
    assert all(v is not None for v in msk_cnt), f"missing MSK cnt: {msk_cnt}"
    assert all(v is not None for v in bsn_cnt), f"missing BSN cnt: {bsn_cnt}"
    msk = [c / total * 100.0 for c in msk_cnt]
    bsn = [c / total * 100.0 for c in bsn_cnt]
    return msk, bsn


# ============================================================================
# Plot
# ============================================================================
def main():
    stages = parse_run_log(LOG_RUN)
    msk_pct, bsn_pct = parse_completion(LOG_SUM)

    # --- Build bar data (minutes) ---
    bar_labels_full = [
        "DIR read\n(+ decompress)",
        "Pass 1\n(outlet\ntop 30)",
        "Pass 2\n(grid\n0,0)",
        "Pass 3\n(grid\n21600,0)",
        "Pass 4\n(grid\n0,21600)",
        "Pass 5\n(grid\n21600,\n21600)",
        "Pass 6\n(anti-\nmeridian)",
        "BSN write\n(+ compress)",
    ]
    keys   = ["DIR_read", "pass1", "pass2", "pass3",
              "pass4", "pass5", "pass6", "BSN_write"]
    times_min = np.array([stages[k] / 60.0 for k in keys])

    total_min = times_min.sum()
    total_h   = total_min / 60.0

    # --- Figure layout ---
    fig, (ax_a, ax_b) = plt.subplots(
        nrows=2, ncols=1,
        figsize=(7.8, 6.3),
        gridspec_kw={"height_ratios": [1.1, 1.0], "hspace": 0.55},
    )

    # ---------- Panel (a) Per-stage wall-clock ----------
    # bar_colors order matches `keys`: [DIR_read, pass1..pass6, BSN_write]
    bar_colors = (
        [C_IO] +                       # DIR_read → light gray
        list(_PASS_VIRIDIS) +          # Pass 1-6 → viridis dark→bright
        [C_IO]                         # BSN_write → light gray
    )
    x_bar = np.arange(len(keys))
    ax_a.bar(
        x_bar, times_min,
        width=0.7,
        color=bar_colors,
        edgecolor=C_EDGE,
        linewidth=0.7,
    )
    for i, t in enumerate(times_min):
        ax_a.text(
            x_bar[i], t + total_min * 0.012,
            f"{t:.1f}",
            ha="center", va="bottom",
            fontsize=8.5, color="black",
        )

    ax_a.set_xticks(x_bar)
    ax_a.set_xticklabels(bar_labels_full, fontsize=8.2)
    ax_a.set_ylabel("Wall-clock time  (min)")
    ax_a.set_ylim(0, times_min.max() * 1.18)
    ax_a.grid(True, axis="y", lw=0.3, alpha=0.3)
    ax_a.spines["top"].set_visible(False)
    ax_a.spines["right"].set_visible(False)

    # Total annotation (upper-right, away from Pass 1's numeric label)
    ax_a.text(
        0.985, 0.96,
        f"Total  {total_min:.1f} min  ({total_h:.2f} h)",
        transform=ax_a.transAxes,
        ha="right", va="top",
        fontsize=10, fontweight="bold",
    )
    ax_a.text(
        -0.06, 1.03,
        "(a)",
        transform=ax_a.transAxes,
        ha="left", va="bottom",
        fontsize=12, fontweight="bold",
    )

    # ---------- Panel (b) Cumulative completion ----------
    passes = np.arange(1, 7)
    ax_b.plot(
        passes, msk_pct,
        color="black", linestyle="-",  linewidth=1.4,
        marker="o", markersize=6,
        markerfacecolor="black", markeredgecolor="black",
        label="MSK == 2  (pixel's basin fully closed)",
    )
    ax_b.plot(
        passes, bsn_pct,
        color="black", linestyle="--", linewidth=1.4,
        marker="s", markersize=6,
        markerfacecolor="white", markeredgecolor="black",
        markeredgewidth=1.1,
        label="BSN $\\neq$ 0  (pixel assigned to a basin)",
    )

    # inline numeric labels (5-decimal %, offset to avoid overlap)
    for xp, y in zip(passes, msk_pct):
        ax_b.text(xp, y - 3.2, f"{y:.5f}%",
                  ha="center", va="top", fontsize=7.5, color="black")
    for xp, y in zip(passes, bsn_pct):
        ax_b.text(xp, y + 1.6, f"{y:.5f}%",
                  ha="center", va="bottom", fontsize=7.5, color="black")

    ax_b.set_xticks(passes)
    ax_b.set_xticklabels([f"Pass {k}" for k in passes])
    ax_b.set_xlabel("Delineation pass")
    ax_b.set_ylabel("Cumulative completion  (%)")
    ax_b.set_ylim(50, 104)
    ax_b.set_xlim(0.7, 6.3)
    ax_b.grid(True, axis="y", lw=0.3, alpha=0.3)
    ax_b.spines["top"].set_visible(False)
    ax_b.spines["right"].set_visible(False)
    ax_b.legend(loc="lower right", frameon=True,
                framealpha=0.95, edgecolor="0.7")

    ax_b.text(
        -0.06, 1.03,
        "(b)",
        transform=ax_b.transAxes,
        ha="left", va="bottom",
        fontsize=12, fontweight="bold",
    )

    # ---------- Save ----------
    fig.subplots_adjust(left=0.09, right=0.98, top=0.95, bottom=0.08)
    stem = "FigA1_Step12_Timing_v1"
    for ext in ("png", "pdf"):
        out = FIG_DIR / f"{stem}.{ext}"
        fig.savefig(out, dpi=300, bbox_inches="tight", facecolor="white")
        print(f"  wrote {out}")

    # ---------- Console sanity ----------
    print("\n--- Per-stage times ---")
    for k, t in zip(keys, times_min):
        print(f"  {k:10s} {t:7.2f} min")
    print(f"  {'TOTAL':10s} {total_min:7.2f} min ({total_h:.2f} h)")
    print("\n--- Coverage (%) ---")
    for k, (m, b) in enumerate(zip(msk_pct, bsn_pct), 1):
        print(f"  pass{k}  MSK={m:6.2f}%  BSN={b:6.2f}%")


if __name__ == "__main__":
    main()
