"""Plot exact Fourier-bin noise correlations for finite spectral windows.

Python 3.14; NumPy and Matplotlib versions are documented in pyproject.toml.
Run in the desired output directory; this generator writes its PNG to cwd.
"""
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def normalized_covariance(window):
    """Unit-variance independent sample noise, coherent amplitude scaling."""
    size = len(window)
    transform = np.exp(-2j * np.pi * np.outer(np.arange(size), np.arange(size)) / size)
    covariance = (transform * window**2) @ transform.conj().T / window.sum()**2
    assert np.max(np.abs(covariance.imag)) < 1e-12
    return covariance.real / covariance[0, 0].real


def main():
    size = 16
    sample = np.arange(size)
    windows = [np.ones(size), 0.5 - 0.5 * np.cos(2 * np.pi * sample / size)]
    titles = ['Rectangular window', 'Periodic Hann window']
    fig, axes = plt.subplots(1, 2, figsize=(9, 4.3), dpi=120, layout='constrained')
    fig.patch.set_facecolor('white')
    for axis, window, title in zip(axes, windows, titles):
        covariance = normalized_covariance(window)
        image = axis.imshow(covariance, origin='lower', cmap='RdBu_r', vmin=-1, vmax=1,
                            interpolation='nearest', extent=(-0.5, size - 0.5, -0.5, size - 0.5))
        axis.set(title=title, xlabel='Fourier bin l', ylabel='Fourier bin k')
        axis.set_xticks([0, 4, 8, 12, 15])
        axis.set_yticks([0, 4, 8, 12, 15])
    colorbar = fig.colorbar(image, ax=axes, shrink=0.84, pad=0.02)
    colorbar.set_label(r'$C_{kl}/C_{00}$ (own-window normalization)')
    fig.suptitle('Windowing correlates Fourier-bin noise: N = 16', fontsize=13)
    fig.savefig(Path('windowed-fourier-bin-noise-covariance.png'), facecolor='white', transparent=False)
    plt.close(fig)


if __name__ == '__main__':
    main()