"""Normalize one ECMWF SEAS5.1 initialization to strict monthly fields."""
from __future__ import annotations

import argparse
import logging
import os
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import xarray as xr

from .config import PROJECT_ROOT
from .download_ecmwf_baseline import parse_init
from .precipitation import aggregate_monthly_precipitation, deaccumulate_daily_precipitation

LOGGER = logging.getLogger(__name__)
INSTANT = {
    "t2m": ("t2m_mod", "degC", 1.0, -273.15),
    "sst": ("sst", "degC", 1.0, -273.15),
    "stl1": ("st1", "degC", 1.0, -273.15),
    "msl": ("msl", "hPa", 0.01, 0.0),
}
FLUX = {"slhf": ("lh", 1.0), "sshf": ("sh", 1.0),
        "tsr": ("rn", 1.0), "ttr": ("olr", -1.0)}


def month_bounds(init: str) -> tuple[np.datetime64, np.datetime64, np.ndarray]:
    parse_init(init)
    start = np.datetime64(f"{init}-01", "ns")
    month = start.astype("datetime64[M]")
    boundary = (month + np.timedelta64(7, "M")).astype("datetime64[ns]")
    months = (month + np.arange(7).astype("timedelta64[M]")).astype("datetime64[ns]")
    return start, boundary, months


def _prepare(source: xr.DataArray) -> xr.DataArray:
    rename = {old: new for old, new in (("number", "member"), ("lat", "latitude"),
              ("lon", "longitude")) if old in source.dims or old in source.coords}
    result = source.rename(rename)
    if "forecast_reference_time" in result.dims:
        if result.sizes["forecast_reference_time"] != 1:
            raise ValueError("Expected exactly one forecast_reference_time")
        result = result.isel(forecast_reference_time=0, drop=True)
    required = {"member", "forecast_period", "latitude", "longitude"}
    missing = required - set(result.dims)
    if missing:
        raise ValueError(f"{source.name} is missing dimensions: {sorted(missing)}")
    if "valid_time" not in result.coords or result.valid_time.dims != ("forecast_period",):
        raise ValueError(f"{source.name} requires 1-D valid_time along forecast_period")
    return result


def _promote_valid_time(dataset: xr.Dataset) -> xr.Dataset:
    if "valid_time" not in dataset:
        raise ValueError("Dataset is missing valid_time")
    valid_time = dataset["valid_time"]
    if valid_time.dims != ("forecast_period",):
        raise ValueError(
            "valid_time must be 1-D along forecast_period; "
            f"found dimensions {valid_time.dims}"
        )
    return dataset.set_coords("valid_time")


def _reference_time(dataset: xr.Dataset, init: str) -> None:
    if "forecast_reference_time" not in dataset:
        raise ValueError("Dataset is missing forecast_reference_time")
    actual = np.asarray(dataset.forecast_reference_time.values).reshape(-1)
    expected = np.datetime64(f"{init}-01T00:00:00", "s")
    if actual.size != 1 or actual[0].astype("datetime64[s]") != expected:
        raise ValueError(f"forecast_reference_time does not match {expected}")


def _lead_hours(source: xr.DataArray) -> np.ndarray:
    values = np.asarray(source.forecast_period.values)
    hours = values / np.timedelta64(1, "h") if np.issubdtype(values.dtype, np.timedelta64) else values.astype(float)
    if hours.ndim != 1 or not np.all(np.isfinite(hours)) or not np.allclose(hours, np.round(hours)):
        raise ValueError("forecast_period must be finite whole hours")
    return np.round(hours).astype(np.int64)


def _add_lead(monthly: xr.DataArray, months: np.ndarray) -> xr.DataArray:
    actual = np.asarray(monthly.valid_time.values).astype("datetime64[ns]")
    if not np.array_equal(actual, months):
        raise ValueError("Monthly valid_time does not equal the seven LT=0..6 month starts")
    return monthly.assign_coords(lead=("valid_time", np.arange(7, dtype=np.int16))).swap_dims({"valid_time": "lead"})


def instantaneous_monthly(source: xr.DataArray, init: str, name: str,
                          units: str, scale: float, offset: float) -> xr.DataArray:
    """Average complete 6-hourly point samples in half-open calendar months."""
    da = _prepare(source)
    source_units = str(da.attrs.get("units", "")).strip()
    required_units = "Pa" if units == "hPa" else "K"
    if source_units != required_units:
        raise ValueError(f"{source.name} units must be {required_units}; found {source_units!r}")
    start, boundary, months = month_bounds(init)
    valid = np.asarray(da.valid_time.values).astype("datetime64[ns]")
    if valid.size > 1 and np.any(np.diff(valid) != np.timedelta64(6, "h")):
        raise ValueError("Instantaneous data contain missing or duplicate 6-hour steps")
    keep = (valid >= start) & (valid < boundary)
    actual = valid[keep]
    expected = np.arange(start, boundary, np.timedelta64(6, "h"))
    if not (np.array_equal(actual, expected) or np.array_equal(actual, expected[1:])):
        raise ValueError("Instantaneous coverage must be 6-hourly over [init, LT7); only step 0 may be absent")
    work = da.isel(forecast_period=np.flatnonzero(keep)).assign_coords(
        sample_time=("forecast_period", actual)).swap_dims({"forecast_period": "sample_time"})
    work = work.drop_vars("forecast_period")
    monthly = (work * scale + offset).resample(sample_time="MS").mean(skipna=False)
    monthly = _add_lead(monthly.rename({"sample_time": "valid_time"}), months).rename(name)
    monthly.attrs = {"units": units, "source_variable": str(source.name),
                     "source_units": source_units,
                     "temporal_aggregation": "mean 6-hourly instantaneous samples over [month_start,next_month_start)",
                     "initial_step_policy": "step 0 optional; no other missing step allowed"}
    return monthly


def daily_flux(source: xr.DataArray, sign: float = 1.0) -> xr.DataArray:
    """De-accumulate J m-2 since initialization to daily-mean W m-2."""
    da = _prepare(source)
    hours = _lead_hours(da)
    if hours.size == 0 or hours[0] != 24 or np.any(np.diff(hours) != 24):
        raise ValueError("Accumulated flux needs consecutive 24-hour steps starting at 24")
    source_units = str(da.attrs.get("units", "")).strip()
    if source_units not in {"J m**-2", "J m-2", "J m^-2", "J m⁻²"}:
        raise ValueError(f"Accumulated flux units must be J m-2; found {source_units!r}")
    increments = xr.concat([da.isel(forecast_period=slice(0, 1)),
                            da.diff("forecast_period", label="upper")], dim="forecast_period")
    result = increments * sign / 86400.0
    result.attrs = {"units": "W m-2", "source_units": source_units,
                    "sign_multiplier": sign,
                    "temporal_processing": "X(24) then X(t)-X(t-24), divided by 86400 s"}
    return result


def daily_mean_monthly(source: xr.DataArray, init: str, name: str) -> xr.DataArray:
    """Average complete 24-hour intervals over LT=0..6."""
    start, boundary, months = month_bounds(init)
    valid = np.asarray(source.valid_time.values).astype("datetime64[ns]")
    interval = valid - np.timedelta64(24, "h")
    if interval.size > 1 and np.any(np.diff(interval) != np.timedelta64(24, "h")):
        raise ValueError("Daily field contains missing or duplicate intervals")
    keep = (interval >= start) & (interval < boundary)
    expected = np.arange(start, boundary, np.timedelta64(24, "h"))
    if not np.array_equal(interval[keep], expected):
        raise ValueError("Daily field does not completely cover LT=0..6")
    work = source.isel(forecast_period=np.flatnonzero(keep)).assign_coords(
        interval_start=("forecast_period", interval[keep])).swap_dims({"forecast_period": "interval_start"})
    for coordinate in ("forecast_period", "valid_time"):
        if coordinate in work.coords:
            work = work.drop_vars(coordinate)
    monthly = work.resample(interval_start="MS").mean(skipna=False).rename({"interval_start": "valid_time"})
    result = _add_lead(monthly, months).rename(name)
    result.attrs = dict(source.attrs)
    result.attrs["temporal_aggregation"] = "mean complete daily intervals over [month_start,next_month_start)"
    return result


def build_monthly_dataset(instantaneous: xr.Dataset, accumulated: xr.Dataset, init: str,
                          instantaneous_source: str = "", accumulated_source: str = "") -> xr.Dataset:
    """Return normalized LT=0..6 fields without collapsing ensemble members."""
    _reference_time(instantaneous, init)
    _reference_time(accumulated, init)
    instantaneous = _promote_valid_time(instantaneous)
    accumulated = _promote_valid_time(accumulated)
    missing_i = sorted(set(INSTANT) - set(instantaneous.data_vars))
    missing_a = sorted({"tp", *FLUX} - set(accumulated.data_vars))
    if missing_i or missing_a:
        raise ValueError(f"Missing variables: instantaneous={missing_i}, accumulated={missing_a}")
    fields = [instantaneous_monthly(instantaneous[key], init, *spec) for key, spec in INSTANT.items()]
    tp_daily = deaccumulate_daily_precipitation(_prepare(accumulated.tp))
    _, _, months = month_bounds(init)
    fields.append(_add_lead(aggregate_monthly_precipitation(tp_daily), months).rename("pr_mod"))
    fields.extend(daily_mean_monthly(daily_flux(accumulated[key], sign), init, name)
                  for key, (name, sign) in FLUX.items())
    result = xr.merge(fields, compat="equals", join="exact")
    init_time = np.datetime64(f"{init}-01T00:00:00", "ns")
    result = result.expand_dims(init_time=[init_time]).transpose(
        "init_time", "member", "lead", "latitude", "longitude")
    result.attrs = {"title": "ECMWF SEAS5.1 normalized monthly fields",
                    "source": "Copernicus CDS seasonal-original-single-levels; ECMWF System 51",
                    "lead_convention": "LT=0 initial month; released products LT=1..6",
                    "temporal_boundary": "half-open months; LT7 00 UTC excluded",
                    "ensemble_processing": "all members retained; no ensemble mean",
                    "instantaneous_source": instantaneous_source,
                    "accumulated_source": accumulated_source,
                    "software": "dr-gcm 0.1.0",
                    "created_utc": datetime.now(timezone.utc).isoformat()}
    return result


def validate_monthly_file(path: Path, init: str) -> dict:
    required = {"t2m_mod", "pr_mod", "sst", "st1", "msl", "lh", "sh", "rn", "olr"}
    with xr.open_dataset(path) as dataset:
        missing = sorted(required - set(dataset.data_vars))
        if missing:
            raise ValueError(f"Missing monthly variables: {missing}")
        expected = {"init_time": 1, "member": 25, "lead": 7}
        wrong = {key: dataset.sizes.get(key) for key, size in expected.items() if dataset.sizes.get(key) != size}
        if wrong:
            raise ValueError(f"Monthly dimension mismatch: {wrong}")
        if not np.array_equal(dataset.lead.values, np.arange(7)):
            raise ValueError("Monthly lead must be 0..6")
        return {"variables": sorted(required), "sizes": dict(dataset.sizes), "bytes": path.stat().st_size}


def run(init: str, instantaneous: Path, accumulated: Path, output: Path,
        overwrite: bool = False, retry_partial: bool = False) -> tuple[Path, dict]:
    if output.exists() and not overwrite:
        return output, validate_monthly_file(output, init)
    partial = output.with_suffix(output.suffix + ".part")
    if partial.exists() and not (overwrite or retry_partial):
        raise FileExistsError(f"Partial output retained: {partial}")
    if overwrite:
        output.unlink(missing_ok=True)
    if overwrite or retry_partial:
        partial.unlink(missing_ok=True)
    output.parent.mkdir(parents=True, exist_ok=True)
    with xr.open_dataset(instantaneous) as instant, xr.open_dataset(accumulated) as accum:
        dataset = build_monthly_dataset(instant, accum, init, str(instantaneous), str(accumulated))
        encoding = {name: {"zlib": True, "complevel": 1, "shuffle": True} for name in dataset.data_vars}
        dataset.to_netcdf(partial, encoding=encoding)
    report = validate_monthly_file(partial, init)
    os.replace(partial, output)
    return output, report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--init", required=True)
    parser.add_argument("--instantaneous", type=Path)
    parser.add_argument("--accumulated", type=Path)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--retry-partial", action="store_true")
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    raw = PROJECT_ROOT / "data/ecmwf/raw_split"
    instant = args.instantaneous or raw / f"ecmwf_seas51_{args.init}_instantaneous.nc"
    accum = args.accumulated or raw / f"ecmwf_seas51_{args.init}_accumulated.nc"
    output = args.output or PROJECT_ROOT / "data/ecmwf/monthly" / f"ecmwf_seas51_{args.init}_monthly.nc"
    target, report = run(args.init, instant, accum, output, args.overwrite, args.retry_partial)
    LOGGER.info("Validated %s: %s", target, report)


if __name__ == "__main__":
    main()