"""
Diagnostics for the CAM aquaplanet CO2 sweep.

Area weighting on the CAM finite-volume grid uses cos(lat) with the standard
FV correction at the polar cells (which are half-width in latitude). We build
weights from the latitude cell edges so the global mean is exact rather than a
naive cos(lat) sum.
"""

import numpy as np
import xarray as xr

# ---------------------------------------------------------------------------
# area weighting
# ---------------------------------------------------------------------------


def gw(ds):
    """Latitude area weights. Prefer CAM's own Gaussian/FV weights (`gw`)."""
    if "gw" in ds:
        return ds["gw"]
    lat = ds["lat"]
    edges = np.empty(lat.size + 1)
    edges[1:-1] = 0.5 * (lat.values[:-1] + lat.values[1:])
    edges[0], edges[-1] = lat.values[0], lat.values[-1]
    w = np.sin(np.deg2rad(edges[1:])) - np.sin(np.deg2rad(edges[:-1]))
    return xr.DataArray(np.abs(w), dims=["lat"], coords={"lat": lat})


def gmean(da, ds=None):
    """Area-weighted global mean over lat/lon (and time if present)."""
    ds = ds if ds is not None else da
    w = gw(ds)
    dims = [d for d in ("time", "lat", "lon") if d in da.dims]
    return da.weighted(w.fillna(0.0)).mean(dim=[d for d in dims if d != "time"]).mean(
        dim="time" if "time" in da.dims else []
    )


def zmean(da):
    """Zonal + time mean."""
    out = da.mean(dim="lon") if "lon" in da.dims else da
    return out.mean(dim="time") if "time" in out.dims else out


# ---------------------------------------------------------------------------
# radiation budget
# ---------------------------------------------------------------------------
# CAM sign conventions:
#   FSNT  net SW flux at top of model, downward positive
#   FLNT  net LW flux at top of model, UPWARD positive
#   so net downward TOA radiation = FSNT - FLNT
#   *C variants are clear-sky; CRE = allsky - clearsky


def toa(ds):
    """TOA radiation diagnostics, global means [W m-2]."""
    d = {}
    d["SOLIN"] = float(gmean(ds["SOLIN"], ds)) if "SOLIN" in ds else np.nan
    d["FSNT"] = float(gmean(ds["FSNT"], ds))
    d["FLNT"] = float(gmean(ds["FLNT"], ds))
    d["NET_TOA"] = d["FSNT"] - d["FLNT"]
    if "FSNTC" in ds and "FLNTC" in ds:
        d["FSNTC"] = float(gmean(ds["FSNTC"], ds))
        d["FLNTC"] = float(gmean(ds["FLNTC"], ds))
        d["NET_TOA_CLR"] = d["FSNTC"] - d["FLNTC"]
        d["SWCRE"] = d["FSNT"] - d["FSNTC"]
        d["LWCRE"] = d["FLNTC"] - d["FLNT"]
        d["CRE"] = d["SWCRE"] + d["LWCRE"]
    for k in ("FSNS", "FLNS", "SHFLX", "LHFLX", "TS", "TREFHT", "TMQ",
              "PRECT", "CLDTOT", "CLDLOW", "CLDHGH", "LWCF", "SWCF"):
        if k in ds:
            d[k] = float(gmean(ds[k], ds))
    if "PRECT" in d:
        d["PRECT_mm_day"] = d["PRECT"] * 86400.0 * 1000.0
    return d


def surface_net(ds):
    """Net downward surface energy flux [W m-2]: SW - LW - SH - LH."""
    v = {k: gmean(ds[k], ds) for k in ("FSNS", "FLNS", "SHFLX", "LHFLX") if k in ds}
    if len(v) < 4:
        return np.nan
    return float(v["FSNS"] - v["FLNS"] - v["SHFLX"] - v["LHFLX"])


# ---------------------------------------------------------------------------
# vertical structure
# ---------------------------------------------------------------------------


def pressure(ds):
    """Mid-level pressure [hPa] from the hybrid coordinate, global mean PS."""
    ps = float(gmean(ds["PS"], ds))
    return (ds["hyam"] * ds["P0"] + ds["hybm"] * ps) / 100.0


def profile(ds, var):
    """Global-mean vertical profile of `var` against pressure [hPa]."""
    p = pressure(ds)
    prof = gmean(ds[var], ds)
    return np.asarray(p), np.asarray(prof)


# ---------------------------------------------------------------------------
# forcing
# ---------------------------------------------------------------------------


def forcing_per_doubling(co2_ppm, net_toa):
    """Least-squares slope of net TOA flux against log2(CO2) [W m-2 / doubling].

    Under fixed SST the change in net downward TOA flux IS the radiative
    forcing (Hansen et al. 2005 'fixed-SST' / effective forcing), because the
    surface temperature is prevented from responding.
    """
    x = np.log2(np.asarray(co2_ppm, float) / co2_ppm[0])
    y = np.asarray(net_toa, float)
    A = np.vstack([x, np.ones_like(x)]).T
    slope, intercept = np.linalg.lstsq(A, y, rcond=None)[0]
    resid = y - (slope * x + intercept)
    return float(slope), float(intercept), resid