- implemented cosine and cosine_quasisep kernels

from typing import Any

import re

# Import SpectralMixture from the dedicated spectral_mixture.py file

# from PyAETNA.kernels.spectral_mixture import SpectralMixture

SUPPORTED_KERNELS = (

    "ExpSquared", "RBF", "Matern32", "Matern32_quasisep", "Matern52", "Matern52_quasisep", 

    "ExpSineSquared", "Cosine", "Cosine_quasisep", "SHO", "SHO_Quality", 

    "Constant", "QuasiPeriodic", "SpectralMixture"

)

# =============================================================================

# Kernel Building - Low-level (works with param values directly)

        - 'ExpSquared' / 'RBF': Squared exponential, params: [sigma, scale]

        - 'Matern32': Matern kernel with nu=3/2, params: [sigma, scale]

        - 'Matern32_quasisep': Quasi-separable Matern 3/2, params: [sigma, scale]

        - 'Matern52': Matern kernel with nu=5/2, params: [sigma, scale]

        - 'Matern52_quasisep': Quasi-separable Matern 5/2, params: [sigma, scale]

        - 'ExpSineSquared': Periodic kernel, params: [sigma, freq, gamma]

                             freq is linear frequency, internally converted to period = 1/freq

        - 'Cosine': Periodic kernel, params: [sigma, freq]

                             freq is linear frequency, internally converted to period = 1/freq

        - 'Cosine_quasisep': Quasi-separable Cosine, params: [sigma, freq]

        - 'Constant': Constant kernel, params: [const]

        - 'QuasiPeriodic': ExpSquared * ExpSineSquared, params: [sigma, se_scale, freq, gamma]

        freq = param_values[1] if len(param_values) > 1 else 1.0

        gamma = param_values[2] if len(param_values) > 2 else 0.5

        period = 1.0 / freq

        kernel = amp**2 * kernels.ExpSineSquared(scale=period, gamma=gamma)

        kernel = sigma**2 * kernels.ExpSineSquared(scale=period, gamma=gamma)

    elif kernel_type_lower == 'cosine':

        # params: [sigma, freq]

        # freq is linear frequency, kernel uses period = 1/freq

        sigma = param_values[0] if len(param_values) > 0 else 1.0

        freq = param_values[1] if len(param_values) > 1 else 1.0

        period = 1.0 / freq

        kernel = sigma**2 * kernels.Cosine(scale=period)

    elif kernel_type_lower == 'sho':

        # Damped simple harmonic oscillator kernel (Foreman-Mackey+2017)

            sigma=jnp.asarray(sigma),

        )

    elif kernel_type_lower == 'cosine_quasisep':

        # params: [sigma, freq]

        # Quasiseparable version: memory-efficient for large 1D time series

        sigma = param_values[0] if len(param_values) > 0 else 1.0

        freq = param_values[1] if len(param_values) > 1 else 1.0

        period = 1.0 / freq

        # tinygp's quasisep.Cosine requires sigma as a direct argument, not as a multiplier

        # We use jnp.asarray to ensure dtype consistency (float32/float64) as seen in SHO

        kernel = kernels.quasisep.Cosine(

            scale=jnp.asarray(period),

            sigma=jnp.asarray(sigma),

        )

    elif kernel_type_lower == 'constant':

        const_val = param_values[0] if len(param_values) > 0 else 1.0

        kernel = kernels.Constant(const_val)

            kernel += sigma_i**2 * kernels.ExpSquared(scale=ls_i) * kernels.Cosine(scale=1.0/f_i)

    else:

        supported = "ExpSquared, RBF, Matern32, Matern32_quasisep, Matern52, Matern52_quasisep, ExpSineSquared, SHO, SHO_Quality, Constant, QuasiPeriodic, SpectralMixture"

        supported = "ExpSquared, RBF, Matern32, Matern32_quasisep, Matern52, Matern52_quasisep, ExpSineSquared, Cosine, Cosine_quasisep, SHO, SHO_Quality, Constant, QuasiPeriodic, SpectralMixture"

        raise ValueError(

            f"Unknown kernel type: '{kernel_type}' (case-insensitive). "

            f"Supported types: {supported}"

def plot_periodograms_comparison(time, y_original, y_detrended, residuals, _log, 

                                  title='Periodogram Comparison', figsize=(15, 12),

                                  n_bootstrap=1000, gp_kernel=None,

                                  gp_x_variable='BJD'):

                                  gp_x_variable='BJD', nominal_uncertainties=None,

                                  map_jitter=None):

    """

    Plot Lomb-Scargle periodograms comparing original data, detrended data, and residuals.

            (time) allows a meaningful PSD comparison with the temporal

            Lomb-Scargle periodogram.  For other variables (e.g. roll angle)

            an annotation is shown instead.

        nominal_uncertainties: (optional) Array of nominal data uncertainties.

            Used to plot the nominal noise floor.

        map_jitter: (optional) MAP estimate of the white noise jitter.

            Used to plot the total white noise floor.

    Returns:

        fig: matplotlib figure

    # Limit f_max to reasonable value for photometry

    f_max = min(f_max, 50.0)  # max 50 cycles/day

    f_range = f_max - f_min

    n_freqs = 10000

    freqs = np.linspace(f_min, f_max, n_freqs)

    # ── Compute Noise Baselines ──

    s_nom = None

    s_jitter = None

    if nominal_uncertainties is not None:

        s_nom = np.mean(nominal_uncertainties**2) / f_range

    if nominal_uncertainties is not None and map_jitter is not None:

        # Total white noise = nominal + jitter

        s_jitter = (np.mean(nominal_uncertainties**2) + map_jitter**2) / f_range

    # ── Compute GP kernel PSD if kernel is provided and x_variable is BJD ──

    # The PSD comparison is only physically meaningful when both the

    # Lomb-Scargle periodogram (computed vs time) and the GP kernel PSD

                ax_lin.axhline(fap_power, color=color, linestyle=':', alpha=0.7, linewidth=0.8,

                              label=f'FAP={fap}%')

        # ── Overlay GP kernel PSD on the second row (GP + noise) ─────────

        # This is the row where data = GP + noise, so the kernel PSD shows

        # the theoretical power the GP *should* absorb.  Comparing the two

        # reveals whether the GP is over- or under-fitting power at each

        # frequency.

        if row == 1 and gp_psd_on_time and gp_psd is not None:

        # ── Overlay Noise Baselines and GP PSD on the second row (GP + noise) ──

        if row == 1:

            # Nominal noise floor

            if s_nom is not None:

                ax_lin.axhline(s_nom, color='gray', linestyle=':', alpha=0.6, linewidth=0.8,

                              label='Nominal Noise')

            # Total white noise floor (Nominal + Jitter)

            if s_jitter is not None:

                ax_lin.axhline(s_jitter, color='tab:blue', linestyle='--', alpha=0.6, linewidth=0.8,

                              label='Total White Noise')

            # GP kernel PSD

            if gp_psd_on_time and gp_psd is not None:

                ax_lin.plot(freqs, gp_psd, color='tab:red', linewidth=1.2,

                            alpha=0.85, label='GP windowed PSD')

                ax_lin.plot(freqs, gp_psd_theoretical, color='tab:green', linewidth=1.2,

                            alpha=0.85, label='GP theoretical PSD')

        elif row == 1 and gp_kernel is not None and not gp_psd_on_time:

                # Total Model PSD = GP + Total White Noise

                if s_jitter is not None:

                    s_total = gp_psd + s_jitter

                    ax_lin.plot(freqs, s_total, color='tab:purple', linewidth=1.2,

                                alpha=0.85, label='Total Model PSD')

                elif s_nom is not None:

                    s_total = gp_psd + s_nom

                    ax_lin.plot(freqs, s_total, color='tab:purple', linewidth=1.2,

                                alpha=0.85, label='Total Model PSD')

            elif gp_kernel is not None and not gp_psd_on_time:

                # GP defined on a non-time variable — annotate instead of plotting

                ax_lin.text(0.02, 0.90,

                            f'GP operates on {gp_x_variable}\n(PSD not comparable)',

        except:

            _log('Failed to set log y-scale on log-frequency panel')

        # ── Overlay GP kernel PSD on row 2, log-scale panel ──────────────

        if row == 1 and gp_psd_on_time and gp_psd is not None:

        # ── Overlay Noise Baselines and GP PSD on row 2, log-scale panel ──────────────

        if row == 1:

            # Nominal noise floor

            if s_nom is not None:

                ax_log.axhline(s_nom, color='gray', linestyle=':', alpha=0.6, linewidth=0.8,

                              label='Nominal Noise')

            # Total white noise floor (Nominal + Jitter)

            if s_jitter is not None:

                ax_log.axhline(s_jitter, color='tab:blue', linestyle='--', alpha=0.6, linewidth=0.8,

                              label='Total White Noise')

            # GP kernel PSD

            if gp_psd_on_time and gp_psd is not None:

                ax_log.plot(freqs, gp_psd, color='tab:red', linewidth=1.2,

                            alpha=0.85, label='GP windowed PSD')

                ax_log.plot(freqs, gp_psd_theoretical, color='tab:green', linewidth=1.2,

                            alpha=0.85, label='GP theoretical PSD')

        elif row == 1 and gp_kernel is not None and not gp_psd_on_time:

                # Total Model PSD = GP + Total White Noise

                if s_jitter is not None:

                    s_total = gp_psd + s_jitter

                    ax_log.plot(freqs, s_total, color='tab:purple', linewidth=1.2,

                                alpha=0.85, label='Total Model PSD')

                elif s_nom is not None:

                    s_total = gp_psd + s_nom

                    ax_log.plot(freqs, s_total, color='tab:purple', linewidth=1.2,

                                alpha=0.85, label='Total Model PSD')

            elif gp_kernel is not None and not gp_psd_on_time:

                ax_log.text(0.02, 0.90,

                            f'GP operates on {gp_x_variable}\n(PSD not comparable)',

                            transform=ax_log.transAxes, fontsize=7, color='gray',