"""Fit and apply reusable SPEI parameters with atomic NetCDF outputs."""
from __future__ import annotations

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

import numpy as np
import xarray as xr

from .spei import fit_loglogistic, fit_thornthwaite_parameters, spei_from_parameters, water_balance

LOGGER = logging.getLogger(__name__)


def _validate_monthly(data: xr.DataArray, label: str) -> None:
    if "time" not in data.dims or not np.issubdtype(data.time.dtype, np.datetime64):
        raise ValueError(f"{label} requires a datetime time dimension")
    months = np.asarray(data.time.values).astype("datetime64[M]")
    if months.size == 0 or np.unique(months).size != months.size:
        raise ValueError(f"{label} has no months or contains duplicates")
    if months.size > 1 and np.any(np.diff(months).astype(int) != 1):
        raise ValueError(f"{label} contains missing monthly steps")


def _units(pr: xr.DataArray, t2m: xr.DataArray) -> None:
    if str(pr.attrs.get("units", "")).strip() not in {"mm", "mm/month", "mm month-1"}:
        raise ValueError("Precipitation units must be monthly mm")
    if str(t2m.attrs.get("units", "")).strip() not in {"degC", "C", "degree_Celsius", "degrees_Celsius"}:
        raise ValueError("Temperature units must be degrees Celsius")


def fit_parameters(pr: xr.DataArray, t2m: xr.DataArray,
                   start: str = "1981-01-01", end: str = "2020-12-31",
                   scales: Iterable[int] = (1, 3, 6)) -> xr.Dataset:
    """Fit PET and log-logistic parameters using only the calibration slice."""
    _validate_monthly(pr, "pr")
    _validate_monthly(t2m, "t2m")
    _units(pr, t2m)
    pr, t2m = xr.align(pr, t2m, join="exact")
    calibration_pr = pr.sel(time=slice(start, end))
    calibration_t2m = t2m.sel(time=slice(start, end))
    if calibration_pr.sizes.get("time", 0) == 0:
        raise ValueError("Calibration period selects no data")
    _validate_monthly(calibration_pr, "calibration_pr")
    thornthwaite = fit_thornthwaite_parameters(calibration_t2m)
    balance = water_balance(calibration_pr, calibration_t2m, thornthwaite)
    loglogistic = fit_loglogistic(balance, scales=scales)
    result = xr.merge([thornthwaite, loglogistic], compat="equals", join="exact")
    result.attrs = {
        "title": "Reusable SPEI calibration parameters",
        "pet_method": "Thornthwaite with calibration climatological heat index",
        "distribution": "three-parameter log-logistic (scipy.stats.fisk MLE)",
        "fit_dimensions": "grid cell, calendar month, timescale",
        "calibration_start": str(calibration_pr.time.min().values),
        "calibration_end": str(calibration_pr.time.max().values),
        "scales": ",".join(str(int(value)) for value in scales),
        "software": "dr-gcm 0.1.0",
        "created_utc": datetime.now(timezone.utc).isoformat(),
    }
    return result


def apply_parameters(pr: xr.DataArray, t2m: xr.DataArray, parameters: xr.Dataset,
                     scales: Iterable[int] = (1, 3, 6)) -> xr.Dataset:
    """Apply saved parameters to a contiguous series including required history."""
    _validate_monthly(pr, "pr")
    _validate_monthly(t2m, "t2m")
    _units(pr, t2m)
    pr, t2m = xr.align(pr, t2m, join="exact")
    balance = water_balance(pr, t2m, parameters)
    result = xr.merge([spei_from_parameters(balance, parameters, int(scale))
                       for scale in scales], compat="equals", join="exact")
    result.attrs = {
        "title": "SPEI from saved calibration parameters",
        "parameter_calibration_start": parameters.attrs.get("calibration_start", ""),
        "parameter_calibration_end": parameters.attrs.get("calibration_end", ""),
        "history_requirement": "input includes preceding observed months required by each timescale",
        "software": "dr-gcm 0.1.0",
        "created_utc": datetime.now(timezone.utc).isoformat(),
    }
    return result


def validate_parameters(path: Path) -> dict:
    required = {"heat_index", "thornthwaite_exponent", "loglogistic_parameters"}
    with xr.open_dataset(path) as dataset:
        missing = sorted(required - set(dataset.data_vars))
        if missing:
            raise ValueError(f"SPEI parameter file is missing: {missing}")
        if not np.array_equal(dataset.timescale.values, [1, 3, 6]):
            raise ValueError("SPEI parameter timescales must be 1,3,6")
        if not np.array_equal(dataset.month.values, np.arange(1, 13)):
            raise ValueError("SPEI parameter months must be 1..12")
        return {"variables": sorted(required), "sizes": dict(dataset.sizes), "bytes": path.stat().st_size}


def validate_output(path: Path) -> dict:
    required = {"spei_1", "spei_3", "spei_6"}
    with xr.open_dataset(path) as dataset:
        missing = sorted(required - set(dataset.data_vars))
        if missing:
            raise ValueError(f"SPEI output is missing: {missing}")
        return {"variables": sorted(required), "sizes": dict(dataset.sizes), "bytes": path.stat().st_size}


def write_atomic(dataset: xr.Dataset, output: Path, validator,
                 overwrite: bool = False, retry_partial: bool = False) -> tuple[Path, dict]:
    if output.exists() and not overwrite:
        return output, validator(output)
    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)
    encoding = {name: {"zlib": True, "complevel": 1, "shuffle": True} for name in dataset.data_vars}
    dataset.to_netcdf(partial, encoding=encoding)
    report = validator(partial)
    os.replace(partial, output)
    return output, report


def _read(path: Path, variable: str) -> xr.DataArray:
    with xr.open_dataset(path) as dataset:
        if variable not in dataset:
            raise KeyError(f"{variable!r} not found in {path}")
        return dataset[variable].load()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    fit = commands.add_parser("fit")
    apply = commands.add_parser("apply")
    for child in (fit, apply):
        child.add_argument("--precipitation", required=True, type=Path)
        child.add_argument("--temperature", required=True, type=Path)
        child.add_argument("--precipitation-variable", default="pr")
        child.add_argument("--temperature-variable", default="t2m")
        child.add_argument("--output", required=True, type=Path)
        child.add_argument("--overwrite", action="store_true")
        child.add_argument("--retry-partial", action="store_true")
    fit.add_argument("--calibration-start", default="1981-01-01")
    fit.add_argument("--calibration-end", default="2020-12-31")
    apply.add_argument("--parameters", required=True, type=Path)
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    pr = _read(args.precipitation, args.precipitation_variable)
    t2m = _read(args.temperature, args.temperature_variable)
    if args.command == "fit":
        dataset = fit_parameters(pr, t2m, args.calibration_start, args.calibration_end)
        target, report = write_atomic(dataset, args.output, validate_parameters,
                                      args.overwrite, args.retry_partial)
    else:
        with xr.open_dataset(args.parameters) as source:
            parameters = source.load()
        dataset = apply_parameters(pr, t2m, parameters)
        target, report = write_atomic(dataset, args.output, validate_output,
                                      args.overwrite, args.retry_partial)
    LOGGER.info("Validated %s: %s", target, report)


if __name__ == "__main__":
    main()