"""NaN-in-mask check for a pair of initial-condition files.

usage: python check_ic.py ATM_IC.nc OCEAN_IC.nc [OCEAN_FORCING.nc]
   or: from check_ic import check; check(atm_ic, ocean_ic, ocean_forcing)  # True on PASS

The masks come from an ocean-forcing file of this package (default:
../forcing_data/ocean-forcing-<first year found>.nc). They are the masks the checkpoints
apply. Rules: a 3-D ocean variable `<name>_<k>` uses mask_<k>; every other ocean variable
uses mask_2d. Passing means:
  - no non-finite value inside the wet mask, and no finite value outside it (ocean);
  - no non-finite value anywhere (atmosphere);
  - sst NaN fraction == 1 - mean(mask_2d) (0.2958 for this grid).
Exit code 0 on pass, 1 on failure.
"""

import glob
import os
import sys

import numpy as np
import xarray as xr


def check(atm_path, ocn_path, mask_path=None):
    """Check an IC pair; prints the details and returns True on PASS."""
    if mask_path is None:
        here = os.path.dirname(os.path.abspath(__file__))
        found = sorted(glob.glob(os.path.join(here, "..", "forcing_data", "ocean-forcing-*.nc")))
        if not found:
            raise FileNotFoundError("no ocean-forcing-*.nc found; pass the mask file")
        mask_path = found[0]
    masks = xr.open_dataset(mask_path)
    atm = xr.open_dataset(atm_path, decode_times=False)
    ocn = xr.open_dataset(ocn_path, decode_times=False)
    ok = True
    print(f"masks from {mask_path}")
    print(f"times: atmosphere {atm.time.values}, ocean {ocn.time.values} ({ocn.time.attrs.get('units')})")
    if not np.array_equal(atm.time.values, ocn.time.values):
        print("FAIL: atmosphere and ocean IC times differ")
        ok = False

    def mask_for(name):
        suffix = name.rsplit("_", 1)[-1]
        if suffix.isdigit() and f"mask_{suffix}" in masks:
            return masks[f"mask_{suffix}"].values > 0
        return masks["mask_2d"].values > 0

    bad_in, wet_out = {}, {}
    for name in ocn.data_vars:
        if "lat" not in ocn[name].dims:
            continue
        m = mask_for(name)
        for i in range(ocn.sizes["time"]):
            v = ocn[name].isel(time=i).values
            b = int((m & ~np.isfinite(v)).sum())
            o = int((~m & np.isfinite(v)).sum())
            if b:
                bad_in[f"{name}[t={i}]"] = b
            if o:
                wet_out[f"{name}[t={i}]"] = o
    atm_bad = {
        n: int((~np.isfinite(atm[n].values)).sum())
        for n in atm.data_vars
        if "lat" in atm[n].dims and (~np.isfinite(atm[n].values)).any()
    }
    print(f"ocean: {len(ocn.data_vars)} variables; non-finite inside wet mask: {bad_in or 'NONE'}")
    print(f"ocean: finite outside wet mask: {wet_out or 'NONE'}")
    print(f"atmosphere: {len(atm.data_vars)} variables; non-finite anywhere: {atm_bad or 'NONE'}")
    expect = 1 - float(masks["mask_2d"].mean())
    for i in range(ocn.sizes["time"]):
        frac = float(np.isnan(ocn["sst"].isel(time=i).values).mean())
        good = abs(frac - expect) < 1e-6
        print(f"sst NaN fraction, IC {i}: {frac:.4f} (expected {expect:.4f}) {'ok' if good else 'FAIL'}")
        ok &= good
    ok &= not bad_in and not wet_out and not atm_bad
    print("PASS" if ok else "FAIL")
    return bool(ok)


def main():
    if len(sys.argv) < 3:
        sys.exit(__doc__)
    sys.exit(0 if check(*sys.argv[1:4]) else 1)


if __name__ == "__main__":
    main()