#!/usr/bin/env python3
"""
TRACER-MIP: Extract Hourly Precipitation from New E3SM-SCREAM Datasets
======================================================================

Processes three E3SM-SCREAM (THREAD) runs:
    E3SM_CNTL              -> /pscratch/sd/y/yunpengs/TRACER-MIP_EAMxx/20220807_ctrl_rainfrac2
    E3SM_CNTL_1.5TKE       -> /pscratch/sd/y/yunpengs/TRACER-MIP_EAMxx/20220807_ctrl_rainfrac2_1.5tke
    E3SM_CNTL_1.5TKE_No_COND -> /pscratch/sd/y/yunpengs/TRACER-MIP_EAMxx/20220807_ctrl_rainfrac2_1.5tke_nocond

Case: 2022-08-07 Houston (48-hour run, Aug 7 00 UTC -> Aug 9 00 UTC)

File layout (verified by Inspect_E3SM_2D.py):
    Stream: 2D.INST.INSTANT, un-prefixed filenames, one file per hour:
        F2010-SCREAMv1.TRACERmipx32subx3v1pg2.001.2D.INST.INSTANT.nmins_x2.YYYY-MM-DD-SSSSS.nc
    Each file contains 30 internal time slices at 2-min cadence.
    Time units: "days since 2022-08-07 00:00:00", noleap calendar.

Source variable: precip_total_surf_mass_flux
    Shape: (time=30, lat=502, lon=583)
    Units: m/s (instantaneous water-equivalent volume flux)

Hourly aggregation (per file = per hour):
    rate_mm_hr = mean(precip_total_surf_mass_flux, over 30 slices) * 1000 * 3600

Timestamp convention:
    The output time coordinate is the END of each hourly window, e.g. a
    file whose internal slices span 00:00-00:58 UTC is reported at
    01:00 UTC. This is consistent with accumulated-precip reporting.

Output: one NetCDF per dataset
    /global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip/E3SM_Yunpeng/
        E3SM_CNTL_Aug7_Precip.nc
        E3SM_CNTL_1.5TKE_Aug7_Precip.nc
        E3SM_CNTL_1.5TKE_No_COND_Aug7_Precip.nc

Output schema:
    Dimensions: time (48), south_north (502), west_east (583)
    Variables:
        time         (time,)                       - hourly timestamps (UTC, end-of-window)
        XLAT         (south_north, west_east)      - 2D latitude
        XLONG        (south_north, west_east)      - 2D longitude
        PRECIP_RATE  (time, south_north, west_east) - hourly precipitation rate [mm hr-1]

Usage:
    python Extract_Precip_E3SM_New.py                       # all datasets
    python Extract_Precip_E3SM_New.py --dataset E3SM_CNTL   # one only
    python Extract_Precip_E3SM_New.py --dryrun              # list files, no write
    python Extract_Precip_E3SM_New.py --overwrite           # rewrite existing outputs
"""

import os
os.environ.setdefault("HDF5_USE_FILE_LOCKING", "FALSE")

import sys
import re
import glob
import time
import argparse
import warnings
import traceback
from datetime import datetime, timedelta

import numpy as np
from netCDF4 import Dataset

warnings.filterwarnings("ignore")


# ============================================================
# CONFIGURATION
# ============================================================
BASE_PATH   = "/pscratch/sd/y/yunpengs/TRACER-MIP_EAMxx"
OUTPUT_DIR  = "/global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip/E3SM_Yunpeng"

DATASETS = {
    "E3SM_CNTL":                "20220807_ctrl_rainfrac2",
    "E3SM_CNTL_1.5TKE":         "20220807_ctrl_rainfrac2_1.5tke",
    "E3SM_CNTL_1.5TKE_No_COND": "20220807_ctrl_rainfrac2_1.5tke_nocond",
}

# Aug 7 case window: extract the full 48 hours (Aug 7 00 UTC -> Aug 9 00 UTC).
# Filename SOD encodes the START of each hourly window; output time is END.
CASE_START = datetime(2022, 8, 7,  0, 0)
CASE_END   = datetime(2022, 8, 9,  0, 0)   # final END-of-window timestamp

# Filename pattern for the 2D INST.INSTANT stream.
FNAME_RE = re.compile(
    r"^F2010-SCREAMv1\.TRACERmipx32subx3v1pg2\.\d+\."
    r"2D\.INST\.INSTANT\.nmins_x\d+\."
    r"(?P<date>\d{4}-\d{2}-\d{2})-(?P<sod>\d{5})\.nc$"
)

SRC_VAR  = "precip_total_surf_mass_flux"
M_TO_MM  = 1000.0
S_TO_HR  = 3600.0
FILL_BIG = 1e30          # fill values in E3SM are ~3.4e+33


# ============================================================
# HELPERS
# ============================================================
def parse_fname(fname):
    m = FNAME_RE.match(fname)
    if not m:
        return None
    sod = int(m.group("sod"))
    start_dt = datetime.strptime(m.group("date"), "%Y-%m-%d") + timedelta(seconds=sod)
    return start_dt


def find_hourly_files(dir_path):
    """Return list of (window_start_dt, window_end_dt, full_path), sorted, in case window."""
    files = []
    for f in sorted(os.listdir(dir_path)):
        if not f.endswith(".nc"):
            continue
        start_dt = parse_fname(f)
        if start_dt is None:
            continue
        end_dt = start_dt + timedelta(hours=1)
        # Keep files whose END-of-window falls within [CASE_START+1h, CASE_END].
        if CASE_START < end_dt <= CASE_END:
            files.append((start_dt, end_dt, os.path.join(dir_path, f)))
    return files


def read_coords(fpath):
    """Read lat (502,) and lon (583,) 1D arrays from a sample file, broadcast to 2D."""
    with Dataset(fpath, "r") as ds:
        lat1d = np.asarray(ds.variables["lat"][:], dtype=np.float32)
        lon1d = np.asarray(ds.variables["lon"][:], dtype=np.float32)
    lon2d, lat2d = np.meshgrid(lon1d, lat1d)   # both (502, 583)
    return lat2d.astype(np.float32), lon2d.astype(np.float32)


def read_hour_mean_precip(fpath):
    """
    Open one hourly file, read precip_total_surf_mass_flux (30, 502, 583),
    mask fill values, and return the time-mean (502, 583) in m/s.

    Returns (mean_array, n_slices_used) or raises on failure.
    """
    with Dataset(fpath, "r") as ds:
        if SRC_VAR not in ds.variables:
            raise KeyError(f"{SRC_VAR} not in {fpath}")
        v = ds.variables[SRC_VAR]
        data = np.asarray(v[:], dtype=np.float64)   # (time, lat, lon)
    # Explicit fill-value masking (netCDF4 auto-mask is lost after np.asarray)
    bad = ~np.isfinite(data) | (np.abs(data) > FILL_BIG)
    if bad.any():
        data = np.where(bad, np.nan, data)
    # nanmean over the time dim; if a slice is all-NaN at a point, nanmean -> NaN
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)
        mean2d = np.nanmean(data, axis=0)
    n_slices = data.shape[0]
    return mean2d.astype(np.float32), n_slices


def write_output(out_path, times_dt, lat2d, lon2d, precip_rate, src_dir):
    """Write one consolidated hourly precipitation file."""
    nt, ny, nx = precip_rate.shape

    # Encode time as "hours since CASE_START"
    time_units = f"hours since {CASE_START.strftime('%Y-%m-%d %H:%M:%S')}"
    time_vals  = np.array(
        [(t - CASE_START).total_seconds() / 3600.0 for t in times_dt],
        dtype=np.float64,
    )

    with Dataset(out_path, "w", format="NETCDF4") as ds:
        # Dimensions
        ds.createDimension("time", nt)
        ds.createDimension("south_north", ny)
        ds.createDimension("west_east", nx)

        # time
        tv = ds.createVariable("time", "f8", ("time",))
        tv.units    = time_units
        tv.calendar = "standard"
        tv.long_name = "End of hourly accumulation window (UTC)"
        tv[:] = time_vals

        # XLAT
        latv = ds.createVariable("XLAT", "f4", ("south_north", "west_east"),
                                 zlib=True, complevel=4)
        latv.units = "degrees_north"
        latv.long_name = "Latitude"
        latv[:] = lat2d

        # XLONG
        lonv = ds.createVariable("XLONG", "f4", ("south_north", "west_east"),
                                 zlib=True, complevel=4)
        lonv.units = "degrees_east"
        lonv.long_name = "Longitude"
        lonv[:] = lon2d

        # PRECIP_RATE
        pr = ds.createVariable(
            "PRECIP_RATE", "f4", ("time", "south_north", "west_east"),
            zlib=True, complevel=4, fill_value=np.float32(np.nan),
        )
        pr.units = "mm hr-1"
        pr.long_name = "Hourly precipitation rate (water-equivalent)"
        pr.description = (
            "Mean of 30 instantaneous 2-min samples of "
            "precip_total_surf_mass_flux (m/s) within each hour, "
            "converted via x1000 (m->mm) and x3600 (s->hr)."
        )
        pr.source_variable = SRC_VAR
        pr[:] = precip_rate

        # Global attrs
        ds.title          = "TRACER-MIP hourly precipitation (E3SM-SCREAM)"
        ds.case           = "20220807 Houston"
        ds.source_dir     = src_dir
        ds.source_stream  = "2D.INST.INSTANT"
        ds.aggregation    = "mean over 30 instantaneous 2-min samples per hour"
        ds.timestamp_convention = "End of accumulation window (UTC)"
        ds.created        = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
        ds.created_by     = "Extract_Precip_E3SM_New.py"


# ============================================================
# PER-DATASET PROCESSING
# ============================================================
def process_dataset(label, subdir, dryrun=False, overwrite=False):
    src_dir  = os.path.join(BASE_PATH, subdir)
    out_path = os.path.join(OUTPUT_DIR, f"{label}_Aug7_Precip.nc")

    print("\n" + "=" * 90)
    print(f"DATASET: {label}")
    print("=" * 90)
    print(f"  src:  {src_dir}")
    print(f"  out:  {out_path}")

    if not os.path.isdir(src_dir):
        print("  [SKIP] source directory does not exist")
        return False

    if os.path.exists(out_path) and not overwrite and not dryrun:
        print("  [SKIP] output exists (use --overwrite to replace)")
        return False

    files = find_hourly_files(src_dir)
    if not files:
        print("  [SKIP] no matching 2D INST.INSTANT files in window")
        return False

    print(f"  Found {len(files)} hourly files in window "
          f"[{CASE_START} -> {CASE_END}]")
    print(f"  First: {os.path.basename(files[0][2])}")
    print(f"  Last:  {os.path.basename(files[-1][2])}")

    if dryrun:
        print("  [DRYRUN] no extraction performed")
        return True

    # Allocate
    nt = len(files)
    lat2d, lon2d = read_coords(files[0][2])
    ny, nx = lat2d.shape
    precip = np.full((nt, ny, nx), np.nan, dtype=np.float32)
    times_dt = []

    t0 = time.time()
    for i, (win_start, win_end, fpath) in enumerate(files):
        try:
            mean_ms, n = read_hour_mean_precip(fpath)
        except Exception as e:
            print(f"  [ERR ] hour {i+1}/{nt} {win_end}: {e}")
            times_dt.append(win_end)
            continue
        precip[i] = mean_ms * M_TO_MM * S_TO_HR
        times_dt.append(win_end)
        if (i + 1) % 6 == 0 or i == nt - 1:
            elapsed = time.time() - t0
            print(f"  [{i+1:>3d}/{nt}] {win_end}  n_slices={n}  "
                  f"max={np.nanmax(precip[i]):.2f} mm/hr  "
                  f"({elapsed:.1f}s elapsed)")

    # Sanity stats
    valid = np.isfinite(precip)
    if valid.any():
        print(f"  Precip stats: min={np.nanmin(precip):.3f}, "
              f"max={np.nanmax(precip):.2f}, "
              f"mean={np.nanmean(precip):.4f} mm/hr "
              f"(valid fraction={valid.mean():.3f})")
    else:
        print("  [WARN] all-NaN precipitation array")

    # Ensure output dir
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    write_output(out_path, times_dt, lat2d, lon2d, precip, src_dir)
    print(f"  [OK] wrote {out_path}")
    return True


# ============================================================
# MAIN
# ============================================================
def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--dataset", choices=list(DATASETS.keys()),
                    help="Process only one dataset (default: all)")
    ap.add_argument("--dryrun", action="store_true",
                    help="List input files, do not extract")
    ap.add_argument("--overwrite", action="store_true",
                    help="Overwrite existing output files")
    args = ap.parse_args()

    targets = ([args.dataset] if args.dataset else list(DATASETS.keys()))
    print(f"Processing {len(targets)} dataset(s): {targets}")
    print(f"Output dir: {OUTPUT_DIR}")
    print(f"Case window: {CASE_START} -> {CASE_END}")

    t_all = time.time()
    n_ok = 0
    for label in targets:
        try:
            if process_dataset(label, DATASETS[label],
                               dryrun=args.dryrun, overwrite=args.overwrite):
                n_ok += 1
        except Exception:
            print(f"\n  [FATAL] {label}")
            traceback.print_exc()

    print("\n" + "=" * 90)
    print(f"DONE. {n_ok}/{len(targets)} datasets processed in "
          f"{time.time() - t_all:.1f} s")
    print("=" * 90)


if __name__ == "__main__":
    main()