from copy import deepcopy
from typing import Optional
import numpy as np
from numpy.typing import NDArray
from scipy.stats import norm
from dynasor.logging_tools import logger
from dynasor.sample import Sample


def get_spherically_averaged_sample_smearing(
        sample: Sample,
        q_norms: NDArray[float],
        q_width: float,
        use_sum: Optional[bool] = False,
        broadening: Optional[str] = 'gaussian',
) -> Sample:
    r"""
    Compute a spherical average over q-points for all the correlation functions in :attr:`sample`.

    Each q-point contributes to the function value at a given :math:`\boldsymbol{q}` with a weight
    determined by a broadening function. For example

    .. math::

        F(q) = \sum_i w(\boldsymbol{q}_i, q) F(\boldsymbol{q}_i)

    where :math:`\sum_i w(\boldsymbol{q}_i, q) = 1` (when :attr:`use_sum` is ``False``).

    Two broadening functions are available via :attr:`broadening`:

    **Gaussian** (``'gaussian'``):

    .. math::

        w(\boldsymbol{q}_i, q) \propto \exp{\left [ -\frac{1}{2} \left ( \frac{|\boldsymbol{q}_i|
        - q}{q_{width}} \right)^2 \right ]}

    where :attr:`q_width` is the standard deviation :math:`\sigma`.

    **Lorentzian** (``'lorentzian'``):

    .. math::

        w(\boldsymbol{q}_i, q) \propto \frac{1}{\left(|\boldsymbol{q}_i| - q\right)^2
        + q_{width}^2}

    where :attr:`q_width` is the half-width at half-maximum :math:`\gamma`.

    Parameters
    ----------
    sample
        Input sample.
    q_norms
        Values of :math:`|\vec{q}|` at which to evaluate the correlation functions.
    q_width
        Width of the broadening function. Standard deviation :math:`\sigma` for Gaussian;
        half-width at half-maximum :math:`\gamma` for Lorentzian.
    use_sum
        Whether to average or sum the sample in each bin.
    broadening
        Broadening function to use. Either ``'gaussian'`` (default) or ``'lorentzian'``.
    """
    if not isinstance(sample, Sample):
        raise ValueError('Input sample is not a Sample object.')

    if broadening not in ('gaussian', 'lorentzian'):
        raise ValueError(f"broadening must be 'gaussian' or 'lorentzian', got '{broadening}'")

    # get q-points
    q_points = sample.q_points
    if q_points.shape[1] != 3:
        raise ValueError('q-points array has the wrong shape.')

    # set up new input dicts for new Sample, remove q_points, add q_norms
    data_dict = dict()
    for key in sample.dimensions:
        if key == 'q_points':
            continue
        data_dict[key] = sample[key]

    if broadening == 'gaussian':
        get_average = _get_gaussian_average
        get_sum = _get_gaussian_sum
    else:
        get_average = _get_lorentzian_average
        get_sum = _get_lorentzian_sum

    for key in sample.available_correlation_functions:
        Z = getattr(sample, key)
        if use_sum:
            averaged_data = get_sum(q_points, Z, q_norms, q_width)
        else:
            averaged_data = get_average(q_points, Z, q_norms, q_width)
        data_dict[key] = averaged_data
    data_dict['q_norms'] = q_norms

    # compose new object
    new_sample = sample.__class__(
        data_dict,
        simulation_data=deepcopy(sample.simulation_data),
        history=deepcopy(sample.history))
    new_sample._append_history(
        'get_spherically_averaged_sample_smearing',
        dict(
            q_width=q_width,
            use_sum=use_sum,
            broadening=broadening,
        ))

    return new_sample


def get_spherically_averaged_sample_binned(
        sample: Sample,
        num_q_bins: int,
        use_sum: Optional[bool] = False,
) -> Sample:
    r"""
    Compute a spherical average over q-points for all the correlation functions in :attr:`sample`.

    Here, a q-binning method is used to conduct the spherical average, meaning all q-points are
    placed into spherical bins (shells).
    The corresponding function is calculated as the average of all q-points in a bin.
    If a q-bin does not contain any q-points, then its value is set to `np.nan`.
    The boundaries of the range, `q_min` and `q_max`, are taken as the minimum and maximum,
    respectively, of `|q_points|`.
    These will be set as bin centers for the first and last bins, respectively.
    The input parameter is the number of q-bins to use :attr:`num_q_bins`.

    Parameters
    ----------
    sample
        Input sample.
    num_q_bins
        Number of q-bins to use.
    use_sum
        Whether to average or sum the sample in each bin.
    """

    if not isinstance(sample, Sample):
        raise ValueError('Input sample is not a Sample object.')

    # get q-points
    q_points = sample.q_points
    if q_points.shape[1] != 3:
        raise ValueError('q-points array has wrong shape.')

    # set up new input dicts for new Sample, remove q_points, add q_norms
    data_dict = dict()
    for key in sample.dimensions:
        if key == 'q_points':
            continue
        data_dict[key] = sample[key]

    # compute spherical average for each correlation function
    for key in sample.available_correlation_functions:
        Z = getattr(sample, key)
        q_bincenters, bin_counts, averaged_data = _get_bin_average(q_points, Z, num_q_bins, use_sum)
        data_dict[key] = averaged_data
    data_dict['q_norms'] = q_bincenters

    # compose new sample
    new_sample = sample.__class__(
        data_dict,
        simulation_data=deepcopy(sample.simulation_data),
        history=deepcopy(sample.history))
    new_sample._append_history(
        'get_spherically_averaged_sample_binned',
        dict(
            num_q_bins=num_q_bins,
            use_sum=use_sum,
        ))

    return new_sample


def _get_gaussian_average(
    q_points: NDArray[float],
    Z: NDArray[float],
    q_norms: NDArray[float],
    q_width: float,
) -> NDArray[float]:
    q_norms_sample = np.linalg.norm(q_points, axis=1)
    # weights shape: (N_q_out, N_qpoints); normalization constant cancels so omit it
    diff = q_norms[:, None] - q_norms_sample[None, :]
    weights = np.exp(-0.5 * (diff / q_width) ** 2)
    norms = weights.sum(axis=1, keepdims=True)
    weights /= np.where(norms != 0, norms, 1.0)
    return weights @ Z


def _get_gaussian_sum(
    q_points: NDArray[float],
    Z: NDArray[float],
    q_norms: NDArray[float],
    q_width: float,
) -> NDArray[float]:
    q_norms_sample = np.linalg.norm(q_points, axis=1)
    diff = q_norms[:, None] - q_norms_sample[None, :]
    weights = np.exp(-0.5 * (diff / q_width) ** 2) / (q_width * np.sqrt(2 * np.pi))
    return weights @ Z


def _gaussian(x: NDArray[float], x0: float, sigma: float) -> NDArray[float]:
    dist = norm(loc=x0, scale=sigma)
    return dist.pdf(x)


def _get_lorentzian_average(
    q_points: NDArray[float],
    Z: NDArray[float],
    q_norms: NDArray[float],
    q_width: float,
) -> NDArray[float]:
    q_norms_sample = np.linalg.norm(q_points, axis=1)
    # weights shape: (N_q_out, N_qpoints); normalization constant cancels so omit it
    diff = q_norms[:, None] - q_norms_sample[None, :]
    weights = 1.0 / (diff ** 2 + q_width ** 2)
    norms = weights.sum(axis=1, keepdims=True)
    weights /= np.where(norms != 0, norms, 1.0)
    return weights @ Z


def _get_lorentzian_sum(
    q_points: NDArray[float],
    Z: NDArray[float],
    q_norms: NDArray[float],
    q_width: float,
) -> NDArray[float]:
    q_norms_sample = np.linalg.norm(q_points, axis=1)
    diff = q_norms[:, None] - q_norms_sample[None, :]
    weights = q_width / (np.pi * (diff ** 2 + q_width ** 2))
    return weights @ Z


def _lorentzian(x: NDArray[float], x0: float, gamma: float) -> NDArray[float]:
    return gamma / (np.pi * ((x - x0) ** 2 + gamma ** 2))


def _get_bin_average(
    q_points: NDArray[float],
    data: NDArray[float],
    num_q_bins: int,
    use_sum: Optional[bool] = False,
) -> tuple[NDArray[float], NDArray[float]]:
    """
    Compute a spherical average over q-points for the data using q-bins.

    If a q-bin does not contain any q-points, then a np.nan is inserted.

    q_min and q_max are determined from min/max of |q_points| and define the bin range.
    These are set as bin centers for the first and last bins, respectively.

    Parameters
    ----------
    q_points
        Array of q-points shape ``(Nq, 3)``.
    data
        Array of shape ``(Nq, N)``, shape cannot be ``(Nq, )``.
    num_q_bins
        Number of radial q-point bins to use.
    use_sum
        Whether or not to sum the data in each bin.

    Returns
    -------
        Tuple comprising the array of |q| bins of shape ``(num_q_bins, )``
        and the averaged data-array.
    """
    N_qpoints = q_points.shape[0]
    N_t = data.shape[1]
    assert q_points.shape[1] == 3
    assert data.shape[0] == N_qpoints

    # q-norms
    q_norms = np.linalg.norm(q_points, axis=1)
    assert q_norms.shape == (N_qpoints,)

    # set up bins
    q_max = np.max(q_norms)
    q_min = np.min(q_norms)
    delta_x = (q_max - q_min) / (num_q_bins - 1)
    q_range = (q_min - delta_x / 2, q_max + delta_x / 2)
    bin_counts, edges = np.histogram(q_norms, bins=num_q_bins, range=q_range)
    q_bincenters = 0.5 * (edges[1:] + edges[:-1])

    # calculate average for each bin
    averaged_data = np.zeros((num_q_bins, N_t))
    for bin_index in range(num_q_bins):
        # find q-indices that belong to this bin
        bin_min = edges[bin_index]
        bin_max = edges[bin_index + 1]
        bin_count = bin_counts[bin_index]
        q_indices = np.where(np.logical_and(q_norms >= bin_min, q_norms < bin_max))[0]
        assert len(q_indices) == bin_count
        logger.debug(f'bin {bin_index} contains {bin_count} q-points')

        # average over q-indices, if no indices then np.nan
        if bin_count == 0:
            logger.warning(f'No q-points for bin {bin_index}')
            data_bin = np.array([np.nan for _ in range(N_t)])
        else:
            if use_sum:
                data_bin = data[q_indices, :].sum(axis=0)
            else:
                data_bin = data[q_indices, :].mean(axis=0)
        averaged_data[bin_index, :] = data_bin

    return q_bincenters, bin_counts, averaged_data
