#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Figure A2: wall-clock time of the four computational kernels.

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

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

from pathlib import Path

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":  10,
    "ytick.labelsize":  9,
    "mathtext.fontset": "cm",
    "axes.linewidth":   0.8,
})


# ============================================================================
# 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)


# ============================================================================
# Timing-file parser
# ============================================================================
def parse_kv(path):
    """Sum io / core / write / total (s) across all lines."""
    sums = {"io": 0.0, "core": 0.0, "write": 0.0, "total": 0.0}
    n = 0
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            for token in line.split():
                if "=" not in token:
                    continue
                k, v = token.split("=", 1)
                v = v.rstrip("s")
                try:
                    v = float(v)
                except ValueError:
                    continue
                if k in sums:
                    sums[k] += v
            n += 1
    return n, sums


def kernel_stats(path):
    """Return (n, io_h, compute_h, total_h)."""
    n, s      = parse_kv(path)
    io_s      = s["io"] + s["write"]
    total_s   = s["total"]
    compute_s = max(0.0, total_s - io_s)
    return n, io_s / 3600.0, compute_s / 3600.0, total_s / 3600.0


# ============================================================================
# Main
# ============================================================================
def main():
    # (display label, sub-label, timing file)
    kernels = [
        ("Derive DFS",
         "dfs_in / dfs_out\n(2 topological indices)",
         "derive_dfs_interval_time.txt"),
        ("Upstream aggregation",
         "mean / min / max / std elv +\nmean / std slp\n(6 pixel attributes)",
         "relief_hypsometry_pixel_time.txt"),
        ("Boundary walk",
         "basin length / basin width /\nperimeter / convexity\n(4 pixel attributes)",
         "exp_bf_boundary_axial_time.txt"),
        ("Arithmetic",
         "circularity / gravelius / elongation /\nlemniscate + relief + HI\n(6 pixel attributes)",
         "exp_bf_arith_shape_time.txt"),
    ]

    labels, sublabels, io_h, cp_h, tot_h, ns = [], [], [], [], [], []
    for label, sub, fname in kernels:
        n, io, cp, tot = kernel_stats(STATS_DIR / fname)
        labels.append(label)
        sublabels.append(sub)
        io_h.append(io)
        cp_h.append(cp)
        tot_h.append(tot)
        ns.append(n)
        print(f"  {label:28s}  n={n:3d}  I/O={io:5.2f}h  "
              f"Compute={cp:5.2f}h  Total={tot:5.2f}h")

    io_h  = np.array(io_h)
    cp_h  = np.array(cp_h)
    tot_h = np.array(tot_h)

    # ---- Plot: 4 vertical stacked bars ----
    fig, ax = plt.subplots(figsize=(7.6, 5.0))

    x     = np.arange(len(labels))
    width = 0.55
    c_io  = "#D9D9D9"        # light gray for I/O
    c_cp  = "#21918c"        # viridis-deep teal for Compute (paper accent)

    ax.bar(x, io_h, width, color=c_io,
           edgecolor="black", lw=0.7,
           label="I/O  (GeoTIFF read + write, incl. (de)compression)")
    ax.bar(x, cp_h, width, bottom=io_h, color=c_cp,
           edgecolor="black", lw=0.7,
           label="Compute  (algorithm + stream setup + allocation)")

    # Total label on top
    for i, t in enumerate(tot_h):
        ax.text(x[i], t + max(tot_h) * 0.06,
                f"Total  {t:.2f} h",
                ha="center", va="bottom",
                fontsize=10, fontweight="bold")

    # Inline numbers (no percentages)
    for i in range(len(labels)):
        if io_h[i] >= 0.3:
            ax.text(x[i], io_h[i] / 2,
                    f"{io_h[i]:.2f} h",
                    ha="center", va="center", fontsize=9, color="black")
        if cp_h[i] >= 0.3:
            ax.text(x[i], io_h[i] + cp_h[i] / 2,
                    f"{cp_h[i]:.2f} h",
                    ha="center", va="center", fontsize=9,
                    color="white", fontweight="bold")
        # (white text on dark gray still readable)
        elif cp_h[i] > 0:
            ax.text(x[i], io_h[i] + cp_h[i] + max(tot_h) * 0.015,
                    f"Compute  {cp_h[i]:.2f} h",
                    ha="center", va="bottom", fontsize=8, color="0.2")

    # Multi-line xticklabels
    xtick_labels = [f"{lab}\n\n{sub}" for lab, sub in zip(labels, sublabels)]
    ax.set_xticks(x)
    ax.set_xticklabels(xtick_labels, fontsize=8.5)

    ax.set_ylabel("Global wall-clock time  (hours)")
    ax.set_ylim(0, max(tot_h) * 1.20)
    ax.legend(loc="upper left", frameon=True, framealpha=0.95,
              edgecolor="0.7")
    ax.grid(True, axis="y", lw=0.3, alpha=0.3)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    fig.subplots_adjust(bottom=0.30, top=0.96, left=0.09, right=0.98)

    stem = "FigA2_Kernel_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}")


if __name__ == "__main__":
    main()