# %% [markdown]
# # Uncertainty of the Mean of Time-Correlated Data
#
# This notebook shows how to use STACIE to compute the error of the mean
# of a time-correlated input sequence,
# meaning not all of its values are statistically independent.
#
# This is a completely self-contained example that generates input sequences
# (with MCMC) and then analyzes them with STACIE.
# Atomic units are used unless otherwise noted.
#
# We suggest experimenting with this notebook by making the following changes:
#
# - Change the number of sequences and their length.
# - Change the correlation time through `PROPOSAL_STEP`.

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

# %%
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import quad
import scipy.constants as sc
from stacie import (
    UnitConfig,
    compute_spectrum,
    estimate_acint,
    LorentzModel,
    plot_extras,
    plot_fitted_spectrum,
    plot_spectrum,
)

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

# %% [markdown]
# ## Data Generation
#
# The data for the analysis are generated by sampling a Kratzer--Feus potential
# of a diatomic molecule {cite:p}`kratzer_1920_die,fues_1926_das`
# at constant temperature.
# This potential is harmonic in $1/r$:
#
# $$
#     U(r) = \frac{K}{2} \left(\frac{r_0^2}{r} - r_0\right)^2
# $$
#
# where $K$ is the force constant and $r_0$ the equilibrium bond length.
# The sampled probability density is the Boltzmann distribution:
#
# $$
#     p_r(r) = \frac{1}{Z} \exp\left(-\frac{U(r)}{k_\text{B}T}\right)
# $$
#
# where the normalization $Z$ is the classical partition function.
#
# In this example,
# the force constant and bond length of the lithium dimer are used,
# with parameters from {cite:t}`zhao_2022_bond` converted to atomic units.
# A high temperature is used to skew the distribution to larger distances.

# %%
K = 0.015
R0 = 5.150
TEMPERATURE = 1000
BOLTZMANN = sc.value("Boltzmann constant") / sc.value("Hartree energy")
BETA = 1 / (BOLTZMANN * TEMPERATURE)
PROPOSAL_STEP = 0.1


def logprob(r):
    """Calculate the logarithm of the probability."""
    energy = 0.5 * K * (R0**2 / r - R0) ** 2
    return -BETA * energy


def plot_potential_dist():
    plt.close("boltzmann")
    _, (ax1, ax2) = plt.subplots(2, 1, num="boltzmann", sharex=True)
    rgrid = np.linspace(0.5 * R0, 2 * R0, 100)
    ax2.sharex(ax1)
    ax1.plot(rgrid, -logprob(rgrid) / BETA)
    ax1.set_ylabel(r"Potential energy [E$_\text{h}$]")
    ax2.plot(rgrid, np.exp(logprob(rgrid)))
    ax2.set_ylabel("Boltzmann factor\n(unnormalized probability density)")
    ax2.set_xlabel("Internuclear distance [a$_0$]")


plot_potential_dist()

# %% [markdown]
# The MCMC implementation below is non-standard in the sense that it
# is vectorized to generate multiple sequences in parallel.


# %%
def sample_mcmc_chain(niter, stride, ndim, burnin, seed=42):
    """Sample independent Markov Chains with the Metropolis algorithm.

    Parameters
    ----------
    niter
        The number of MCMC iterations to run.
    stride
        The number of iterations between samples returned,
        i.e. the thinning interval.
    ndim
        The number of independent Markov chains to run.
    burnin
        The number of iterations to discard at the beginning.
    seed
        The random number generator seed.

    Returns
    -------
    result
        A 2D array of shape (ndim, niter // stride) containing the sampled sequences.
    """
    rng = np.random.default_rng(seed)
    result = np.zeros((ndim, niter // stride))
    r_old = np.full(ndim, R0)
    lp_old = logprob(r_old)
    irow = 0
    istep = 0
    while irow < result.shape[1]:
        r_new = r_old + rng.normal(0, PROPOSAL_STEP, ndim)
        lp_new = logprob(r_new)
        accept = lp_new > lp_old
        mask = ~accept
        nrnd = mask.sum()
        if nrnd > 0:
            accept[mask] = rng.uniform(0, 1, nrnd) < np.exp(lp_new[mask] - lp_old[mask])
        r_old[accept] = r_new[accept]
        lp_old[accept] = lp_new[accept]
        if burnin > 0:
            burnin -= 1
            continue
        if istep % stride == 0:
            result[:, irow] = r_new
            irow += 1
        istep += 1
    return result


sequences = sample_mcmc_chain(10240, 5, 50, 200)
print(f"(nseq, nstep) = {sequences.shape}")
mean_mc = sequences.mean()
print(f"Monte Carlo E[r] ≈ {mean_mc:.5f} > R0 = {R0:.5f}")

# %% [markdown]
# Because of the finite temperature and the anharmonicity of the potential,
# the average distance is greater than the equilibrium bond length.


# %%
# Plot the beginning of a few sequences.
# The atomic unit of length is the Bohr radius, $\mathrm{a}_0$.
def plot_chains():
    plt.close("chains")
    _, ax = plt.subplots(num="chains")
    ax.plot(sequences[0][:500], label="Chain 1")
    ax.plot(sequences[1][:500], label="Chain 2")
    ax.plot(sequences[2][:500], label="Chain 3")
    ax.set_xlabel("Step")
    ax.set_ylabel(r"Bond length [a$_0$]")
    ax.set_title("Markov Chain samples")
    ax.legend()


plot_chains()

# %% [markdown]
#
# The sequences in the plot are clearly time-correlated.
# The following cells show how STACIE can be used
# to compute the {term}`uncertainty` of this average,
# taking into account that not all samples
# are independent due to time correlations.

# %% [markdown]
# ## Uncertainty Quantification
#
# The spectrum is calculated using settings that are appropriate for error estimation.
# See the [Error Estimates](../properties/error_estimates.md) section
# for the justification of the `prefactors` and `include_zero_freq` keyword arguments.
# Since we are analyzing MCMC data, the `timestep` argument is not specified,
# corresponding to a dimensionless time step of 1.

# %%
# Compute and plot the power spectrum.
spectrum = compute_spectrum(
    sequences,
    prefactors=2.0 / sequences.size,
    include_zero_freq=False,
)

# %% [markdown]
# The `UnitConfig` object contains settings
# that are reused by most plotting functions.
# The integral has units of length squared, $\mathrm{a}_0^2$.
# (It is the variance of the mean.)

# %%
uc = UnitConfig(
    time_fmt=".1f",
    acint_fmt=".1e",
    acint_unit_str=r"a$^2_0$",
)
plt.close("spectrum")
_, ax = plt.subplots(num="spectrum")
plot_spectrum(ax, uc, spectrum, 180)

# %% [markdown]
# From the spectrum, one can already visually estimate the variance of the mean:
# the limit to zero frequency is about $6 \times 10^{-5}\,\mathrm{a}_0^2$.
# By normalizing the spectrum with the total simulation time,
# the spectrum has the correct unit of length squared.
# In the following cell, a model is fitted to the spectrum to get a more precise estimate.

# %%
result = estimate_acint(spectrum, LorentzModel(), verbose=True, uc=uc)

# %% [markdown]
# The spectrum is normalized such that the integral of the autocorrelation function
# is equal to the variance of the mean.
# Because STACIE estimates errors of the autocorrelation integral,
# it can thus also estimate errors of errors of means.
#
# The error of the mean and its uncertainty are printed in the following cell.

# %%
error_mc = np.sqrt(result.acint)
print(f"Error of the mean = {error_mc:.5f}")
error_of_error_mc = 0.5 * result.acint_std / error_mc
print(f"Uncertainty of the error of the mean = {error_of_error_mc:.5f}")

# %% [markdown]
# It is also interesting to visualize the fitted spectrum and some intermediate results.

# %%
# Plot of the sampling and fitted model spectrum.
plt.close("fitted")
_, ax = plt.subplots(num="fitted")
plot_fitted_spectrum(ax, uc, result)

# %% [markdown]
# The Lorentz model can clearly explain the spectrum,
# even well beyond the width of the peak at zero frequency.

# %%
# Plot additional intermediate results as a function of the frequency cutoff.
plt.close("extras")
_, axs = plt.subplots(2, 2, num="extras")
plot_extras(axs, uc, result)

# %% [markdown]
# The extra plots reveal several interesting challenges of the analysis:
#
# - The cutoff weight (top left panel) remains high up to the highest cutoff frequency considered.
#   If one is only interested in the zero-frequency limit of the spectrum,
#   there is little to be gained by including many data points in the fit at high frequencies,
#   well past the width of the peak at zero frequency.
#   These will not make the autocorrelation integral more precise,
#   but bear the risk of introducing some bias due to underfitting.
#   One may manually impose a maximum frequency cutoff with the `fcut_max` argument
#   of the `estimate_acint()` function.
#
# - The risk for some bias at high cutoff frequencies is also visible in the Z-score
#   associated with the cutoff criterion (green curve in the lower left panel).
#   For higher cutoff frequencies, the Z-score slowly increases to values above 2,
#   where the cutoff weight is still significant.
#
#   The reason for the higher Z-score is that the input time series is not normally distributed,
#   due to the asymmetry of the Kratzer--Fues potential.
#   As a result, the MC chain cannot be described by a Gaussian process,
#   and the uncertainty of the spectrum amplitudes is not exactly Gamma-distributed.
#   You can verify this hypothesis by rerunning this example with `TEMPERATURE = 100`
#   and `PROPOSAL_STEP = 0.03`.
#   This will result in a more symmetric distribution of bond lengths.
#   By lowering the proposal step, the correlation time remains about the same.
#   With these settings, a lower criterion Z-score is obtained at high cutoff frequencies.
#
#

# %% [markdown]
# ## Precise Mean With Numerical Quadrature
#
# Because the probability density sampled by the MC chain is one-dimensional,
# it is feasible to compute the mean using numerical quadrature,
# which is much more accurate than the Monte Carlo estimate.
# (For production simulations,
# Monte Carlo is only advantageous for high-dimensional problems.)
#
# As shown in the code below,
# the difference between the quadrature and Monte Carlo estimates
# is on the order of the estimated uncertainty of the MC result.

# %%
numer_quad = quad(lambda r: r * np.exp(logprob(r)), 0, 50)[0]
denom_quad = quad(lambda r: np.exp(logprob(r)), 0, 50)[0]
mean_quad = numer_quad / denom_quad
print(f"Quadrature  E[r]   ≈ {mean_quad:8.5f}")
print(f"Monte Carlo E[r]   ≈ {mean_mc:8.5f}")
print(f"|Difference|       = {abs(mean_quad - mean_mc):8.5f}")
print(f"Estimated MC error = {error_mc:8.5f}")

# %% [markdown]
# ## Autocorrelation time
#
# The Lorentz model estimates the *exponential correlation time* {cite:p}`sokal_1997_monte`
# from the width of the peak at zero frequency in the spectrum.
# It may differ from the *integrated autocorrelation time*.
# Only if the autocorrelation function is nothing but an exponentially decaying function,
# both should match.

# %%
print("Autocorrelation times:")
print(f"exponential: {result.corrtime_exp:5.2f} ± {result.corrtime_exp_std:5.2f}")
print(f"integrated:  {result.corrtime_int:5.2f} ± {result.corrtime_int_std:5.2f}")

# %% [markdown]
# Here, the deviation between the two autocorrelation times
# falls within the uncertainty of the estimates.
# This is the expected result, since the Lorentzian model is able to explain the whole spectrum.

# %%  [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(mean_mc - 5.28666) > 1e-3:
    raise ValueError(f"Wrong mean_mc: {mean_mc:.5f}")
if abs(error_mc - 0.00794) > 1e-3:
    raise ValueError(f"Wrong error_mc: {error_mc:.5f}")
