#!/usr/bin/env python3
"""
Regrid a NorESM aerosol-PPE OAT campaign from the unstructured spectral-element
(ne16pg3 physgrid, `ncol`) history output onto a regular lat/lon grid, and cache
the time-mean field of every mapped AeroCom variable for each ensemble member.

This is the data-prep step behind the OAT web interface (build_oat_site.py):
rendering ADF-style maps needs lat/lon fields, while the raw h0a output is on
the SE grid. Regridding is done once here and cached, so the site build (and any
later re-styling of the plots) is fast and needs no access to the raw campaign.

The AeroCom variable definitions are imported directly from aerocom_oat_checker
rather than duplicated, so the maps and the global-mean CSV always agree on what
each variable means. Variables whose NorESM source fields are absent from a
campaign's history stream (e.g. no SWCF/LWCF/FSNT/PRECT in oatv5) are skipped
with a warning instead of aborting.

Derived quantities (ssa550, angstrm550_865, ...) are evaluated per timestep on
the native SE grid and only then time-averaged and regridded, matching the order
of operations in aerocom_oat_checker.compute_member_means -- so the area-weighted
global mean of a cached map reproduces that variable's value in the checker CSV.

Run with the ml-notebook conda environment, e.g.:
    /nird/scratch/ovewh/ml-notebook/bin/python oat_web/prepare_oat_data.py \\
        --ppe-dir aerosol_ppe/oatv5_aerosolPPE_20260818
"""

import argparse
import sys
from pathlib import Path

import numpy as np
import xarray as xr
from scipy import sparse

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from aerocom_oat_checker import (  # noqa: E402
    DEFAULT_HIST_GLOBS,
    VARIABLE_MAP,
    build_member_table,
)

DEFAULT_WEIGHT_FILE = "/nird/datalake/NS2345K/ovewh/map_ne16pg3_to_1.9x2.5_nomask_scripgrids_c250425.nc"

# netCDF variable names can't safely carry the '.' in AeroCom names like
# 'ccns.3'; the original name is kept in each variable's `aerocom_name` attr.
def _storage_name(var):
    return var.replace(".", "_")


class Regridder:
    """Sparse ESMF/SCRIP weight-file regridder, SE `ncol` -> (lat, lon)."""

    def __init__(self, weight_file):
        with xr.open_dataset(weight_file) as w:
            n_a = w.sizes["n_a"]
            n_b = w.sizes["n_b"]
            # SCRIP row/col indices are 1-based (Fortran).
            self.matrix = sparse.csr_matrix(
                (w["S"].values, (w["row"].values - 1, w["col"].values - 1)),
                shape=(n_b, n_a),
            )
            n_lon, n_lat = (int(x) for x in w["dst_grid_dims"].values)
            self.shape = (n_lat, n_lon)
            self.lat = w["yc_b"].values.reshape(self.shape)[:, 0]
            self.lon = w["xc_b"].values.reshape(self.shape)[0, :]
            # A destination cell with frac_b == 0 receives no source data at all;
            # leaving it as a hard 0 would paint a fake value onto the map.
            self.frac_b = w["frac_b"].values
        self.n_src = n_a

    def __call__(self, values):
        """Regrid a 1-D array over `ncol` to a 2-D (lat, lon) array."""
        out = self.matrix.dot(np.nan_to_num(values, nan=0.0))
        out = np.where(self.frac_b > 1e-8, out, np.nan)
        return out.reshape(self.shape)


def member_time_mean_fields(member_dir: Path, hist_globs, regridder, variables):
    """Return {aerocom_var: (lat, lon) array} of time-mean fields for one member."""
    files = []
    for pattern in hist_globs:
        matched = sorted(member_dir.glob(pattern))
        if matched:
            files = matched
            break
    if not files:
        print(f"  warning: no history files matching {hist_globs} in {member_dir}", file=sys.stderr)
        return None

    fields = {}
    with xr.open_mfdataset(files, combine="by_coords") as ds:
        for var in variables:
            spec = VARIABLE_MAP[var]
            if spec["status"] == "unmapped":
                continue
            try:
                da = spec["compute"](ds) if "compute" in spec else ds[spec["noresm"]]
                # Derived quantities can divide by ~0 (e.g. ssa550 where there is
                # no aerosol); those become inf, which would otherwise poison the
                # time mean and saturate the colour scale.
                values = da.mean(dim="time", skipna=True).values
            except KeyError as e:
                print(f"  warning: {e} not found in {member_dir.name}, skipping variable", file=sys.stderr)
                continue
            values = np.where(np.isfinite(values), values, np.nan)
            fields[var] = regridder(values)
    return fields


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--ppe-dir", type=Path, required=True)
    parser.add_argument("--report-json", type=Path, default=None,
                        help="default: paramfile/<ppe-dir name>/check_member_report.json")
    parser.add_argument("--weight-file", default=DEFAULT_WEIGHT_FILE)
    parser.add_argument("--hist-glob", nargs="+", default=DEFAULT_HIST_GLOBS)
    parser.add_argument("--variables", nargs="+", default=list(VARIABLE_MAP.keys()))
    parser.add_argument("-o", "--output", type=Path, default=None,
                        help="default: oat_web/data/<ppe-dir name>_maps.nc")
    args = parser.parse_args()

    ppe_name = args.ppe_dir.name
    report_json = args.report_json or Path("paramfile") / ppe_name / "check_member_report.json"
    output = args.output or Path(__file__).resolve().parent / "data" / f"{ppe_name}_maps.nc"
    output.parent.mkdir(parents=True, exist_ok=True)

    regridder = Regridder(args.weight_file)
    print(f"regridder: {regridder.n_src} ncol -> {regridder.shape[0]} lat x {regridder.shape[1]} lon")

    member_table = build_member_table(report_json)

    labels, member_ids, per_member = [], [], []
    for member_id, label in sorted(member_table.items()):
        member_dir = args.ppe_dir / f"ensemble_member.{member_id:03d}"
        if not member_dir.is_dir():
            print(f"  warning: {member_dir} not found, skipping", file=sys.stderr)
            continue
        print(f"regridding member {member_id:03d} ({label})...", flush=True)
        fields = member_time_mean_fields(member_dir, args.hist_glob, regridder, args.variables)
        if fields is None:
            continue
        labels.append(label)
        member_ids.append(member_id)
        per_member.append(fields)

    # Only keep variables that every retained member actually produced, so the
    # site build can assume a complete (member x lat x lon) block per variable.
    available = [v for v in args.variables if all(v in f for f in per_member)]
    dropped = [v for v in args.variables
               if v not in available and VARIABLE_MAP[v]["status"] != "unmapped"]
    if dropped:
        print(f"\nnote: {len(dropped)} variable(s) unavailable for this campaign: {', '.join(dropped)}")

    data_vars = {}
    for var in available:
        stack = np.stack([f[var] for f in per_member])
        spec = VARIABLE_MAP[var]
        attrs = {"aerocom_name": var, "status": spec["status"]}
        if "note" in spec:
            attrs["note"] = spec["note"]
        attrs["source"] = spec.get("noresm") or getattr(spec.get("compute"), "__name__", "?")
        data_vars[_storage_name(var)] = (("experiment", "lat", "lon"), stack, attrs)

    ds = xr.Dataset(
        data_vars,
        coords={
            "experiment": ("experiment", labels),
            "member_id": ("experiment", member_ids),
            "lat": ("lat", regridder.lat),
            "lon": ("lon", regridder.lon),
        },
        attrs={
            "ppe": ppe_name,
            "description": "time-mean AeroCom diagnostics regridded from ne16pg3 SE to regular lat/lon",
            "weight_file": args.weight_file,
        },
    )
    ds.to_netcdf(output)
    print(f"\nwrote {output} ({len(labels)} experiments x {len(available)} variables)")


if __name__ == "__main__":
    main()