"""Fitting operations for DataArrays and Datasets."""

from __future__ import annotations

import inspect
import warnings
from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence
from inspect import Parameter
from types import MappingProxyType
from typing import (
    Any,
    Literal,
    Union,
)

import numpy as np

# remove once numpy 2.0 is the oldest supported version
try:
    from numpy.exceptions import RankWarning
except ImportError:
    from numpy import RankWarning  # type: ignore[no-redef,attr-defined,unused-ignore]

from xarray.computation.apply_ufunc import apply_ufunc
from xarray.computation.computation import _ensure_numeric, where
from xarray.core.dataarray import DataArray
from xarray.core.duck_array_ops import is_duck_dask_array, least_squares
from xarray.core.types import Dims, ErrorOptions
from xarray.core.variable import Variable
from xarray.structure.alignment import broadcast


def _get_func_args(func, param_names):
    """Use `inspect.signature` to try accessing `func` args. Otherwise, ensure
    they are provided by user.
    """
    func_args: Union[dict[str, Parameter], MappingProxyType[str, Parameter]]
    try:
        func_args = inspect.signature(func).parameters
    except ValueError as err:
        func_args = {}  # type: ignore[assignment,unused-ignore]
        if not param_names:
            raise ValueError(
                "Unable to inspect `func` signature, and `param_names` was not provided."
            ) from err
    if param_names:
        params = param_names
    else:
        params = list(func_args)[1:]
        if any(
            (p.kind in [p.VAR_POSITIONAL, p.VAR_KEYWORD]) for p in func_args.values()
        ):
            raise ValueError(
                "`param_names` must be provided because `func` takes variable length arguments."
            )
    return params, func_args


def _initialize_curvefit_params(params, p0, bounds, func_args):
    """Set initial guess and bounds for curvefit.
    Priority: 1) passed args 2) func signature 3) scipy defaults
    """

    def _initialize_feasible(lb, ub):
        # Mimics functionality of scipy.optimize.minpack._initialize_feasible
        lb_finite = np.isfinite(lb)
        ub_finite = np.isfinite(ub)
        p0 = where(
            lb_finite,
            where(
                ub_finite,
                0.5 * (lb + ub),  # both bounds finite
                lb + 1,  # lower bound finite, upper infinite
            ),
            where(
                ub_finite,
                ub - 1,  # lower bound infinite, upper finite
                0,  # both bounds infinite
            ),
        )
        return p0

    param_defaults = dict.fromkeys(params, 1)
    bounds_defaults = dict.fromkeys(params, (-np.inf, np.inf))
    for p in params:
        if p in func_args and func_args[p].default is not func_args[p].empty:
            param_defaults[p] = func_args[p].default
        if p in bounds:
            lb, ub = bounds[p]
            bounds_defaults[p] = (lb, ub)
            param_defaults[p] = where(
                (param_defaults[p] < lb) | (param_defaults[p] > ub),
                _initialize_feasible(lb, ub),
                param_defaults[p],
            )
        if p in p0:
            param_defaults[p] = p0[p]
    return param_defaults, bounds_defaults


def polyfit(
    obj,
    dim: Hashable,
    deg: int,
    skipna: bool | None = None,
    rcond: np.floating[Any] | float | None = None,
    w: Hashable | Any = None,
    full: bool = False,
    cov: bool | Literal["unscaled"] = False,
):
    """
    Least squares polynomial fit.

    This replicates the behaviour of `numpy.polyfit` but differs by skipping
    invalid values when `skipna = True`.

    Parameters
    ----------
    obj : Dataset or DataArray
        Object to perform the polyfit on
    dim : hashable
        Coordinate along which to fit the polynomials.
    deg : int
        Degree of the fitting polynomial.
    skipna : bool or None, optional
        If True, removes all invalid values before fitting each 1D slices of the array.
        Default is True if data is stored in a dask.array or if there is any
        invalid values, False otherwise.
    rcond : float or None, optional
        Relative condition number to the fit.
    w : hashable or Any, optional
        Weights to apply to the y-coordinate of the sample points.
        Can be an array-like object or the name of a coordinate in the dataset.
    full : bool, default: False
        Whether to return the residuals, matrix rank and singular values in addition
        to the coefficients.
    cov : bool or "unscaled", default: False
        Whether to return to the covariance matrix in addition to the coefficients.
        The matrix is not scaled if `cov='unscaled'`.

    Returns
    -------
    Dataset
        A single dataset which contains (for each "var" in the input dataset):

        [var]_polyfit_coefficients
            The coefficients of the best fit for each variable in this dataset.
        [var]_polyfit_residuals
            The residuals of the least-square computation for each variable (only included if `full=True`)
            When the matrix rank is deficient, np.nan is returned.
        [dim]_matrix_rank
            The effective rank of the scaled Vandermonde coefficient matrix (only included if `full=True`)
            The rank is computed ignoring the NaN values that might be skipped.
        [dim]_singular_values
            The singular values of the scaled Vandermonde coefficient matrix (only included if `full=True`)
        [var]_polyfit_covariance
            The covariance matrix of the polynomial coefficient estimates (only included if `full=False` and `cov=True`)

    Warns
    -----
    RankWarning
        The rank of the coefficient matrix in the least-squares fit is deficient.
        The warning is not raised with in-memory (not dask) data and `full=True`.

    See Also
    --------
    numpy.polyfit
    numpy.polyval
    xarray.polyval
    """
    variables: dict[Hashable, Variable] = {}
    skipna_da = skipna

    x = np.asarray(_ensure_numeric(obj.coords[dim]).astype(np.float64))

    xname = f"{obj[dim].name}_"
    order = int(deg) + 1
    degree_coord_values = np.arange(order)[::-1]
    lhs = np.vander(x, order)

    if rcond is None:
        rcond = x.shape[0] * np.finfo(x.dtype).eps

    # Weights:
    if w is not None:
        if isinstance(w, Hashable):
            w = obj.coords[w]
        w = np.asarray(w)
        if w.ndim != 1:
            raise TypeError("Expected a 1-d array for weights.")
        if w.shape[0] != lhs.shape[0]:
            raise TypeError(f"Expected w and {dim} to have the same length")
        lhs *= w[:, np.newaxis]

    # Scaling
    scale = np.sqrt((lhs * lhs).sum(axis=0))
    lhs /= scale

    from xarray.core import utils

    degree_dim = utils.get_temp_dimname(obj.dims, "degree")

    rank = np.linalg.matrix_rank(lhs)

    if full:
        rank = Variable(dims=(), data=rank)
        variables[xname + "matrix_rank"] = rank
        _sing = np.linalg.svd(lhs, compute_uv=False)
        variables[xname + "singular_values"] = Variable(
            dims=(degree_dim,),
            data=np.concatenate([np.full((order - rank.data,), np.nan), _sing]),
        )

    # If we have a coordinate get its underlying dimension.
    (true_dim,) = obj.coords[dim].dims

    other_coords = {
        dim: obj._variables[dim]
        for dim in set(obj.dims) - {true_dim}
        if dim in obj._variables
    }
    present_dims: set[Hashable] = set()
    for name, var in obj._variables.items():
        if name in obj._coord_names or name in obj.dims:
            continue
        if true_dim not in var.dims:
            continue

        if is_duck_dask_array(var._data) and (rank != order or full or skipna is None):
            # Current algorithm with dask and skipna=False neither supports
            # deficient ranks nor does it output the "full" info (issue dask/dask#6516)
            skipna_da = True
        elif skipna is None:
            skipna_da = bool(np.any(var.isnull()))

        if var.ndim > 1:
            rhs = var.transpose(true_dim, ...)
            other_dims = rhs.dims[1:]
            scale_da = scale.reshape(-1, *((1,) * len(other_dims)))
        else:
            rhs = var
            scale_da = scale
            other_dims = ()

        present_dims.update(other_dims)
        if w is not None:
            rhs = rhs * w.reshape(-1, *((1,) * len(other_dims)))

        with warnings.catch_warnings():
            if full:  # Copy np.polyfit behavior
                warnings.simplefilter("ignore", RankWarning)
            else:  # Raise only once per variable
                warnings.simplefilter("once", RankWarning)

            coeffs, residuals = least_squares(
                lhs, rhs.data, rcond=rcond, skipna=skipna_da
            )

        from xarray.core.dataarray import _THIS_ARRAY

        if name is _THIS_ARRAY:
            # When polyfit is called on a DataArray, ensure the resulting
            # dataset is backwards compatible with previous behavior
            name = ""
        elif isinstance(name, str):
            name = f"{name}_"
        else:
            # For other non-string names
            name = ""

        variables[name + "polyfit_coefficients"] = Variable(
            data=coeffs / scale_da, dims=(degree_dim,) + other_dims
        )

        if full or (cov is True):
            variables[name + "polyfit_residuals"] = Variable(
                data=residuals if var.ndim > 1 else residuals.squeeze(),
                dims=other_dims,
            )

        fac: Variable | int
        if cov:
            Vbase = np.linalg.inv(np.dot(lhs.T, lhs))
            Vbase /= np.outer(scale, scale)
            if cov == "unscaled":
                fac = 1
            else:
                if x.shape[0] <= order:
                    raise ValueError(
                        "The number of data points must exceed order to scale the covariance matrix."
                    )
                fac = variables[name + "polyfit_residuals"] / (x.shape[0] - order)
            variables[name + "polyfit_covariance"] = (
                Variable(data=Vbase, dims=("cov_i", "cov_j")) * fac
            )

    return type(obj)(
        data_vars=variables,
        coords={
            degree_dim: degree_coord_values,
            **{
                name: coord
                for name, coord in other_coords.items()
                if name in present_dims
            },
        },
        attrs=obj.attrs.copy(),
    )


def curvefit(
    obj,
    coords: str | DataArray | Iterable[str | DataArray],
    func: Callable[..., Any],
    reduce_dims: Dims = None,
    skipna: bool = True,
    p0: Mapping[str, float | DataArray] | None = None,
    bounds: Mapping[str, tuple[float | DataArray, float | DataArray]] | None = None,
    param_names: Sequence[str] | None = None,
    errors: ErrorOptions = "raise",
    kwargs: dict[str, Any] | None = None,
):
    """
    Curve fitting optimization for arbitrary functions.

    Wraps `scipy.optimize.curve_fit` with `apply_ufunc`.

    Parameters
    ----------
    obj : Dataset or DataArray
        Object to perform the curvefit on
    coords : hashable, DataArray, or sequence of hashable or DataArray
        Independent coordinate(s) over which to perform the curve fitting. Must share
        at least one dimension with the calling object. When fitting multi-dimensional
        functions, supply `coords` as a sequence in the same order as arguments in
        `func`. To fit along existing dimensions of the calling object, `coords` can
        also be specified as a str or sequence of strs.
    func : callable
        User specified function in the form `f(x, *params)` which returns a numpy
        array of length `len(x)`. `params` are the fittable parameters which are optimized
        by scipy curve_fit. `x` can also be specified as a sequence containing multiple
        coordinates, e.g. `f((x0, x1), *params)`.
    reduce_dims : str, Iterable of Hashable or None, optional
        Additional dimension(s) over which to aggregate while fitting. For example,
        calling `ds.curvefit(coords='time', reduce_dims=['lat', 'lon'], ...)` will
        aggregate all lat and lon points and fit the specified function along the
        time dimension.
    skipna : bool, default: True
        Whether to skip missing values when fitting. Default is True.
    p0 : dict-like, optional
        Optional dictionary of parameter names to initial guesses passed to the
        `curve_fit` `p0` arg. If the values are DataArrays, they will be appropriately
        broadcast to the coordinates of the array. If none or only some parameters are
        passed, the rest will be assigned initial values following the default scipy
        behavior.
    bounds : dict-like, optional
        Optional dictionary of parameter names to tuples of bounding values passed to the
        `curve_fit` `bounds` arg. If any of the bounds are DataArrays, they will be
        appropriately broadcast to the coordinates of the array. If none or only some
        parameters are passed, the rest will be unbounded following the default scipy
        behavior.
    param_names : sequence of hashable, optional
        Sequence of names for the fittable parameters of `func`. If not supplied,
        this will be automatically determined by arguments of `func`. `param_names`
        should be manually supplied when fitting a function that takes a variable
        number of parameters.
    errors : {"raise", "ignore"}, default: "raise"
        If 'raise', any errors from the `scipy.optimize_curve_fit` optimization will
        raise an exception. If 'ignore', the coefficients and covariances for the
        coordinates where the fitting failed will be NaN.
    kwargs : optional
        Additional keyword arguments to passed to scipy curve_fit.

    Returns
    -------
    Dataset
        A single dataset which contains:

        [var]_curvefit_coefficients
            The coefficients of the best fit.
        [var]_curvefit_covariance
            The covariance matrix of the coefficient estimates.

    See Also
    --------
    Dataset.polyfit
    scipy.optimize.curve_fit
    """
    from scipy.optimize import curve_fit

    if p0 is None:
        p0 = {}
    if bounds is None:
        bounds = {}
    if kwargs is None:
        kwargs = {}

    reduce_dims_: list[Hashable]
    if not reduce_dims:
        reduce_dims_ = []
    elif isinstance(reduce_dims, str) or not isinstance(reduce_dims, Iterable):
        reduce_dims_ = [reduce_dims]
    else:
        reduce_dims_ = list(reduce_dims)

    if isinstance(coords, str | DataArray) or not isinstance(coords, Iterable):
        coords = [coords]
    coords_: Sequence[DataArray] = [
        obj[coord] if isinstance(coord, str) else coord for coord in coords
    ]

    # Determine whether any coords are dims on self
    for coord in coords_:
        reduce_dims_ += [c for c in obj.dims if coord.equals(obj[c])]
    reduce_dims_ = list(set(reduce_dims_))
    preserved_dims = list(set(obj.dims) - set(reduce_dims_))
    if not reduce_dims_:
        raise ValueError(
            "No arguments to `coords` were identified as a dimension on the calling "
            "object, and no dims were supplied to `reduce_dims`. This would result "
            "in fitting on scalar data."
        )

    # Check that initial guess and bounds only contain coordinates that are in preserved_dims
    for param, guess in p0.items():
        if isinstance(guess, DataArray):
            unexpected = set(guess.dims) - set(preserved_dims)
            if unexpected:
                raise ValueError(
                    f"Initial guess for '{param}' has unexpected dimensions "
                    f"{tuple(unexpected)}. It should only have dimensions that are in data "
                    f"dimensions {preserved_dims}."
                )
    for param, (lb, ub) in bounds.items():
        for label, bound in zip(("Lower", "Upper"), (lb, ub), strict=True):
            if isinstance(bound, DataArray):
                unexpected = set(bound.dims) - set(preserved_dims)
                if unexpected:
                    raise ValueError(
                        f"{label} bound for '{param}' has unexpected dimensions "
                        f"{tuple(unexpected)}. It should only have dimensions that are in data "
                        f"dimensions {preserved_dims}."
                    )

    if errors not in ["raise", "ignore"]:
        raise ValueError('errors must be either "raise" or "ignore"')

    # Broadcast all coords with each other
    coords_ = broadcast(*coords_)
    coords_ = [coord.broadcast_like(obj, exclude=preserved_dims) for coord in coords_]
    n_coords = len(coords_)

    params, func_args = _get_func_args(func, param_names)
    param_defaults, bounds_defaults = _initialize_curvefit_params(
        params, p0, bounds, func_args
    )
    n_params = len(params)

    def _wrapper(Y, *args, **kwargs):
        # Wrap curve_fit with raveled coordinates and pointwise NaN handling
        # *args contains:
        #   - the coordinates
        #   - initial guess
        #   - lower bounds
        #   - upper bounds
        coords__ = args[:n_coords]
        p0_ = args[n_coords + 0 * n_params : n_coords + 1 * n_params]
        lb = args[n_coords + 1 * n_params : n_coords + 2 * n_params]
        ub = args[n_coords + 2 * n_params :]

        x = np.vstack([c.ravel() for c in coords__])
        y = Y.ravel()
        if skipna:
            mask = np.all([np.any(~np.isnan(x), axis=0), ~np.isnan(y)], axis=0)
            x = x[:, mask]
            y = y[mask]
            if y.size == 0:
                popt = np.full([n_params], np.nan)
                pcov = np.full([n_params, n_params], np.nan)
                return popt, pcov
        x = np.squeeze(x)

        try:
            popt, pcov = curve_fit(func, x, y, p0=p0_, bounds=(lb, ub), **kwargs)
        except RuntimeError:
            if errors == "raise":
                raise
            popt = np.full([n_params], np.nan)
            pcov = np.full([n_params, n_params], np.nan)

        return popt, pcov

    from xarray.core.dataarray import _THIS_ARRAY

    result = type(obj)()
    for name, da in obj.data_vars.items():
        if name is _THIS_ARRAY:
            # When curvefit is called on a DataArray, ensure the resulting
            # dataset is backwards compatible with previous behavior
            var_name = ""
        else:
            var_name = f"{name}_"

        input_core_dims = [reduce_dims_ for _ in range(n_coords + 1)]
        input_core_dims.extend(
            [[] for _ in range(3 * n_params)]
        )  # core_dims for p0 and bounds

        popt, pcov = apply_ufunc(
            _wrapper,
            da,
            *coords_,
            *param_defaults.values(),
            *[b[0] for b in bounds_defaults.values()],
            *[b[1] for b in bounds_defaults.values()],
            vectorize=True,
            dask="parallelized",
            input_core_dims=input_core_dims,
            output_core_dims=[["param"], ["cov_i", "cov_j"]],
            dask_gufunc_kwargs={
                "output_sizes": {
                    "param": n_params,
                    "cov_i": n_params,
                    "cov_j": n_params,
                },
            },
            output_dtypes=(np.float64, np.float64),
            exclude_dims=set(reduce_dims_),
            kwargs=kwargs,
        )
        result[var_name + "curvefit_coefficients"] = popt
        result[var_name + "curvefit_covariance"] = pcov

    result = result.assign_coords({"param": params, "cov_i": params, "cov_j": params})
    result.attrs = obj.attrs.copy()

    return result