# FILE: parsivel_daily_terrace.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Parsivel monthly TXT → daily NetCDF (UNBC TERRACE) — lean deps, CF-friendly fills

Same engine as UQAM version; station metadata + defaults differ.
"""

import argparse, csv, os, re, sys
from datetime import datetime, timezone, date
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

import numpy as np
from netCDF4 import Dataset

# ---------- Station / paths ----------
IN_DIR   = Path("/instruments/TERRACE/UNBC_Rooftop/parsivel")
PATTERN  = "**/master/*.txt"
OUT_DIR  = Path("/station/instruments/disdrodb/PAR_UNBC_TERRACE")

CAMPAIGN = "MECHE / UNBC Terrace"
SITE_NAME = "UNBC Rooftop (Terrace)"
LAT, LON, ALT_M = 54.5186111, -128.6038889, 67.0   # adjust if needed

SKIP_EXISTING  = True
VERIFY_SAMPLES = False
DEFAULT_COMP   = 4

# class tables reused from above
D_MID = np.array([0.062,0.187,0.312,0.437,0.562,0.687,0.812,0.937,1.062,1.187,1.375,1.625,1.875,2.125,2.375,2.750,3.250,3.750,4.250,4.750,5.500,6.500,7.500,8.500,9.500,11.000,13.000,15.000,17.000,19.000,21.500,24.500], dtype=np.float32)
D_SPR = np.array([0.125,0.125,0.125,0.125,0.125,0.125,0.125,0.125,0.125,0.125,0.250,0.250,0.250,0.250,0.250,0.500,0.500,0.500,0.500,0.500,1.000,1.000,1.000,1.000,1.000,2.000,2.000,2.000,2.000,2.000,3.000,3.000], dtype=np.float32)
V_MID = np.array([0.05,0.15,0.25,0.35,0.45,0.55,0.65,0.75,0.85,0.95,1.15,1.35,1.55,1.75,1.95,2.25,2.75,3.25,3.75,4.25,5.00,6.00,7.00,8.00,9.00,11.00,13.00,15.00,17.00,19.00,22.00,26.00], dtype=np.float32)
D_LOW = (D_MID - 0.5*D_SPR).astype(np.float32); D_UP = (D_MID + 0.5*D_SPR).astype(np.float32)
V_SPR = np.array([0.10,0.10,0.10,0.10,0.10,0.10,0.10,0.10,0.10,0.10,0.20,0.20,0.20,0.20,0.20,0.50,0.50,0.50,0.50,0.50,1.00,1.00,1.00,1.00,1.00,2.00,2.00,2.00,2.00,2.00,3.00,4.00], dtype=np.float32)
V_LOW = (V_MID - 0.5*V_SPR).astype(np.float32); V_UP  = (V_MID + 0.5*V_SPR).astype(np.float32)

_NUM_KEEP = re.compile(r"[^\d\.\+\-eE]")
SPEC_RE1 = re.compile(r"^V(\d{1,2})D(\d{1,2})$")
SPEC_RE2 = re.compile(r"^D(\d{1,2})V(\d{1,2})$")
TIME_CANDIDATES = ("Timestamp","timestamp","time","DateTime","datetime","Date","UTC","Time_UTC")

EXTRA_MAP = {
    "rainfall_rate_32bit": ("rainfall_rate_32bit","Intensity of precipitation","Intensity of precipitation (mm/h)"),
    "reflectivity_32bit":  ("reflectivity_32bit","Radar reflectivity","Radar reflectivity (dBz)"),
    "mor_visibility":      ("mor_visibility","MOR Visibility (m)"),
    "laser_amplitude":     ("laser_amplitude","Signal amplitude of Laserband","Sensor voltage (V)"),
    "number_particles_validated": ("number_particles_validated","Number of detected particles"),
    "sensor_temperature":  ("sensor_temperature","Temperature in sensor (C)"),
    "kinetic_energy":      ("kinetic_energy","Kinetic Energy"),
    "weather_code_metar_4678": ("weather_code_metar_4678","Weather code METAR"),
    "weather_code_nws":    ("weather_code_nws","Weather code NWS"),
    "error_code":          ("error_code","Error code"),
}

# — utilities (same as above, verbatim) —
def info(m: str): print(f"[info {datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%SZ')}] {m}")
def canon(s: str) -> str: return re.sub(r"[^0-9a-z]+", "", s.lower())
def find_col(header: List[str], *aliases: str) -> Optional[int]:
    idx = {canon(c): i for i, c in enumerate(header)}
    for a in aliases:
        ca = canon(a)
        if ca in idx: return idx[ca]
    for a in aliases:
        ca = canon(a)
        for i, c in enumerate(header):
            if ca in canon(c): return i
    return None
def detect_sep(path: Path) -> str:
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        lines = [f.readline() for _ in range(10)]
    c = sum(l.count(",") for l in lines if l)
    s = sum(l.count(";") for l in lines if l)
    return "," if c >= s else ";"
def parse_utc_timestamp(s: str):
    if not s: return None
    t = s.strip()
    if t.endswith("Z"): t = t[:-1] + "+00:00"
    try:
        dt = datetime.fromisoformat(t)
        dt = dt.replace(tzinfo=dt.tzinfo or timezone.utc)
        return dt.astimezone(timezone.utc)
    except Exception: pass
    for fmt in ("%Y-%m-%d %H:%M:%S","%Y/%m/%d %H:%M:%S","%Y-%m-%d %H:%M","%Y-%m-%dT%H:%M:%S","%Y-%m-%dT%H:%M"):
        try: return datetime.strptime(t, fmt).replace(tzinfo=timezone.utc)
        except Exception: pass
    return None
def to_float(s: str) -> float:
    if s is None: return np.nan
    t = s.strip()
    if t == "" or t.lower() in ("nan","none","null"): return np.nan
    if re.search(r"\d,\d", t) and not re.search(r"\d\.\d", t):
        t = t.replace(".", "").replace(",", ".")
    else:
        t = _NUM_KEEP.sub("", t)
    try: return float(t)
    except Exception: return np.nan
def month_from_file(fp: Path) -> Optional[Tuple[int,int]]:
    try:
        ym = fp.parent.parent.name
        m = re.match(r"^(\d{4})[_-](\d{2})$", ym)
        if m: return int(m.group(1)), int(m.group(2))
    except Exception: pass
    return None
def today_iso_toronto() -> str:
    try:
        import zoneinfo
        tz = zoneinfo.ZoneInfo("America/Toronto")
        return datetime.now(tz).date().isoformat()
    except Exception:
        return date.today().isoformat()
def month_paths(in_root: Path, month: Optional[str], pattern: str) -> List[Path]:
    if month:
        ym = month.replace("-", "_"); base = in_root / ym
        return sorted(base.rglob(pattern))
    return sorted(in_root.rglob(pattern))
def existing_days_for_month(out_root: Path, year: int, month: int, id_tag: str) -> Set[str]:
    ydir = out_root / f"{year:04d}" / f"{year:04d}{month:02d}"
    out = set()
    if ydir.exists():
        for p in ydir.glob(f"{id_tag}_{year:04d}{month:02d}*.nc"):
            ymd = p.stem.split("_")[-1]
            try: out.add(f"{ymd[:4]}-{ymd[4:6]}-{ymd[6:8]}")
            except Exception: pass
    return out
def read_time_len_fast(nc_path: Path) -> Optional[int]:
    try:
        with Dataset(str(nc_path), "r") as ds:
            if "time" in ds.dimensions: return len(ds.dimensions["time"])
            if "epoch_time" in ds.variables: return int(ds.variables["epoch_time"].shape[0])
        return None
    except Exception: return None
def plan_days(counts: Dict[str,int], out_root: Path, id_tag: str, fp: Path,
              skip_existing: bool, verify_samples: bool, forced_days: Set[str]) -> Set[str]:
    today = today_iso_toronto()
    counts = {d:c for d,c in counts.items() if d < today}
    forced_days = {d for d in forced_days if d < today}
    to_do: Set[str] = set()
    ym = month_from_file(fp); existing: Set[str] = set()
    if ym:
        y, m = ym; existing = {d for d in existing_days_for_month(out_root, y, m, id_tag) if d < today}
    missing = set(counts.keys()) - existing
    if skip_existing and not verify_samples and not forced_days and not missing: return set()
    to_do.update(missing)
    if verify_samples and ym:
        y, m = ym
        for day in (set(counts.keys()) & existing):
            ymd = day.replace("-", "")
            out_fp = out_root / f"{y:04d}" / f"{y:04d}{m:02d}" / f"{id_tag}_{ymd}.nc"
            have = read_time_len_fast(out_fp)
            if have is None or have != counts[day]: to_do.add(day)
    to_do.update(forced_days)
    if skip_existing and not verify_samples: to_do = (missing | (forced_days - existing))
    return to_do

# --- writer helpers (same CF fill logic as UQAM) ---
F32_FILL = np.float32(-9999.0)
def _add_scalar(ds, name, dtype, value, **attrs):
    v = ds.createVariable(name, dtype); v[...] = value
    for k, val in attrs.items(): v.setncattr(k, val); return v
def _add_1d(ds, name, dtype, dims, data=None, fill=None, z=DEFAULT_COMP, **attrs):
    kw = {"zlib": True, "complevel": z, "shuffle": True}
    v = ds.createVariable(name, dtype, dims, fill_value=fill, **kw) if fill is not None \
        else ds.createVariable(name, dtype, dims, **kw)
    if data is not None: v[:] = data
    for k, val in attrs.items(): v.setncattr(k, val)
    return v
def _add_nd(ds, name, dtype, dims, data=None, fill=None, z=DEFAULT_COMP, **attrs):
    kw = {"zlib": True, "complevel": z, "shuffle": True}
    v = ds.createVariable(name, dtype, dims, fill_value=fill, **kw) if fill is not None \
        else ds.createVariable(name, dtype, dims, **kw)
    if data is not None: v[:] = data
    for k, val in attrs.items(): v.setncattr(k, val)
    return v

def write_day_netcdf(out_fp: Path,
                     epochs: np.ndarray,
                     t_strings: List[str],
                     raw3d: Optional[np.ndarray],
                     fieldN: Optional[np.ndarray],
                     fieldV: Optional[np.ndarray],
                     extras: Dict[str, np.ndarray],
                     complevel: int,
                     id_tag: str) -> None:
    out_fp.parent.mkdir(parents=True, exist_ok=True)
    with Dataset(str(out_fp), "w", format="NETCDF4") as ds:
        n = int(len(epochs))
        ds.createDimension("time", n)
        ds.createDimension("diameter_classes", 32)
        ds.createDimension("velocity_classes", 32)

        _add_scalar(ds, "latitude",  "f4", LAT, units="degrees_north")
        _add_scalar(ds, "longitude", "f4", LON, units="degrees_east")
        _add_scalar(ds, "altitude",  "f4", ALT_M, units="m")

        _add_1d(ds, "lower_diameter_class_limits", "f4", ("diameter_classes",), D_LOW, F32_FILL, z=complevel, units="mm")
        _add_1d(ds, "upper_diameter_class_limits", "f4", ("diameter_classes",), D_UP,  F32_FILL, z=complevel, units="mm")
        _add_1d(ds, "lower_velocity_class_limits", "f4", ("velocity_classes",), V_LOW, F32_FILL, z=complevel, units="m s-1")
        _add_1d(ds, "upper_velocity_class_limits", "f4", ("velocity_classes",), V_UP,  F32_FILL, z=complevel, units="m s-1")

        _add_1d(ds, "epoch_time", "i8", ("time",), epochs.astype(np.int64, copy=False), None, z=complevel,
                units="seconds since 1970-01-01 00:00:00 UTC")
        ts = ds.createVariable("time_as_string", str, ("time",)); ts[:] = np.asarray(t_strings, dtype=object)

        if raw3d is not None:
            _add_nd(ds, "raw_data", "f4", ("time","diameter_classes","velocity_classes"),
                    raw3d.astype(np.float32, copy=False), F32_FILL, z=complevel, long_name="drop size-velocity spectra")
        if fieldN is not None:
            _add_nd(ds, "fieldN", "f4", ("time","diameter_classes"),
                    fieldN.astype(np.float32, copy=False), F32_FILL, z=complevel, long_name="sum over velocity bins", units="1")
        if fieldV is not None:
            _add_nd(ds, "fieldV", "f4", ("time","diameter_classes"),
                    fieldV.astype(np.float32, copy=False), F32_FILL, z=complevel, long_name="velocity-weighted mean", units="m s-1")

        for name, arr in extras.items():
            if arr.dtype.kind in ("U","O"):
                vv = ds.createVariable(name, str, ("time",)); vv[:] = np.asarray(arr, dtype=object)
            else:
                arr32 = arr.astype(np.float32, copy=False)
                _add_1d(ds, name, "f4", ("time",), arr32, F32_FILL, z=complevel)

        ds.setncattr("title", "OTT Parsivel2 disdrometer data (UNBC Terrace)")
        ds.setncattr("institution", "UNBC / MECHE")
        ds.setncattr("source", "CANADA")
        ds.setncattr("history", f"produced on {datetime.utcnow().strftime('%Y-%m-%d')} by Terrace pipeline")
        ds.setncattr("sensor_type", "OTT Hydromet Parsivel2 optical disdrometer")
        ds.setncattr("site_name", SITE_NAME)
        ds.setncattr("sensor_name", id_tag)
        ds.setncattr("logging_software", "Custom Terrace pipeline")
        ds.setncattr("project_name", CAMPAIGN)
        ds.setncattr("data_license", "CC BY 4.0 https://creativecommons.org/licenses/by/4.0/")

def process_month(fp: Path, out_root: Path, id_tag: str, complevel: int,
                  skip_existing: bool, verify_samples: bool,
                  forced_days: Set[str], write_empty_forced: bool) -> List[Path]:
    sep = detect_sep(fp)
    with open(fp, "r", encoding="utf-8", errors="ignore", newline="") as f:
        r = csv.reader(f, delimiter=sep)
        try: header = next(r)
        except StopIteration: return []

    t_idx = find_col(header, *TIME_CANDIDATES)
    if t_idx is None: 
        info(f"[skip] time column not found in {fp}")
        return []

    spec_cols = []
    extra_idx: Dict[str, Tuple[int,str]] = {}
    for i, h in enumerate(header):
        m1 = SPEC_RE1.match(h); m2 = SPEC_RE2.match(h)
        if m1:
            v, d = int(m1.group(1))-1, int(m1.group(2))-1
            if 0<=d<32 and 0<=v<32: spec_cols.append((i,d,v)); continue
        if m2:
            d, v = int(m2.group(1))-1, int(m2.group(2))-1
            if 0<=d<32 and 0<=v<32: spec_cols.append((i,d,v)); continue

    for name, aliases in EXTRA_MAP.items():
        col = find_col(header, *aliases)
        if col is None: continue
        dtype = "f4"
        extra_idx[name] = (col, dtype)

    # pre-count per day
    counts: Dict[str,int] = {}
    with open(fp, "r", encoding="utf-8", errors="ignore", newline="") as f:
        r = csv.reader(f, delimiter=sep); next(r, None)
        for row in r:
            if not row or t_idx >= len(row): continue
            dt = parse_utc_timestamp(row[t_idx]); 
            if dt is None: continue
            d = dt.date().isoformat()
            counts[d] = counts.get(d, 0) + 1

    id_tag = out_root.name
    days_to_do = plan_days(counts, out_root, id_tag, fp, skip_existing, verify_samples, forced_days)
    if not days_to_do:
        info(f"[up-to-date] {fp}")
        return []

    # buffers
    day_buffers: Dict[str, Dict] = {}
    for day, n in counts.items():
        if day not in days_to_do: continue
        buf = {
            "n": n, "pos": 0,
            "epochs": np.empty(n, dtype=np.int64),
            "tstr":   np.empty(n, dtype=object),
            "raw":    np.full((n, 32, 32), np.nan, dtype=np.float32) if spec_cols else None,
            "extras": {},
        }
        for name, (col, dtype) in extra_idx.items():
            if dtype == "f4":
                buf["extras"][name] = np.full(n, np.nan, dtype=np.float32)
            else:
                buf["extras"][name] = np.full(n, np.nan, dtype=np.float64)
        day_buffers[day] = buf

    # stream rows
    with open(fp, "r", encoding="utf-8", errors="ignore", newline="") as f:
        r = csv.reader(f, delimiter=sep); next(r, None)
        for row in r:
            if not row or t_idx >= len(row): continue
            dt = parse_utc_timestamp(row[t_idx]); 
            if dt is None: continue
            day = dt.date().isoformat()
            if day not in day_buffers: continue
            buf = day_buffers[day]; n = buf.get("n", 0)
            if n == 0: continue
            i = buf["pos"]
            if i >= n:
                grow = 256
                new_epochs = np.empty(n + grow, dtype=buf["epochs"].dtype); new_epochs[:n] = buf["epochs"]; buf["epochs"] = new_epochs
                new_tstr   = np.empty(n + grow, dtype=object); new_tstr[:n] = buf["tstr"]; buf["tstr"] = new_tstr
                if buf["raw"] is not None:
                    old = buf["raw"]; new_raw = np.full((old.shape[0] + grow, 32, 32), np.nan, dtype=old.dtype)
                    new_raw[:old.shape[0]] = old; buf["raw"] = new_raw
                for ex, arr in list(buf["extras"].items()):
                    new = np.full(arr.shape[0] + grow, np.nan, dtype=arr.dtype); new[:arr.shape[0]] = arr
                    buf["extras"][ex] = new
                buf["n"] = n + grow; n = buf["n"]

            buf["epochs"][i] = int(dt.timestamp())
            buf["tstr"][i]   = dt.strftime("%Y-%m-%d %H:%M:%S")

            if buf["raw"] is not None:
                grid = buf["raw"][i]
                for (col,d,v) in spec_cols:
                    if col < len(row):
                        try:
                            grid[d,v] = np.float32(to_float(row[col]))
                        except Exception:
                            grid[d,v] = np.float32(np.nan)

            for name, (col, dtype) in extra_idx.items():
                if col >= len(row): continue
                buf["extras"][name][i] = np.float32(to_float(row[col]))

            buf["pos"] = i + 1

    # finalize
    written: List[Path] = []
    for day, buf in day_buffers.items():
        n = int(buf.get("pos", 0)) if buf.get("n", 0) else 0
        ymd = day.replace("-", "")
        y = int(ymd[:4]); m = int(ymd[4:6])
        out_dir = out_root / f"{y:04d}" / f"{y:04d}{m:02d}"
        out_fp = out_dir / f"{OUT_DIR.name}_{ymd}.nc"

        if n == 0:
            # write empty only if explicitly forced (to avoid “today” premature files)
            if day in days_to_do:
                write_day_netcdf(out_fp, np.array([], dtype=np.int64), [], None, None, None, {}, DEFAULT_COMP, OUT_DIR.name)
                info(f"wrote (empty): {out_fp}"); written.append(out_fp)
            continue

        epochs = buf["epochs"][:n].astype(np.int64, copy=False)
        tstr   = list(buf["tstr"][:n])
        raw3d  = buf["raw"][:n] if buf["raw"] is not None else None

        fieldN = None; fieldV = None
        if raw3d is not None:
            fieldN = np.nansum(raw3d, axis=2).astype(np.float32)
            V = V_MID.astype(np.float64)[None, None, :]
            num = np.nansum(raw3d * V, axis=2)
            den = np.nansum(raw3d, axis=2)
            with np.errstate(invalid="ignore", divide="ignore"):
                fv = (num / den); fv[~np.isfinite(fv)] = np.nan
            fieldV = fv.astype(np.float32)

        extras_out: Dict[str, np.ndarray] = {}
        for name, arr in buf.get("extras", {}).items():
            extras_out[name] = arr[:n]

        write_day_netcdf(out_fp, epochs, tstr, raw3d, fieldN, fieldV, extras_out, DEFAULT_COMP, OUT_DIR.name)
        info(f"wrote: {out_fp}"); written.append(out_fp)

    return written

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--in-dir", type=Path, default=IN_DIR)
    p.add_argument("--pattern", type=str, default=PATTERN)
    p.add_argument("--out-dir", type=Path, default=OUT_DIR)
    p.add_argument("--month", type=str, default=None, help="YYYY-MM or omit for all under in-dir")
    p.add_argument("--no-skip-existing", action="store_true")
    p.add_argument("--verify-samples", action="store_true")
    p.add_argument("--force-days", nargs="*", default=[], help="YYYY-MM-DD ...")
    p.add_argument("--write-empty-forced", action="store_true")
    p.add_argument("--complevel", type=int, default=DEFAULT_COMP)
    p.add_argument("--fast-skip", action="store_true")
    args = p.parse_args()

    in_dir = args.in_dir
    out_dir = args.out_dir
    verify = args.verify_samples or args.fast_skip
    skip = not args.no_skip_existing
    forced = set(args.force_days)

    txts = month_paths(in_dir, args.month, args.pattern)
    if not txts:
        info(f"No input under {in_dir} (pattern={args.pattern}, month={args.month})")
        return 0

    for fp in txts:
        info(f"TXT: {fp}")
        process_month(fp, out_dir, out_dir.name, args.complevel, skip, verify, forced, args.write_empty_forced)
    return 0

if __name__ == "__main__":
    sys.exit(main())