"""Classify SPEI drought/wetness levels and compute ensemble probabilities."""
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

LOGGER = logging.getLogger(__name__)
CATEGORY_MEANINGS = "D4 D3 D2 D1 D0 W1 W2 W3 W4"
CATEGORY_VALUES = np.arange(-4, 5, dtype=np.int8)
DROUGHT_TESTS = {
    "probability_d1": lambda values: values <= -0.5,
    "probability_d2": lambda values: values <= -1.0,
    "probability_d3": lambda values: values <= -1.5,
    "probability_d4": lambda values: values < -2.0,
}


def classify_spei(spei: xr.DataArray) -> xr.DataArray:
    """Return mutually exclusive D4..W4 category codes while preserving missing."""
    finite = np.isfinite(spei)
    category = xr.zeros_like(spei, dtype=np.int8)
    category = xr.where(spei >= 0.5, 1, category)
    category = xr.where(spei >= 1.0, 2, category)
    category = xr.where(spei >= 1.5, 3, category)
    category = xr.where(spei >= 2.0, 4, category)
    category = xr.where(spei <= -0.5, -1, category)
    category = xr.where(spei <= -1.0, -2, category)
    category = xr.where(spei <= -1.5, -3, category)
    category = xr.where(spei < -2.0, -4, category)
    result = category.where(finite).rename("drought_category")
    result.attrs = {
        "long_name": "mutually exclusive SPEI drought and wetness category",
        "units": "1",
        "flag_values": CATEGORY_VALUES,
        "flag_meanings": CATEGORY_MEANINGS,
        "classification": (
            "D4: SPEI<-2; D3: -2<=SPEI<=-1.5; D2: -1.5<SPEI<=-1; "
            "D1: -1<SPEI<=-0.5; D0: -0.5<SPEI<0.5; "
            "W1: 0.5<=SPEI<1; W2: 1<=SPEI<1.5; "
            "W3: 1.5<=SPEI<2; W4: SPEI>=2"
        ),
    }
    return result


def drought_probabilities(spei_components: xr.DataArray,
                          component_dimension: str = "component",
                          require_complete: bool = True) -> xr.Dataset:
    """Compute nested P(D1)..P(D4) fractions from SPEI components.

    The supplied component dimension must already contain every component that
    belongs in the denominator, including SPEI_ens when required.  Partial
    missing ensembles are rejected by default rather than silently changing
    the denominator; all-missing cells remain missing.
    """
    if component_dimension not in spei_components.dims:
        raise ValueError(f"SPEI components are missing dimension {component_dimension!r}")
    total = spei_components.sizes[component_dimension]
    if total < 1:
        raise ValueError("SPEI component dimension must not be empty")
    finite = np.isfinite(spei_components)
    valid_count = finite.sum(component_dimension)
    if require_complete:
        partial = (valid_count > 0) & (valid_count < total)
        if bool(partial.any().item()):
            raise ValueError("Partial missing SPEI ensemble found; denominator would vary")
    denominator = valid_count if not require_complete else xr.full_like(valid_count, total)
    output = {}
    for name, test in DROUGHT_TESTS.items():
        numerator = test(spei_components).where(finite, False).sum(component_dimension)
        probability = (numerator / denominator).where(valid_count > 0).rename(name)
        probability.attrs = {
            "long_name": f"ensemble probability of {name[-2:].upper()} drought threshold",
            "units": "1",
            "valid_min": 0.0,
            "valid_max": 1.0,
            "denominator": f"all {total} supplied SPEI components",
        }
        output[name] = probability
    count = valid_count.rename("valid_component_count")
    count.attrs = {"long_name": "number of finite SPEI components", "units": "1"}
    output[count.name] = count
    result = xr.Dataset(output)
    result.attrs = {
        "method": "threshold exceedance count divided by supplied component count",
        "thresholds": "D1<=-0.5; D2<=-1.0; D3<=-1.5; D4<-2.0",
        "partial_missing_policy": "rejected" if require_complete else "finite components only",
    }
    return result


def build_products(spei_ensemble: xr.DataArray, spei_components: xr.DataArray,
                   component_dimension: str = "component") -> xr.Dataset:
    """Combine deterministic category and probabilistic drought products."""
    category = classify_spei(spei_ensemble)
    probabilities = drought_probabilities(spei_components, component_dimension)
    result = xr.merge([category.to_dataset(), probabilities], compat="equals", join="exact")
    result.attrs = {
        "title": "SPEI drought category and ensemble probability products",
        "method": "confirmed README.tp thresholds and component-count probabilities",
        "software": "dr-gcm 0.1.0",
        "created_utc": datetime.now(timezone.utc).isoformat(),
    }
    return result


def validate_product(path: Path) -> dict:
    required = {"drought_category", "probability_d1", "probability_d2",
                "probability_d3", "probability_d4", "valid_component_count"}
    with xr.open_dataset(path) as dataset:
        missing = sorted(required - set(dataset.data_vars))
        if missing:
            raise ValueError(f"Drought product is missing variables: {missing}")
        for name in sorted(required - {"drought_category", "valid_component_count"}):
            values = dataset[name]
            if bool(((values < 0) | (values > 1)).any().item()):
                raise ValueError(f"{name} falls outside [0,1]")
        return {"variables": sorted(required), "sizes": dict(dataset.sizes),
                "bytes": path.stat().st_size}


def run(input_path: Path, output: Path, ensemble_variable: str = "spei_ens",
        component_variable: str = "spei_components", component_dimension: str = "component",
        overwrite: bool = False, retry_partial: bool = False) -> tuple[Path, dict]:
    if output.exists() and not overwrite:
        return output, validate_product(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)
    with xr.open_dataset(input_path) as source:
        missing = [name for name in (ensemble_variable, component_variable) if name not in source]
        if missing:
            raise KeyError(f"Input is missing variables: {missing}")
        result = build_products(source[ensemble_variable], source[component_variable], component_dimension)
        result.attrs.update({"source_file": str(input_path),
                             "ensemble_variable": ensemble_variable,
                             "component_variable": component_variable})
        encoding = {name: {"zlib": True, "complevel": 1, "shuffle": True}
                    for name in result.data_vars}
        result.to_netcdf(partial, encoding=encoding)
    report = validate_product(partial)
    os.replace(partial, output)
    return output, report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--ensemble-variable", default="spei_ens")
    parser.add_argument("--component-variable", default="spei_components")
    parser.add_argument("--component-dimension", default="component")
    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")
    target, report = run(args.input, args.output, args.ensemble_variable,
                         args.component_variable, args.component_dimension,
                         args.overwrite, args.retry_partial)
    LOGGER.info("Validated %s: %s", target, report)


if __name__ == "__main__":
    main()