"""Utility functions: Gaussian PDF helpers, Gaussian fitting, circular/angular helpers, and simulated data generation."""

import jax
import jax.numpy as jnp

_TAU = 2 * jnp.pi

__all__ = [
    "gaussian_pdf",
    "log_gaussian_pdf",
    "gaussian_norm_const",
    "fit_gaussian",
    "fit_gaussian_vmap",
    "make_simulated_dataset",
]


def _wrap_minuspi_pi(theta: jnp.ndarray) -> jnp.ndarray:
    """Wrap angles to [-pi, pi).

    Parameters
    ----------
    theta : jnp.ndarray
        Angles in radians (any range)

    Returns
    -------
    jnp.ndarray
        Angles wrapped to [-pi, pi)
    """
    return jnp.mod(theta + jnp.pi, _TAU) - jnp.pi


def _bin_indices_minuspi_pi(theta: jnp.ndarray, n_bins: int) -> jnp.ndarray:
    """Map theta in radians to integer bin indices [0, n_bins).

    Maps angles to bin indices where bin 0 corresponds to [-pi, -pi + delta).

    Parameters
    ----------
    theta : jnp.ndarray
        Angles in radians (any range)
    n_bins : int
        Number of bins

    Returns
    -------
    jnp.ndarray
        Integer bin indices in [0, n_bins)
    """
    theta = _wrap_minuspi_pi(theta)
    u = (theta + jnp.pi) * (n_bins / _TAU)  # in [0, n_bins)
    idx = jnp.floor(u).astype(jnp.int32)
    # guard against theta == pi mapping to n_bins (shouldn't happen for [-pi,pi) but safe)
    return jnp.clip(idx, 0, n_bins - 1)


def _circular_conv_fft_1d(x: jnp.ndarray, k: jnp.ndarray) -> jnp.ndarray:
    """Circular convolution via FFT for 1D arrays.

    Parameters
    ----------
    x : jnp.ndarray
        Input array of length N
    k : jnp.ndarray
        Kernel array of length N

    Returns
    -------
    jnp.ndarray
        Circular convolution of x and k, same length as input
    """
    return jnp.fft.ifft(jnp.fft.fft(x) * jnp.fft.fft(k)).real


def gaussian_pdf(
    x: jnp.ndarray,
    mu: jnp.ndarray,
    sigma: jnp.ndarray,
) -> jnp.ndarray:
    """Calculates the gaussian pdf of a multivariate normal distribution of mean mu and covariance sigma at x

    Parameters
    ----------
    x: (D,) array
        The position at which to evaluate the pdf
    mu: (D,) array
        The mean of the distribution
    sigma: (D, D) array
        The covariance of the distribution

    Returns
    -------
    pdf: float
        The probability density at x
    """
    assert x.ndim == 1
    assert mu.ndim == 1
    assert sigma.ndim == 2
    assert x.shape[0] == mu.shape[0]
    assert x.shape[0] == sigma.shape[0]
    assert sigma.shape[0] == sigma.shape[1]

    x = x - mu
    norm_const = gaussian_norm_const(sigma)
    return norm_const * jnp.exp(-0.5 * jnp.sum(x @ jnp.linalg.inv(sigma) * x, axis=-1))


def log_gaussian_pdf(
    x: jnp.ndarray,
    mu: jnp.ndarray,
    sigma: jnp.ndarray,
) -> jnp.ndarray:
    """Calculates the log of the gaussian pdf of a multivariate normal distribution of mean mu and covariance sigma at x

    Parameters
    ----------
    x: (D,) array
        The position at which to evaluate the pdf
    mu: (D,) array
        The mean of the distribution
    sigma: (D, D) array
        The covariance of the distribution

    Returns
    -------
    log_pdf: float
        The log probability density at x
    """
    assert x.ndim == 1
    assert mu.ndim == 1
    assert sigma.ndim == 2
    assert x.shape[0] == mu.shape[0]
    assert x.shape[0] == sigma.shape[0]
    assert sigma.shape[0] == sigma.shape[1]

    x = x - mu
    norm_const = gaussian_norm_const(sigma)
    return jnp.log(norm_const) - 0.5 * jnp.sum(x @ jnp.linalg.inv(sigma) * x)


def gaussian_norm_const(sigma: jnp.ndarray) -> jnp.ndarray:
    """Calculates the normalizing constant of a multivariate normal distribution with covariance sigma

    Parameters
    ----------
    sigma: jnp.ndarray, shape (D, D)
        The covariance matrix of the distribution

    Returns
    -------
    norm_const: jnp.ndarray, shape (1,)
        The normalizing constant
    """
    assert sigma.ndim == 2
    D = sigma.shape[0]
    return 1 / jnp.sqrt((2 * jnp.pi) ** D * jnp.linalg.det(sigma))


def fit_gaussian(x, likelihood):
    """Fits a multivariate-Gaussian to the likelihood function P(spikes | x) in x-space.

    Parameters
    ----------
    x : jnp.ndarray, shape (N_bins,D)
        The position bins in which the likelihood is calculated
    likelihood : jnp.ndarray, shape (N_bins,)
        The combined likelihood (not log-likelihood) of the neurons firing at each position bin

    Returns
    -------
    mu : jnp.ndarray, shape (D,)
        The mean of the Gaussian
    mode : jnp.ndarray, shape (D,)
        The mode of the Gaussian
    covariance : jnp.ndarray, shape (D, D)
        The covariance of the Gaussian
    """
    assert x.ndim == 2
    assert likelihood.ndim == 1
    assert x.shape[0] == likelihood.shape[0]

    mu = (x.T @ likelihood) / likelihood.sum()
    mode = x[jnp.argmax(likelihood)]
    covariance = ((x - mu) * likelihood[:, None]).T @ (x - mu) / likelihood.sum()
    return mu, mode, covariance


def fit_gaussian_vmap(x, likelihoods):
    """Fits a multivariate-Gaussian to each row of a batch of likelihood arrays.

    This is the vmapped version of ``fit_gaussian``: it accepts likelihoods of
    shape ``(T, N_bins)`` and returns batched means, modes, and covariances.

    Parameters
    ----------
    x : jnp.ndarray, shape (N_bins, D)
        The position bins (shared across all time steps).
    likelihoods : jnp.ndarray, shape (T, N_bins)
        Likelihood values at each bin for each time step.

    Returns
    -------
    means : jnp.ndarray, shape (T, D)
        The mean of each fitted Gaussian.
    modes : jnp.ndarray, shape (T, D)
        The mode of each fitted Gaussian.
    covariances : jnp.ndarray, shape (T, D, D)
        The covariance of each fitted Gaussian.
    """
    return jax.vmap(fit_gaussian, in_axes=(None, 0))(x, likelihoods)


def make_simulated_dataset(time_mins=60, n_cells=100, firing_rate=10, random_seed=None, **kwargs):
    """Makes a simulated dataset for an agent randomly foraging a 1 m square box. Data generated with the RatInABox package and defaults to place cells. Returns the data as jax arrays.

    Parameters
    ----------
    time_mins: int
        The number of minutes to simulate the agent for
    n_cells : int
        The number of place cells to simulate. Default is 100.
    firing_rate : float
        The maximum firing rate of each place cell in Hz. Default is 10.
    random_seed : int, optional
        If provided, sets the numpy random seed for reproducibility.
    kwargs: dict
        Additional arguments to pass to the RatInABox simulation

    Returns
    -------
    time: jnp.ndarray, shape (N,)
        The time points of the simulation
    position: jnp.ndarray, shape (N, dims)
        The position of the agent at each time point
    spikes: jnp.ndarray, shape (N, N_cells)
        The spikes of the place cells at each time point
    """
    import tqdm as tqdm

    from ratinabox.Environment import Environment
    from ratinabox.Agent import Agent
    from ratinabox.Neurons import PlaceCells

    if random_seed is not None:
        import numpy as np

        np.random.seed(random_seed)

    env_params = kwargs.get("env_params", {})
    agent_params = kwargs.get("agent_params", {"dt": 0.1})
    place_cell_params = kwargs.get("place_cell_params", {"n": n_cells, "max_fr": firing_rate, "widths": 0.1})

    env = Environment(params=env_params)
    agent = Agent(env, params=agent_params)
    place_cells = PlaceCells(agent, params=place_cell_params)

    for i in tqdm.tqdm(range(int(60 * time_mins / agent.dt))):
        agent.update()
        place_cells.update()

    time = jnp.array(agent.history["t"])
    position = jnp.array(agent.history["pos"])
    spikes = jnp.array(place_cells.history["spikes"])

    return time, position, spikes
