"""
Subset CONUS404 by user-defined lat/lon box and variable list,
looping one day at a time over a range of day-of-years.
"""

from datetime import datetime, timedelta
from pathlib import Path

import numpy as np
import xarray as xr
import planetary_computer
import pystac_client

# -----------------------------------------------------------------------------
# USER INPUTS
# -----------------------------------------------------------------------------
VARIABLES  = ["T2", "U10", "V10", "Q2", "PSFC", "GLW"]

LAT_MIN, LAT_MAX = 24.0, 51.0
LON_MIN, LON_MAX = -126.0, -66.0

YEAR       = 2020
DOY_START  = 1        # e.g. June 1
DOY_END    = 30        # inclusive

OUT_DIR    = Path("./conus404_out")
ASSET_KEY  = "zarr-abfs"
# -----------------------------------------------------------------------------

OUT_DIR.mkdir(parents=True, exist_ok=True)

# 1. Open the STAC asset once
catalog = pystac_client.Client.open(
    "https://planetarycomputer.microsoft.com/api/stac/v1",
    modifier=planetary_computer.sign_inplace,
)
asset = catalog.get_collection("conus404").assets[ASSET_KEY]

ds = xr.open_zarr(
    asset.href,
    storage_options=asset.extra_fields["xarray:storage_options"],
    **asset.extra_fields["xarray:open_kwargs"],
)

try:
    # 2. Validate variables
    missing = [v for v in VARIABLES if v not in ds.data_vars]
    if missing:
        raise KeyError(f"Variables not found in dataset: {missing}")
    ds = ds[VARIABLES]

    # 3. Compute spatial index bounds ONCE (grid is static across time)
    lat2d = ds["lat"]
    lon2d = ds["lon"]
    mask = (
        (lat2d >= LAT_MIN) & (lat2d <= LAT_MAX) &
        (lon2d >= LON_MIN) & (lon2d <= LON_MAX)
    )
    if not bool(mask.any()):
        raise ValueError("Lat/lon box does not intersect the CONUS404 grid.")

    sn_idx = np.where(mask.any(dim="west_east").values)[0]
    we_idx = np.where(mask.any(dim="south_north").values)[0]
    sn_slice = slice(int(sn_idx.min()), int(sn_idx.max()) + 1)
    we_slice = slice(int(we_idx.min()), int(we_idx.max()) + 1)

    ds_space = ds.isel(south_north=sn_slice, west_east=we_slice)

    # 4. Loop one day at a time
    for doy in range(DOY_START, DOY_END + 1):
        day_start = datetime(YEAR, 1, 1) + timedelta(days=doy - 1)
        day_end   = day_start + timedelta(days=1) - timedelta(seconds=1)

        t0 = day_start.strftime("%Y-%m-%dT%H:%M:%S")
        t1 = day_end.strftime("%Y-%m-%dT%H:%M:%S")

        out_file = OUT_DIR / f"conus404_{YEAR}_doy{doy:03d}.nc"
        if out_file.exists():
            print(f"[skip] {out_file.name} already exists")
            continue

        print(f"[{YEAR} DOY {doy:03d}] {t0}  ->  {t1}")

        ds_day = ds_space.sel(time=slice(t0, t1))
        if ds_day.sizes["time"] == 0:
            print(f"  no timesteps found, skipping")
            continue

        encoding = {
            v: {"zlib": True, "complevel": 4, "shuffle": True}
            for v in ds_day.data_vars
        }
        ds_day.to_netcdf(out_file, encoding=encoding)
        print(f"  wrote {out_file}  (nt={ds_day.sizes['time']})")

finally:
    ds.close()

print("Done.")