from typing import Callable, Iterator, Tuple

import numpy as np

from stats.averageOverLogLog import AverageOverLogLog  # type: ignore
from stats.pdf import MakeLogPdf  # type: ignore
from stats.psd import MakeLogPsd  # type: ignore
from taudiff import event_series_generator, get_tau_series


def simulate_tau_pdf(
    model: Callable[..., Iterator[float]],
    *,
    sigma: float = 1e-4,
    hurst: float = 0.5,
    boundary: float = 1e-4,
    n_tau_points: int = 67108864,
    fold_mid: bool = False,
    seed: int = -1,
) -> np.ndarray:
    """Obtain PDF of inter-event times (in event time space).

    Input:
        model:
            Generator to use when generating the series. Use generators
            from the `taudiff` module or create your own, which would
            replace the interface.
        sigma: (default: 1e-4)
            Sigma parameter of the iterative equation driving the
            process. Sigma parameter controls how fast the inter-event
            time values change.
        hurst: (default: 0.5)
            Hurst index of the inter-event time series. If driving noise
            is not fractional Gaussian noise, then this is an effective
            Hurst index (the "faking" process should still generate the
            same stationary inter-event time distribution as the process
            with fGn).
        boundary: (default: 1e-4)
            Fine-tunes soft boundary conditions the process is allowed
            to approach. Inter-event times confined to interval
            [0+boundary, 1-boundary].
        n_tau_points: (default: 67108864)
            The desired length of the inter-event time series. Note that
            the length applies to the event time space (k-space) and not
            physical time space. It is advisable to use numbers which
            would be powers of 2.
        fold_mid: (default: False)
            Whether to "fold" the series around the mid value (tau=0.5)
            before evaluating inter-event time PDF.
        seed: (default: -1)
            RNG seed. If negative value is passed (which is the
            default), then seed will be randomly generated by
            `np.random.rand(2**20)`.

    Output:
        Numpy array first column of which stores mid points of the boxes
        and the second column storing estimated probability density for
        that box. Note that log base 10 of the actual values is
        returned.

    Examples:
        ```
        >> from taudiff import beta_model
        >> from sim_utils import simulate_tau_pdf
        >> pdf = simulate_tau_pdf(beta_model)
        ```
    """
    tau_series = get_tau_series(
        model,
        sigma=sigma,
        hurst=hurst,
        boundary=boundary,
        n_points=n_tau_points,
        seed=seed,
    )

    if fold_mid:
        tau_series[tau_series > 0.5] = 1 - tau_series[tau_series > 0.5]

    return np.log10(MakeLogPdf(tau_series))


def simulate_event_stats(
    model: Callable[..., Iterator[float]],
    *,
    sigma: float = 1e-4,
    hurst: float = 0.5,
    boundary: float = 1e-4,
    dt: float = 1,
    n_batches: int = 64,
    n_batch_points: int = 1048576,
    seed: int = -1,
) -> Tuple[np.ndarray, np.ndarray]:
    """Obtain PDF and PSD of event count times (in physical time space).

    PDF and PSD are obtained by averaging over results of `n_batches`
    independet simulations (each of length `n_batch_points`).

    Input:
        model:
            Generator to use when generating the series. Use generators
            from the `taudiff` module or create your own, which would
            replace the interface.
        sigma: (default: 1e-4)
            Sigma parameter of the iterative equation driving the
            process. Sigma parameter controls how fast the inter-event
            time values change.
        hurst: (default: 0.5)
            Hurst index of the inter-event time series. If driving noise
            is not fractional Gaussian noise, then this is an effective
            Hurst index (the "faking" process should still generate the
            same stationary inter-event time distribution as the process
            with fGn).
        boundary: (default: 1e-4)
            Fine-tunes soft boundary conditions the process is allowed
            to approach. Inter-event times confined to interval
            [0+boundary, 1-boundary].
        dt: (default: 1)
            Discretization time step for the event count time series.
        n_batches: (default: 64)
            Number of times to sample single time series.
        n_batch_points: (default: 1048576)
            Number of points in a single sample time series. It is
            advisable to use numbers which would be powers of 2.
        seed: (default: -1)
            RNG seed. If negative value is passed (which is the
            default), then seed will be randomly generated by
            `np.random.rand(2**20)`.

    Output:
        Two numpy arrays: first stores obtained PDF, second stores
        obtained PSD. Both arrays contain two columns: N and p(N) in the
        first one, f and S(f) in the second one. Note that log base 10
        of the actual values is returned.

    Examples:
        ```
        >> from taudiff import beta_model
        >> from sim_utils import simulate_event_stats
        >> pdf, psd = simulate_event_stats(beta_model)
        ```
    """
    esg = event_series_generator(
        model,
        sigma=sigma,
        hurst=hurst,
        boundary=boundary,
        dt=dt,
        n_batch_points=n_batch_points,
        seed=seed,
    )

    pdfs = ()
    psds = ()

    for _ in range(n_batches):
        event_series = next(esg)
        pdfs = pdfs + (MakeLogPdf(event_series),)  # type: ignore
        psds = psds + (MakeLogPsd(event_series),)  # type: ignore
    del esg, event_series

    pdf = np.log10(AverageOverLogLog(pdfs))
    psd = np.log10(AverageOverLogLog(psds))

    return pdf, psd
