# %% [markdown]
# # Minimal Example
#
# The main goal of this example is to demonstrate how to use STACIE
# with a minimal, self-contained example.
# First, the properties of a basic Markov process are discussed,
# and then data is generated using this process.
# (A detailed derivation of the analytical results is provided in the last section.)
# The Markov chains are analyzed using
# two [models](../theory/model.md),
# followed by some comments on their applicability.
#
# A secondary goal is to thoroughly discuss the plots generated by STACIE,
# which can help detect problems with the analysis or input data.

# %% [markdown]
# ## Library Imports and Matplotlib Configuration

# %%
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from stacie import (
    compute_spectrum,
    estimate_acint,
    ExpPolyModel,
    LorentzModel,
    UnitConfig,
    plot_extras,
    plot_fitted_spectrum,
)

# %%
mpl.rc_file("matplotlibrc")
# %config InlineBackend.figure_formats = ["svg"]

# %% [markdown]
# ## The Markov Process
#
# We use the following simple [discrete-time Markov process](https://en.wikipedia.org/wiki/Discrete-time_Markov_chain):
#
# $$
#   \hat{x}_{n+1} = \alpha \hat{x}_n + \beta \hat{z}_n
# $$
#
# where $\hat{z}_n$ are uncorrelated standard normal random variables,
# and $\alpha$ and $\beta$ are real constants.
# The parameter $\alpha$ controls the autocorrelation of the process,
# with $0 < \alpha < 1$.
#
# One can show that the autocorrelation function of this process is given by
#
# $$
#   c_\Delta = \frac{\beta^2}{1 - \alpha^2}\, \alpha^{|\Delta|}
# $$
#
# The variance, autocorrelation integral (with $F=1$ and $h=1$) and integrated correlation time are respectively:
#
# $$
#   \begin{aligned}
#     \sigma^2 &= \frac{\beta^2}{1 - \alpha^2}
#     \\
#     \mathcal{I} &= \frac{1}{2}\left(\frac{\beta}{1 - \alpha}\right)^2
#     \\
#     \tau_{\text{int}} &= \frac{\mathcal{I}}{F\,c_0} = \frac{1}{2}\frac{1 + \alpha}{1 - \alpha}
#   \end{aligned}
# $$
#
# Derivations of these equations can be found in the final section of this notebook.
#
# The example below uses the parameters $\alpha = 31/33$ and $\beta = \sqrt{8/1089}$,
# for which we expect $\mathcal{I} = 1$ and $\tau_{\text{int}} = 16$.
#

# %% [markdown]
# ## Data generation
#
# The following code cell implements 64 independent realizations
# of the Markov process of 32768 steps each.
# This implementation vectorizes over independent sequences,
# which is much faster than generating them one by one.

# %%
nseq = 64
nstep = 1024 * 32
alpha = 31 / 33
beta = np.sqrt(8 / 1089)
std = beta / np.sqrt(1 - alpha**2)
rng = np.random.default_rng(0)
sequences = np.zeros((nseq, nstep))
sequences[:, 0] = rng.normal(0, std, nseq)
for i in range(1, nstep):
    sequences[:, i] = alpha * sequences[:, i - 1] + rng.normal(0, beta, nseq)

# %% [markdown]
# ## Analysis With STACIE, Using the ExpPoly Model
#
# The following code cell estimates the autocorrelation integral using the `ExpPolyModel`.
# Because the autocorrelation decays exponentially,
# the spectrum features a Lorentzian peak at zero frequency.
# Hence, we use degrees $S=\{0, 2\}$ for the polynomial, which ensures a zero-derivative at the origin.

# %%
spectrum = compute_spectrum(sequences)
result_exppoly = estimate_acint(spectrum, ExpPolyModel([0, 2]), verbose=True)

# %% [markdown]
# With the `verbose=True` option, STACIE prints the results of the analysis.
# The first section of the screen output shows the progress of the cutoff frequency scan
# and includes the following columns:
# - `neff`: the number of effective spectrum data points used in the fit.
# - `criterion`: the value of the cutoff criterion used for the weighted average
#    over solutions at different cutoff frequencies.
# - `fcut`: the cutoff frequency used for the fit.
#
# The second section summarizes the analysis,
# and can be easily related to the concepts in the [theory section](../theory/index.md).
# These results already reveal that STACIE reproduces the expected results.
#
# The next code cell plots the model fitted to the spectrum.
# %%
uc = UnitConfig()
plt.close("fitted_exppoly")
fig, ax = plt.subplots(num="fitted_exppoly")
plot_fitted_spectrum(ax, uc, result_exppoly)

# %% [markdown]
# This plot displays a lot of information:
# - The spectrum is shown as **blue dots**.
# - The fitted model is plotted as a **solid green line**.
#   Its parameters are the weighted average over multiple cutoff frequencies.
# - The **red dotted line** shows the weighted average of the switching functions,
#   used to identify the low-frequency region of the spectrum.
# - The **green band** shows the expected {term}`uncertainty` of the sampling spectrum,
#   as a 95% confidence interval.
#   Most of the blue data points should fall within this band,
#   at least in the region where the model is fitted to the data
#   (red dotted line close to 1.0).
# - The **green dashed lines** are the 95% confidence intervals of the fitted model.
#   This confidence interval should be narrow in the low-frequency region.
# - The **black vertical line** corresponds to the weighted average
#   of the cutoff frequencies used in the fit.
# - The **plot title** summarizes key information about the analysis, including:
#     - The model used to fit the data.
#     - The autocorrelation integral and its uncertainty.
#     - The integrated correlation time and its uncertainty.
# - The **legend** shows to the confidence level used for plotting.
#
# The next plot shows some intermediate results,
# which can help you understand the fitting process or detect problems.
# %%
plt.close("extras_exppoly")
fig, axs = plt.subplots(2, 2, num="extras_exppoly")
plot_extras(axs, uc, result_exppoly)

# %% [markdown]
# This plot contains four panels with extra results:
# 1. The **top left panel** shows the weight assigned to each cutoff frequency,
#    based on the CV2L criterion.
#    Things to look for:
#    - If the cutoff weight is large for the lowest cutoffs,
#      then the input sequences are likely too short.
#      This typically also results in a low number of effective data points.
#      Increasing the number of steps in the inputs will increase the frequency resolution of
#      the spectrum, allowing for better fits with lower cutoff frequencies.
#    - If the cutoff weight is large for the highest cutoffs,
#      there is most likely also a problem with the analysis.
#      There can be multiple causes for this:
#      - The input sequences are much longer than necessary.
#        In this case, you can increase the `neff_max` option of
#        [`estimate_acint()`](#stacie.estimate.estimate_acint)
#        to fit the model with higher cutoffs.
#        However, this can be expensive, so it is recommended
#        to use more and shorter sequences instead.
#        This can be done by preprocessing the data with the
#        [`split()`](#stacie.utils.split) function before computing the spectrum.
#        Even better is to plan ahead and avoid this situation.
#      - The data is block-averaged with a block size that is too large,
#        which limits the available frequency range.
# 2. The **top right panel** shows the autocorrelation integral for each cutoff frequency.
#    The dots indicate the extent to which each point contributes to the final result.
#    (Black is high weight, white is low weight.)
#    Things to look for:
#    - If the autocorrelation integral shows sharp jumps at low cutoff frequencies,
#      the model is most likely overfitting the data at low frequencies.
#      These points are practically always given low cutoff weights, so you can ignore them.
#      However, if you want to exclude them from the analysis,
#      you can increase the `neff_min` option of
#      [`estimate_acint()`](#stacie.estimate.estimate_acint).
# 3. The **bottom left panel** the Z-scores of the regression cost and the cutoff criterion,
#    as a function of the cutoff frequency.
#    The Z-score is the number of standard deviations a value deviates from its mean.
#    For ill-behaved fits, the Z-scores easily exceed 2.
#    When providing sufficient inputs, high Z-scores should only occur where the cutoff weight is low.
#    If the Z-scores are high for cutoff frequencies with high cutoff weights,
#    the input data is insufficient for reliable error estimation or the model is not appropriate.
#    In this case, it is recommended to use a different model
#    or to increase the length of the input sequences.
# 4. The **bottom right panel** shows the eigenvalues of the Hessian matrix of the fit,
#    in a preconditioned parameter space, at each cutoff frequency.
#    A large spread of the eigenvalues indicates that the fit is not well constrained.
#    Such a large spread typically results in overfitting artifacts.

# %% [markdown]
# ## Analysis With STACIE, Using the Lorentz Model
#
# In this example, we know *a priori* that the autocorrelation function decays exponentially.
# Therefore, the `LorentzModel` should be able to perfectly explain the spectrum,
# up to the statistical noise in the data.

# %%
# Analysis
result_lorentz = estimate_acint(spectrum, LorentzModel(), verbose=True)

# Plotting
plt.close("fitted_lorentz")
fig, ax = plt.subplots(num="fitted_lorentz")
plot_fitted_spectrum(ax, uc, result_lorentz)
plt.close("extras_lorentz")
fig, axs = plt.subplots(2, 2, num="extras_lorentz")
plot_extras(axs, uc, result_lorentz)

# %% [markdown]
# The extra plots reveal some noteworthy features:
# - At the lowest cutoff frequencies, the model clearly overfits the data.
#   There is a large spread on the eigenvalues,
#   and the autocorrelation integral fluctuates significantly for low cutoff frequencies.
#   Although the results at these cutoff frequencies are unreliable, they are given low weights,
#   so you don't need to intervene manually to exclude these results.
# - The cutoff weight is maximal for the highest cutoff frequencies.
#   This is typically a sign that the input sequences are too long, but this is not the case here.
#   While the Lorentz model would also yield excellent results with shorter sequences,
#   this would still result in a high frequency cutoff.
#   This is because the Lorentz model fits the data perfectly;
#   more data points will always result in a lower cutoff criterion.
#   However, in more realistic cases involving data from complex simulations or measurements,
#   this is unlikely to happen.

# %%  [markdown]
# ## Regression Tests
#
# If you are experimenting with this notebook, you can ignore any exceptions below.
# The tests are only meant to pass for the notebook in its original form.

# %%
if abs(result_exppoly.acint - 1.0) > 0.03:
    raise ValueError(f"Wrong acint: {result_exppoly.acint:.4e}")
if abs(result_exppoly.corrtime_int - 16.0) > 0.5:
    raise ValueError(f"Wrong corrtime_int: {result_exppoly.corrtime_int:.4e}")
if abs(result_lorentz.acint - 1.0) > 0.03:
    raise ValueError(f"Wrong acint: {result_lorentz.acint:.4e}")
if abs(result_lorentz.corrtime_int - 16.0) > 0.5:
    raise ValueError(f"Wrong corrtime_int: {result_lorentz.corrtime_int:.4e}")
if abs(result_lorentz.corrtime_exp - 16.0) > 0.5:
    raise ValueError(f"Wrong corrtime_exp: {result_lorentz.corrtime_exp:.4e}")

# %% [markdown]
# ## Derivation of the Autocorrelation Integral
#
# The Markov process is defined by the following equation:
#
# $$
#   \hat{x}_{n+1} = \alpha \hat{x}_n + \beta \hat{z}_n
# $$
#
# The stationary distribution of the process is Gaussian with zero mean.
# The initial state is slowly reduced by repetitive application of the factor $\alpha$.
# The only part that remains after a long time
# is the additive contributions from the normal noise $\hat{z}_n$.
#
# Since the two terms on the right-hand side of the equation are independent,
# the variance of the stationary distribution can be found by solving:
#
# $$
#   c_0 = \alpha^2 c_0 + \beta^2
# $$
#
# This gives:
#
# $$
#   c_0 = \frac{\beta^2}{1 - \alpha^2}
# $$
#
# The covariance of two neighboring points is given by:
#
# $$
#   \cov[\hat{x}_n, \hat{x}_{n+1}]
#   = \alpha \cov[\hat{x}_n, \hat{x}_n] = \alpha\, c_0
# $$
#
# This can easily be generalized to points separated by $\Delta>0$ steps through induction:
#
# $$
#   \cov[\hat{x}_n, \hat{x}_{n+\Delta}]
#   = \alpha\, \cov[\hat{x}_n, \hat{x}_{n+\Delta-1}]
#   = \alpha^{\Delta}\, c_0
# $$
#
# with a similar result for $\Delta<0$.
# Combining these results gives:
#
# $$
#   c_\Delta = \frac{\beta^2}{1 - \alpha^2}\, \alpha^{|\Delta|}
# $$
#
# The autocorrelation integral is defined by a simple quadrature rule (with $F=1$ and $h=1$):
#
# $$
#   \mathcal{I} = \frac{1}{2} \sum_{\Delta=-\infty}^{\infty} \frac{\beta^2}{1 - \alpha^2}\, \alpha^{|\Delta|}
# $$
#
# This can be rewritten easily using properties of geometric series.
# One must be careful not to double-count the $\Delta=0$ term.
# This can be accomplished by isolating this term
# and rewriting the remaining terms in the sum with a shifted index,
# $n=|\Delta|-1$.
# After replacing the index, we use $\sum_{\Delta=1}^\infty \alpha^{\Delta}=\alpha\sum_{n=0}^{\infty}\alpha^n$.
#
# $$
#   \begin{aligned}
#     \mathcal{I}
#     &= \frac{\beta^2}{(1 - \alpha^2)}\left(
#          \frac{1}{2} + \alpha \sum_{n=0}^{\infty} \alpha^n
#        \right)
#     \\
#     &= \frac{\beta^2}{(1 - \alpha^2)}\left(
#          \frac{1}{2} + \frac{\alpha}{1 - \alpha}
#        \right)
#     \\
#     &= \frac{\beta^2}{(1 - \alpha^2)}\left(
#          \frac{1 + \alpha}{2(1 - \alpha)}
#        \right)
#     \\
#     &= \frac{1}{2}\left(\frac{\beta}{1 - \alpha}\right)^2
#   \end{aligned}
# $$
