#!/usr/bin/env python3
"""Plot July FATES aboveground carbon density for all available years.

Dependencies: numpy, xarray, netCDF4, matplotlib, shapely>=2, pyshp.
Run: python plot_fates_basin_trends.py
Use --start-year 1985 to restrict the analysis to 1985 onwards. Changes
and smoothing use only the selected years, without pre-period data.

Reads exactly one h0a (monthly mean) NetCDF file per year, once for both
basins. Only FATES_VEGC_ABOVEGROUND and grid metadata are loaded.
These are July samples, NOT annual means or fitted linear trends.
Also plots year-to-year changes: value(year) - value(year - 1), labelled
by the later year. The CSV includes these differences (blank in year one).
An additional plot shows centred five-year rolling means of those changes,
labelled by the middle year. Only complete windows are used (no padding);
the CSV leaves smoothed values blank where a complete window is unavailable.

Watersheds: HydroBASINS v1c, level 3, Amazon 6030007000 / Congo 1030020040.
Downloads the public continental shapefile ZIPs once and caches them.
Source and license: https://www.hydrosheds.org/products/hydrobasins
Reference: Lehner & Grill (2013), https://doi.org/10.1002/hyp.9740

Grid-cell centres determine basin membership; boundary-cell fractional
overlap is not computed. Means are weighted by grid area * land fraction.
Stock units are kg C m-2 of land, not total basin carbon. Changes are in
Mg C ha-1 yr-1: consecutive July samples are one year apart, and kg/m2
is multiplied by 10,000 m2/ha / 1,000 kg/Mg = 10.
"""

import argparse
import calendar
import csv
import io
from pathlib import Path
import re
import urllib.request
import warnings
import zipfile

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import shapefile
import shapely
from shapely.geometry import shape
import xarray as xr


DATA_DIR = Path(
    "/datalake/NS9560K/kjetisaa/"
    "ihist.f09_g17.ctsm5.4.042_noresm_v3_Trendy26.T26_S2_pt2.2026-08-22/"
    "lnd/hist"
)
RUN_LABEL = "TRENDY26 S2"
VARIABLE = "FATES_VEGC_ABOVEGROUND"
BASINS = {
    "Amazon": ("sa", 6030007000, "#CC3311"),
    "Congo": ("af", 1030020040, "#4477AA"),
}


def load_basin(cache_dir: Path, continent: str, basin_id: int):
    """Read a named full watershed, not its enclosing level-2 region."""
    stem = f"hybas_{continent}_lev03_v1c"
    archive = cache_dir / f"{stem}.zip"
    if not archive.exists():
        url = f"https://data.hydrosheds.org/file/hydrobasins/standard/{stem}.zip"
        print(f"Downloading {url}", flush=True)
        temporary = archive.with_suffix(".zip.part")
        try:
            with urllib.request.urlopen(url, timeout=120) as response:
                with temporary.open("wb") as target:
                    while chunk := response.read(1024 * 1024):
                        target.write(chunk)
            with zipfile.ZipFile(temporary) as check:
                bad_member = check.testzip()
                if bad_member is not None:
                    raise ValueError(f"Corrupt ZIP member: {bad_member}")
            temporary.replace(archive)
        finally:
            temporary.unlink(missing_ok=True)

    with zipfile.ZipFile(archive) as zipped:
        with shapefile.Reader(
            shp=io.BytesIO(zipped.read(f"{stem}.shp")),
            shx=io.BytesIO(zipped.read(f"{stem}.shx")),
            dbf=io.BytesIO(zipped.read(f"{stem}.dbf")),
        ) as reader:
            for record in reader.iterShapeRecords():
                if record.record["HYBAS_ID"] == basin_id:
                    geometry = shape(record.shape.__geo_interface__)
                    if not geometry.is_valid:
                        warnings.warn(
                            f"Repairing invalid HydroBASINS geometry {basin_id}: "
                            f"{shapely.is_valid_reason(geometry)}",
                            stacklevel=2,
                        )
                        geometry = shapely.make_valid(geometry)
                    if geometry.is_empty or not geometry.is_valid:
                        raise ValueError(f"Invalid basin geometry: {basin_id}")
                    return geometry
    raise ValueError(f"HYBAS_ID {basin_id} not found in {archive}")


def find_files(data_dir: Path, month: int) -> dict[int, Path]:
    """Check all years before opening any model files."""
    candidates = {}
    available_years = set()
    for path in data_dir.glob("*.clm2.h0a.*.nc"):
        match = re.search(r"\.(\d{4})-(\d{2})\.nc$", path.name)
        if match:
            year, file_month = map(int, match.groups())
            available_years.add(year)
            if file_month == month:
                candidates.setdefault(year, []).append(path)
    if not available_years:
        raise FileNotFoundError(f"No monthly h0a files found in {data_dir}")
    years = range(min(available_years), max(available_years) + 1)
    if len(years) < 2:
        raise ValueError("At least two consecutive years are needed for differences")
    selected = []
    for year in years:
        matches = candidates.get(year, [])
        if len(matches) != 1:
            raise FileNotFoundError(
                f"Expected one h0a file for {year}-{month:02d}, "
                f"found {len(matches)} in {data_dir}"
            )
        selected.append(matches[0])
    return dict(zip(years, selected))


def read_series(files: dict[int, Path], basins: dict) -> dict[str, list[float]]:
    series = {name: [] for name in basins}
    weights = {}
    reference_lat = reference_lon = None
    reference_valid = None
    for year, path in files.items():
        with xr.open_dataset(
            path, engine="netcdf4", decode_times=False, cache=False
        ) as dataset:
            variable = dataset[VARIABLE]
            if set(variable.dims) != {"time", "lat", "lon"}:
                raise ValueError(f"Unexpected dimensions in {path}: {variable.dims}")
            if variable.sizes["time"] != 1:
                raise ValueError(f"Expected one monthly record in {path}")
            if variable.attrs.get("units") != "kg m-2":
                raise ValueError(f"Unexpected carbon units in {path}: {variable.attrs}")
            lat = dataset["lat"].values
            lon = dataset["lon"].values
            values = variable.isel(time=0).transpose("lat", "lon").values
            if reference_lat is None:
                reference_lat, reference_lon = lat.copy(), lon.copy()
                area = dataset["area"].transpose("lat", "lon").values
                landfrac = dataset["landfrac"].transpose("lat", "lon").values
                lon2d, lat2d = np.meshgrid((lon + 180) % 360 - 180, lat)
                points = shapely.points(lon2d, lat2d)
                land_area = area * landfrac
                land = (
                    np.isfinite(land_area) & (area > 0)
                    & (landfrac > 0) & (landfrac <= 1)
                )
                for name, geometry in basins.items():
                    mask = shapely.covers(geometry, points) & land
                    if not np.any(mask):
                        raise ValueError(f"No land grid-cell centres inside {name}")
                    weights[name] = np.where(mask, land_area, 0.0)
                    print(f"{name}: {np.count_nonzero(mask)} land cells", flush=True)
                reference_valid = np.isfinite(values)
            else:
                if not (
                    np.array_equal(lat, reference_lat)
                    and np.array_equal(lon, reference_lon)
                ):
                    raise ValueError(f"Grid changed in {path}")
            for name, weight in weights.items():
                mask = weight > 0
                valid = mask & np.isfinite(values)
                if not np.any(valid):
                    raise ValueError(f"No valid {VARIABLE} values for {name}, {year}")
                if not np.array_equal(
                    np.isfinite(values)[mask], reference_valid[mask]
                ):
                    raise ValueError(
                        f"Missing-data footprint changed for {name} in {year}; "
                        "refusing to compare different spatial coverage."
                    )
                if year == min(files) and np.any(mask & ~np.isfinite(values)):
                    warnings.warn(
                        f"{name}: excluding {np.count_nonzero(mask & ~np.isfinite(values))} "
                        "land cells with missing carbon data; coverage must stay constant.",
                        stacklevel=2,
                    )
                series[name].append(
                    float(np.average(values[valid], weights=weight[valid]))
                )
        print(
            f"{year}: " + ", ".join(
                f"{name}={series[name][-1]:.4f}" for name in basins
            ),
            flush=True,
        )
    return series


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path, default=DATA_DIR)
    parser.add_argument("--month", type=int, choices=range(1, 13), default=7)
    parser.add_argument("--start-year", type=int, help="First stock year to include")
    parser.add_argument("--output-dir", type=Path, default=Path.cwd())
    parser.add_argument(
        "--cache-dir", type=Path,
        default=Path(__file__).resolve().parent / "hydrobasins_cache",
    )
    args = parser.parse_args()
    files = find_files(args.data_dir, args.month)
    if args.start_year is not None:
        files = {year: path for year, path in files.items() if year >= args.start_year}
        if len(files) < 2:
            parser.error("--start-year must leave at least two available stock years")
    years = list(files)
    year_range = f"{years[0]}-{years[-1]}"
    print(f"Sampling {len(years)} years: {year_range}", flush=True)
    args.cache_dir.mkdir(parents=True, exist_ok=True)
    basins = {
        name: load_basin(args.cache_dir, continent, basin_id)
        for name, (continent, basin_id, _) in BASINS.items()
    }
    series = read_series(files, basins)
    differences = {
        name: np.diff(values) * 10
        for name, values in series.items()
    }
    smoothed_changes = {
        name: xr.DataArray(values, dims="year").rolling(
            year=5, center=True, min_periods=5
        ).mean().values
        for name, values in differences.items()
    }
    args.output_dir.mkdir(parents=True, exist_ok=True)
    stem = f"fates_aboveground_carbon_{args.month:02d}_{years[0]}_{years[-1]}"
    csv_path = args.output_dir / f"{stem}.csv"
    with csv_path.open("w", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow([
            "year", "month", "Amazon_kgC_m2", "Congo_kgC_m2",
            "Amazon_change_MgC_ha_yr", "Congo_change_MgC_ha_yr",
            "Amazon_change_5yr_mean_MgC_ha_yr", "Congo_change_5yr_mean_MgC_ha_yr",
        ])
        for index, year in enumerate(years):
            writer.writerow(
                [
                    year, args.month, series["Amazon"][index], series["Congo"][index],
                    differences["Amazon"][index - 1] if index else "",
                    differences["Congo"][index - 1] if index else "",
                    *[
                        smoothed_changes[name][index - 1]
                        if index and np.isfinite(smoothed_changes[name][index - 1])
                        else ""
                        for name in BASINS
                    ],
                ]
            )

    for plot_kind in ("stock", "change", "smoothed_change"):
        is_change = plot_kind != "stock"
        is_smoothed = plot_kind == "smoothed_change"
        if is_smoothed and len(years) < 6:
            warnings.warn(
                "Skipping five-year-smoothed change plot: at least six stock years are required.",
                stacklevel=2,
            )
            continue
        fig, ax = plt.subplots(figsize=(10, 5.5))
        plot_years = years[1:] if is_change else years
        plot_values = differences if is_change else series
        if is_smoothed:
            plot_years = years[3:-2]
            plot_values = {name: values[2:-2] for name, values in smoothed_changes.items()}
        for name, (_, _, colour) in BASINS.items():
            ax.plot(
                plot_years, plot_values[name], label=name, color=colour,
                linewidth=2, marker="o", markersize=3,
            )
        if is_change:
            ax.axhline(0, color="black", linewidth=0.8, linestyle="--")
        ax.set(
            xlabel=(
                "Year (centre of five-year window)" if is_smoothed else
                "Year (later year of each difference)" if is_change else "Year"
            ),
            ylabel=(
                r"Year-to-year carbon change (Mg C ha$^{-1}$ yr$^{-1}$)"
                if is_change else
                r"Aboveground vegetation carbon (kg C m$^{-2}$ land)"
            ),
            title=(
                "FATES aboveground vegetation carbon: Amazon and Congo\n"
                f"{RUN_LABEL}; {calendar.month_name[args.month]} monthly means, {year_range}"
                + (
                    "; centred five-year mean of annual changes" if is_smoothed else
                    "; year minus previous year" if is_change else ""
                )
            ),
            xlim=(plot_years[0], plot_years[-1]) if len(plot_years) > 1
            else (plot_years[0] - 0.5, plot_years[0] + 0.5),
        )
        ax.grid(alpha=0.25)
        ax.legend(frameon=False)
        fig.text(
            0.5, 0.01,
            "Land-area-weighted means; HydroBASINS v1c level-3 watersheds "
            "(Lehner & Grill, 2013); grid-cell-centre masks",
            ha="center", fontsize=8,
        )
        fig.tight_layout(rect=(0, 0.035, 1, 1))
        suffix = "_year_to_year_change" if is_change else ""
        if is_smoothed:
            suffix += "_5yr_mean"
        png_path = args.output_dir / f"{stem}{suffix}.png"
        fig.savefig(png_path, dpi=200)
        plt.close(fig)
        print(f"Saved {png_path}")
    print(f"Saved {csv_path}")


if __name__ == "__main__":
    main()