import numpy as np
import xarray as xr


def _to_180_lon_value(lon: float) -> float:
    lon180 = ((float(lon) + 180.0) % 360.0) - 180.0
    if np.isclose(lon180, -180.0) and float(lon) > 0:
        return 180.0
    return lon180


def _to_360_lon_value(lon: float) -> float:
    return float(lon) % 360.0


def to_360(da: xr.DataArray) -> xr.DataArray:
    if "lon" not in da.coords:
        raise ValueError("Input DataArray must contain 'lon' coordinate.")
    if float(da.lon.min()) < 0:
        da = da.assign_coords(lon=(da.lon % 360.0)).sortby("lon")
    return da


def to_180(da: xr.DataArray) -> xr.DataArray:
    if "lon" not in da.coords:
        raise ValueError("Input DataArray must contain 'lon' coordinate.")
    if float(da.lon.max()) > 180:
        da = da.assign_coords(lon=((da.lon + 180.0) % 360.0) - 180.0).sortby("lon")
    return da


def _slice_lat(da: xr.DataArray, lat_min: float, lat_max: float) -> xr.DataArray:
    if "lat" not in da.coords:
        raise ValueError("Input DataArray must contain 'lat' coordinate.")

    lat_vals = da["lat"].values
    if lat_vals[0] <= lat_vals[-1]:
        return da.sel(lat=slice(lat_min, lat_max))
    return da.sel(lat=slice(lat_max, lat_min))


def _slice_lon_360(da: xr.DataArray, lon_min: float, lon_max: float) -> xr.DataArray:
    """
    Slice longitudes in 0..360 convention.
    Wrapped domain if lon_min > lon_max.
    """
    out = to_360(da)
    lo = _to_360_lon_value(lon_min)
    hi = _to_360_lon_value(lon_max)

    if lo <= hi:
        return out.sel(lon=slice(lo, hi))

    # Wrapped interval: [lo, 360) U [0, hi]
    eps = 1e-6
    part1 = out.sel(lon=slice(lo, 360.0 - eps))
    part2 = out.sel(lon=slice(0.0, hi))
    return xr.concat([part1, part2], dim="lon").sortby("lon")


def _slice_lon_180(da: xr.DataArray, lon_min: float, lon_max: float) -> xr.DataArray:
    """
    Slice longitudes in -180..180 convention.
    Wrapped domain if lon_min > lon_max after conversion.
    """
    out = to_180(da)
    lo = _to_180_lon_value(lon_min)
    hi = _to_180_lon_value(lon_max)

    if lo <= hi:
        return out.sel(lon=slice(lo, hi))

    # Wrapped interval: [lo, 180] U [-180, hi]
    part1 = out.sel(lon=slice(lo, 180.0))
    part2 = out.sel(lon=slice(-180.0, hi))
    return xr.concat([part1, part2], dim="lon").sortby("lon")


def subset_latlon(
    da: xr.DataArray,
    *,
    lat_min: float,
    lat_max: float,
    lon_min: float,
    lon_max: float,
    lon_mode: str,
    name: str,
) -> xr.DataArray:
    """
    Subset a lat/lon field with portable wrapped-longitude support.

    Wrapped logic:
        if lon_min > lon_max, interpret as crossing the longitude seam.

    Examples:
        lon_mode='360', lon_min=120, lon_max=275  -> normal Pacific subset
        lon_mode='360', lon_min=120, lon_max=20   -> wrapped Pacific+Atlantic subset
        lon_mode='180', lon_min=150, lon_max=-80  -> wrapped dateline-crossing subset
    """
    if "lat" not in da.dims or "lon" not in da.dims:
        raise ValueError(f"{name} must contain lat/lon dims. Found dims={da.dims}")

    out = da.sortby("lat").sortby("lon")
    out = _slice_lat(out, lat_min, lat_max)

    if lon_mode == "360":
        out = _slice_lon_360(out, lon_min, lon_max)
    elif lon_mode == "180":
        out = _slice_lon_180(out, lon_min, lon_max)
    else:
        raise ValueError(f"lon_mode must be '180' or '360', got {lon_mode!r}")

    if out.size == 0 or out.sizes.get("lat", 0) == 0 or out.sizes.get("lon", 0) == 0:
        raise ValueError(
            f"{name} subset is empty. Requested "
            f"lat=({lat_min}, {lat_max}), lon=({lon_min}, {lon_max}), lon_mode={lon_mode}"
        )

    return out


def create_dry_mask_frequency(
    data: xr.DataArray,
    *,
    sample_dim: str = "time",
    threshold: float = 0.5,
    min_fraction: float = 2 / 3,
) -> xr.DataArray:
    """
    Identify climatologically dry grid cells.

    Returns True where:
        data < threshold in more than min_fraction of samples
    """
    if sample_dim not in data.dims:
        raise ValueError(f"{sample_dim!r} not found in data dims={data.dims}")

    dry_occurrence = data < threshold
    dry_fraction = dry_occurrence.mean(dim=sample_dim, skipna=True)
    return (dry_fraction > min_fraction).rename("dry_mask")
