from __future__ import annotations

from pathlib import Path
from typing import Any, Iterable
import json
import numpy as np
import xarray as xr

DEFAULT_MISSING_VALUE = -99.0


def as_float_array(x: Any) -> np.ndarray:
    return np.asarray(x, dtype=float)


def invalid_to_nan(x: Any, missing_value: float = DEFAULT_MISSING_VALUE, *, atol: float = 1e-6) -> np.ndarray:
    """Return float array with missing sentinel and inf converted to NaN.

    This is only the first validity layer.  Physical and mathematical domain
    checks must still be applied by the calling function.
    """
    a = as_float_array(x).copy()
    a[~np.isfinite(a)] = np.nan
    if missing_value is not None:
        a[np.isclose(a, float(missing_value), atol=atol, rtol=0.0)] = np.nan
    return a


def nan_to_missing(x: Any, missing_value: float = DEFAULT_MISSING_VALUE, dtype=np.float32) -> np.ndarray:
    a = np.asarray(x, dtype=float)
    out = np.where(np.isfinite(a), a, float(missing_value))
    return out.astype(dtype)


def safe_mean(x: Any, axis=None, min_count: int = 1, dtype=np.float32) -> np.ndarray:
    """Mean ignoring NaN without RuntimeWarning on all-missing slices.

    Returns NaN where the number of finite values is below min_count.
    """
    a = as_float_array(x)
    valid = np.isfinite(a)
    count = np.sum(valid, axis=axis)
    total = np.sum(np.where(valid, a, 0.0), axis=axis)
    out = np.full(np.shape(total), np.nan, dtype=float)
    ok = count >= int(min_count)
    np.divide(total, count, out=out, where=ok)
    if np.ndim(out) == 0:
        return np.array(out, dtype=dtype).item()
    return out.astype(dtype)


def safe_std(x: Any, axis=None, min_count: int = 2, dtype=np.float32) -> np.ndarray:
    a = as_float_array(x)
    mean = safe_mean(a, axis=axis, min_count=1, dtype=float)
    if axis is None:
        valid = np.isfinite(a)
        if np.sum(valid) < min_count:
            return np.array(np.nan, dtype=dtype).item()
        return np.array(np.sqrt(np.mean((a[valid] - mean) ** 2)), dtype=dtype).item()
    expanded = np.expand_dims(mean, axis=axis)
    valid = np.isfinite(a) & np.isfinite(expanded)
    count = np.sum(valid, axis=axis)
    ss = np.sum(np.where(valid, (a - expanded) ** 2, 0.0), axis=axis)
    out = np.full(np.shape(ss), np.nan, dtype=float)
    ok = count >= int(min_count)
    np.divide(ss, count, out=out, where=ok)
    out = np.sqrt(out)
    return out.astype(dtype)


def safe_divide(num: Any, den: Any, *, min_den: float = 1e-12, dtype=np.float32) -> np.ndarray:
    n = as_float_array(num)
    d = as_float_array(den)
    out = np.full(np.broadcast_shapes(n.shape, d.shape), np.nan, dtype=float)
    n2, d2 = np.broadcast_arrays(n, d)
    ok = np.isfinite(n2) & np.isfinite(d2) & (np.abs(d2) > float(min_den))
    np.divide(n2, d2, out=out, where=ok)
    return out.astype(dtype)


def valid_temperature_c(t: Any, *, lower: float = -80.0, upper: float = 60.0, missing_value: float = DEFAULT_MISSING_VALUE) -> np.ndarray:
    a = invalid_to_nan(t, missing_value=missing_value)
    return np.where(np.isfinite(a) & (a >= lower) & (a <= upper), a, np.nan)


def valid_precip_mm(p: Any, *, max_reasonable: float = 10000.0, missing_value: float = DEFAULT_MISSING_VALUE) -> np.ndarray:
    a = invalid_to_nan(p, missing_value=missing_value)
    return np.where(np.isfinite(a) & (a >= 0.0) & (a <= max_reasonable), a, np.nan)


def finite_fraction(x: Any) -> float:
    a = np.asarray(x)
    if a.size == 0:
        return float('nan')
    return float(np.isfinite(a.astype(float)).sum() / a.size)


def write_status(path: str | Path, **kwargs: Any) -> None:
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    data = dict(kwargs)
    p.write_text(json.dumps(data, indent=2, ensure_ascii=False, default=str), encoding='utf-8')


def write_dataset_netcdf(ds: xr.Dataset, path: str | Path, *, missing_value: float = DEFAULT_MISSING_VALUE) -> None:
    """Write Dataset with a consistent floating missing value.

    Internal NaNs remain NaNs in memory.  NetCDF encoding advertises -99 as
    missing value; readers can decode to NaN. Integer variables are left alone
    except their attrs are annotated when they already contain -99 codes.
    """
    encoding = {}
    for name, da in ds.data_vars.items():
        if np.issubdtype(da.dtype, np.floating):
            da.attrs.pop('_FillValue', None)
            da.attrs['missing_value'] = float(missing_value)
            encoding[name] = {'_FillValue': float(missing_value)}
        elif np.issubdtype(da.dtype, np.integer):
            da.attrs['missing_value'] = int(missing_value)
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    ds.to_netcdf(path, encoding=encoding)