# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""
Built-in distribution-creation functions.
"""

from warnings import warn

import numpy as np

from astropy import units as u

from .core import Distribution

__all__ = ["normal", "poisson", "uniform"]


def normal(
    center, *, std=None, var=None, ivar=None, n_samples, cls=Distribution, **kwargs
):
    """
    Create a Gaussian/normal distribution.

    Parameters
    ----------
    center : `~astropy.units.Quantity`
        The center of this distribution
    std : `~astropy.units.Quantity` or None
        The standard deviation/σ of this distribution. Shape must match and unit
        must be compatible with ``center``, or be `None` (if ``var`` or ``ivar``
        are set).
    var : `~astropy.units.Quantity` or None
        The variance of this distribution. Shape must match and unit must be
        compatible with ``center``, or be `None` (if ``std`` or ``ivar`` are set).
    ivar : `~astropy.units.Quantity` or None
        The inverse variance of this distribution. Shape must match and unit
        must be compatible with ``center``, or be `None` (if ``std`` or ``var``
        are set).
    n_samples : int
        The number of Monte Carlo samples to use with this distribution
    cls : class
        The class to use to create this distribution.  Typically a
        `Distribution` subclass.

    Remaining keywords are passed into the constructor of the ``cls``

    Returns
    -------
    distr : `~astropy.uncertainty.Distribution` or object
        The sampled Gaussian distribution.
        The type will be the same as the parameter ``cls``.

    """
    center = np.asanyarray(center)
    if var is not None:
        if std is None:
            std = np.asanyarray(var) ** 0.5
        else:
            raise ValueError("normal cannot take both std and var")
    if ivar is not None:
        if std is None:
            std = np.asanyarray(ivar) ** -0.5
        else:
            raise ValueError("normal cannot take both ivar and and std or var")
    if std is None:
        raise ValueError("normal requires one of std, var, or ivar")
    else:
        std = np.asanyarray(std)

    randshape = np.broadcast(std, center).shape + (n_samples,)
    samples = (
        center[..., np.newaxis] + np.random.randn(*randshape) * std[..., np.newaxis]
    )
    return cls(samples, **kwargs)


COUNT_UNITS = (
    u.count,
    u.electron,
    u.dimensionless_unscaled,
    u.chan,
    u.bin,
    u.vox,
    u.bit,
    u.byte,
)


def poisson(center, n_samples, cls=Distribution, **kwargs):
    """
    Create a Poisson distribution.

    Parameters
    ----------
    center : `~astropy.units.Quantity`
        The center value of this distribution (i.e., λ).
    n_samples : int
        The number of Monte Carlo samples to use with this distribution
    cls : class
        The class to use to create this distribution.  Typically a
        `Distribution` subclass.

    Remaining keywords are passed into the constructor of the ``cls``

    Returns
    -------
    distr : `~astropy.uncertainty.Distribution` or object
        The sampled Poisson distribution.
        The type will be the same as the parameter ``cls``.
    """
    # we convert to arrays because np.random.poisson has trouble with quantities
    has_unit = False
    if hasattr(center, "unit"):
        has_unit = True
        poissonarr = np.asanyarray(center.value)
    else:
        poissonarr = np.asanyarray(center)
    randshape = poissonarr.shape + (n_samples,)

    samples = np.random.poisson(poissonarr[..., np.newaxis], randshape)
    if has_unit:
        if center.unit == u.adu:
            warn(
                "ADUs were provided to poisson.  ADUs are not strictly count"
                "units because they need the gain to be applied. It is "
                "recommended you apply the gain to convert to e.g. electrons."
            )
        elif center.unit not in COUNT_UNITS:
            warn(
                f"Unit {center.unit} was provided to poisson, which is not one of"
                f' {COUNT_UNITS}, and therefore suspect as a "counting" unit.  Ensure'
                " you mean to use Poisson statistics."
            )

        # re-attach the unit
        samples = samples * center.unit

    return cls(samples, **kwargs)


def uniform(
    *,
    lower=None,
    upper=None,
    center=None,
    width=None,
    n_samples,
    cls=Distribution,
    **kwargs,
):
    """
    Create a Uniform distriution from the lower and upper bounds.

    Note that this function requires keywords to be explicit, and requires
    either ``lower``/``upper`` or ``center``/``width``.

    Parameters
    ----------
    lower : array-like
        The lower edge of this distribution. If a `~astropy.units.Quantity`, the
        distribution will have the same units as ``lower``.
    upper : `~astropy.units.Quantity`
        The upper edge of this distribution. Must match shape and if a
        `~astropy.units.Quantity` must have compatible units with ``lower``.
    center : array-like
        The center value of the distribution. Cannot be provided at the same
        time as ``lower``/``upper``.
    width : array-like
        The width of the distribution.  Must have the same shape and compatible
        units with ``center`` (if any).
    n_samples : int
        The number of Monte Carlo samples to use with this distribution
    cls : class
        The class to use to create this distribution.  Typically a
        `Distribution` subclass.

    Remaining keywords are passed into the constructor of the ``cls``

    Returns
    -------
    distr : `~astropy.uncertainty.Distribution` or object
        The sampled uniform distribution.
        The type will be the same as the parameter ``cls``.
    """
    if center is None and width is None:
        lower = np.asanyarray(lower)
        upper = np.asanyarray(upper)
        if lower.shape != upper.shape:
            raise ValueError("lower and upper must have consistent shapes")
    elif upper is None and lower is None:
        center = np.asanyarray(center)
        width = np.asanyarray(width)
        lower = center - width / 2
        upper = center + width / 2
    else:
        raise ValueError(
            "either upper/lower or center/width must be given "
            "to uniform - other combinations are not valid"
        )

    newshape = lower.shape + (n_samples,)
    if lower.shape == tuple() and upper.shape == tuple():
        width = upper - lower  # scalar
    else:
        width = (upper - lower)[:, np.newaxis]
        lower = lower[:, np.newaxis]
    samples = lower + width * np.random.uniform(size=newshape)

    return cls(samples, **kwargs)