#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Figure 6: global dfs_in and dfs_out at 1/20 degree, and the example basin 8595 at 3 arcsec.

Input: ../Data/dfs_in_global_020th_v11.tif, dfs_out_global_020th_v11.tif, rgn_global_200th_v11.gpkg, ex8595_*.tif, ex8595_record.json.
Output: ../Figures/Fig06_DFS_Interval_Global_v11.png and .pdf.

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

import os

import json

import numpy as np
import rasterio
from rasterio.features import shapes as raster_shapes
from shapely.geometry import shape as shp_shape
from shapely.ops import unary_union
from rasterio.enums import Resampling
import geopandas as gpd

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from matplotlib.gridspec import GridSpec
from mpl_toolkits.axes_grid1 import make_axes_locatable


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


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

G_DFSIN  = os.path.join(DATA, "dfs_in_global_020th_v11.tif")
G_DFSOUT = os.path.join(DATA, "dfs_out_global_020th_v11.tif")
G_RGN    = os.path.join(DATA, "rgn_global_200th_v11.gpkg")
E_DFSIN  = os.path.join(DATA, "ex8595_dfs_in.tif")
E_DFSOUT = os.path.join(DATA, "ex8595_dfs_out.tif")
E_BSN    = os.path.join(DATA, "ex8595_bsn.tif")
E_REC    = os.path.join(DATA, "ex8595_record.json")

FOCUS_BSN_ID = 8595      # global_basin_id (the v1.1 basin raster holds global ids)


# ============================================================================
# Region boundaries (drawn on global panels a/c)
# ============================================================================
gdf_regions = gpd.read_file(G_RGN)
assert len(gdf_regions) == 90, len(gdf_regions)


# ============================================================================
# Resolve example-basin geometry (focus mask, outlet, polygon, bbox)
# ============================================================================
with rasterio.open(E_BSN) as ds:
    bsn_arr = ds.read(1)
    bsn_xform = ds.transform
FOCUS_MASK = (bsn_arr == FOCUS_BSN_ID)

REC = json.load(open(E_REC))
r_out, c_out = REC["outlet_row_in_window"], REC["outlet_col_in_window"]
OUTLET_LON, OUTLET_LAT = rasterio.transform.xy(bsn_xform, r_out, c_out)

_polys = [shp_shape(g) for g, v in raster_shapes(FOCUS_MASK.astype(np.uint8), mask=FOCUS_MASK,
                                                    transform=bsn_xform) if int(v) == 1]
gdf_basin = gpd.GeoDataFrame(geometry=[unary_union(_polys)], crs="EPSG:4326")

basin_bounds = gdf_basin.total_bounds
BLON0, BLAT0, BLON1, BLAT1 = basin_bounds

# Red locator rectangle on globe — bigger pad so it reads at 360°×145°
PAD_DEG = 4.0
RBOX_LON0 = BLON0 - PAD_DEG; RBOX_LON1 = BLON1 + PAD_DEG
RBOX_LAT0 = BLAT0 - PAD_DEG; RBOX_LAT1 = BLAT1 + PAD_DEG

# Basin zoom panels b/d: subtract the outlet's dfs_in, the basin's first
# number (the outlet is entered first): (b) runs 0 .. n - 1, (d) from the
# first source pixel's number up to n - 1; the outlet's interval is [0, n - 1].
with rasterio.open(E_DFSIN) as ds:
    _din = ds.read(1)
    DFSIN_OFFSET = int(_din[r_out, c_out])
    assert DFSIN_OFFSET == int(_din[FOCUS_MASK].min()) == REC["dfs_in_first"]
DFSOUT_OFFSET = DFSIN_OFFSET


# ============================================================================
# Helpers
# ============================================================================
def human_fmt(v, _):
    """Compact tick formatter: 1.2M, 350k, 480, etc."""
    av = abs(v)
    if av >= 1e9:
        x = v / 1e9
        return f"{x:.1f}B"
    if av >= 1e6:
        x = v / 1e6
        return f"{x:.0f}M" if x == int(x) else f"{x:.1f}M"
    if av >= 1e3:
        x = v / 1e3
        return f"{x:.0f}k" if x == int(x) else f"{x:.1f}k"
    return f"{v:.0f}"


def plot_global(fig, ax, tif_path, cbar_label):
    with rasterio.open(tif_path) as ds:
        # Subsample × 2 for plotting speed (already 5 km; 10 km display is fine)
        data = ds.read(1, out_shape=(ds.height // 2, ds.width // 2),
                       resampling=Resampling.nearest).astype(np.float32)
        nd = ds.nodata
        valid = np.isfinite(data)
        if nd is not None:
            valid &= (data != nd)
        data = np.ma.masked_where(~valid, data)
        ext = (ds.bounds.left, ds.bounds.right,
               ds.bounds.bottom, ds.bounds.top)
        vmin = float(data.min()) if valid.any() else 0.0
        vmax = float(data.max()) if valid.any() else 1.0

    im = ax.imshow(data, extent=ext, origin="upper", cmap="viridis",
                   vmin=vmin, vmax=vmax, interpolation="nearest")

    # Region boundaries: thin white lines so the per-region DFS index
    # segmentation is explicit on top of the viridis raster.
    gdf_regions.boundary.plot(ax=ax, color="white", linewidth=0.3,
                              alpha=0.8, zorder=4)

    div = make_axes_locatable(ax)
    cax = div.append_axes("right", size="2.2%", pad=0.08)
    cb  = fig.colorbar(im, cax=cax)
    cb.ax.yaxis.set_major_formatter(mticker.FuncFormatter(human_fmt))
    cb.ax.tick_params(labelsize=7, length=2)
    cb.outline.set_linewidth(0.4)

    rect = plt.Rectangle((RBOX_LON0, RBOX_LAT0),
                         RBOX_LON1 - RBOX_LON0, RBOX_LAT1 - RBOX_LAT0,
                         lw=1.8, edgecolor="red", facecolor="none", zorder=5)
    ax.add_patch(rect)

    ax.set_xlim(-180, 180); ax.set_ylim(-60, 85)
    ax.set_xticks([-180, -90, 0, 90, 180])
    ax.set_yticks([-60, -30, 0, 30, 60])
    ax.set_xlabel("Longitude [°]")
    ax.set_ylabel("Latitude [°]")
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)


def plot_basin(fig, ax, tif_path, cbar_label, offset=0):
    with rasterio.open(tif_path) as ds:
        data = ds.read(1).astype(np.float64)
        nd = ds.nodata
        valid = np.isfinite(data) & FOCUS_MASK
        if nd is not None:
            valid &= (data != nd)
        data = data - offset                     # local-relative values
        data = np.ma.masked_where(~valid, data)
        ext = (ds.bounds.left, ds.bounds.right,
               ds.bounds.bottom, ds.bounds.top)
        vmin = float(data.min()) if valid.any() else 0.0
        vmax = float(data.max()) if valid.any() else 1.0

    im = ax.imshow(data, extent=ext, origin="upper", cmap="viridis",
                   vmin=vmin, vmax=vmax, interpolation="nearest")
    gdf_basin.boundary.plot(ax=ax, color="black", linewidth=0.8, zorder=5)

    ax.plot(OUTLET_LON, OUTLET_LAT, marker="v", markersize=7,
            markerfacecolor="#E53935", markeredgecolor="black",
            markeredgewidth=0.7, linestyle="", zorder=10)

    div = make_axes_locatable(ax)
    cax = div.append_axes("right", size="3%", pad=0.08)
    cb  = fig.colorbar(im, cax=cax)
    cb.ax.yaxis.set_major_formatter(mticker.FuncFormatter(human_fmt))
    cb.ax.tick_params(labelsize=7, length=2)
    cb.outline.set_linewidth(0.4)
    cb.set_label(cbar_label, fontsize=7.5, labelpad=3)

    xmin, ymin, xmax, ymax = basin_bounds
    px = (xmax - xmin) * 0.04; py = (ymax - ymin) * 0.04
    ax.set_xlim(xmin - px, xmax + px)
    ax.set_ylim(ymin - py, ymax + py)
    ax.set_xticks([]); ax.set_yticks([])
    for s in ax.spines.values():
        s.set_visible(False)


# ============================================================================
# Main
# ============================================================================
def main():
    fig = plt.figure(figsize=(7.4, 4.4))
    gs  = GridSpec(2, 2, figure=fig,
                   width_ratios=[1.65, 1], hspace=0.32, wspace=0.16)

    panels = [
        (G_DFSIN,  E_DFSIN,
         r"(a) Global $\mathit{dfs\_in}$",
         r"(b) Example basin, $\mathit{dfs\_in}$",
         r"$\mathit{dfs\_in}$",
         r"$\mathit{dfs\_in} - \mathit{dfs\_in}_{\mathrm{outlet}}$",
         DFSIN_OFFSET),
        (G_DFSOUT, E_DFSOUT,
         r"(c) Global $\mathit{dfs\_out}$",
         r"(d) Example basin, $\mathit{dfs\_out}$",
         r"$\mathit{dfs\_out}$",
         r"$\mathit{dfs\_out} - \mathit{dfs\_in}_{\mathrm{outlet}}$",
         DFSOUT_OFFSET),
    ]

    for i, (g_tif, e_tif, t_l, t_r, cb_g, cb_b, off) in enumerate(panels):
        ax_l = fig.add_subplot(gs[i, 0])
        ax_r = fig.add_subplot(gs[i, 1])
        plot_global(fig, ax_l, g_tif, cb_g)
        ax_l.set_title(t_l, loc="left", pad=4)
        plot_basin(fig, ax_r, e_tif, cb_b, offset=off)
        ax_r.set_title(t_r, loc="left", pad=4)

    fig.subplots_adjust(left=0.06, right=0.96, top=0.94, bottom=0.07)

    for ext in ("png", "pdf"):
        out = os.path.join(OUT_DIR, f"Fig06_DFS_Interval_Global_v11.{ext}")
        if os.path.exists(out):
            raise SystemExit(f"{out} is there already")
        fig.savefig(out, dpi=300, bbox_inches="tight", facecolor="white")
        print(f"Saved: {out}")
    print(f"  Outlet: lon={OUTLET_LON:.4f}°, lat={OUTLET_LAT:.4f}°")
    print(f"  Basin bounds (lon0,lat0,lon1,lat1) = "
          f"({BLON0:.3f}, {BLAT0:.3f}, {BLON1:.3f}, {BLAT1:.3f})")
    print(f"  Focus pixels (global_basin_id={FOCUS_BSN_ID}): {int(FOCUS_MASK.sum())}")
    plt.close(fig)


if __name__ == "__main__":
    main()