# Musterloesung
#
# Mock-Likelihood fuer den Gravitationswellenhintergrund (GWB), wie ihn
# Pulsar-Timing-Arrays (PTAs, z.B. NANOGrav) messen.
#
# Die "Daten" bestehen aus 14 Frequenz-Bins f_k = k / T_obs (k = 1, ..., 14),
# mit T_obs = 16 Jahren (Laufzeit des NANOGrav-15yr-Datensatzes). Fuer jeden
# Bin gibt es eine (Gauss'sche) Messung von log10(h^2 Omega_gw(f_k)) mit
# Mittelwert mu_k und Unsicherheit sigma_k -- genau wie Datenpunkte mit
# Fehlerbalken in einem gewoehnlichen chi^2-Fit:
#
#   chi^2 = sum_k ( (log10 Omega_model(f_k) - mu_k) / sigma_k )^2 ,
#
# nur dass wir chi^2 gleich als (negative, doppelte) log-Likelihood
# interpretieren. Jeder der 14 Frequenz-Bins liefert also eine eigene,
# unabhaengige Gauss-Likelihood fuer log10(h^2 Omega_gw(f_k)); die
# Gesamt-Likelihood ist das Produkt dieser 14 Teil-Likelihoods (die
# Gesamt-log-Likelihood also ihre Summe).
#
# mu_log10 und sigma_log10 sind keine frei erfundenen Zahlen, sondern eine
# vereinfachte Gauss-Naeherung an den echten NANOGrav-15yr-"free spectrum"-
# Datensatz (Zenodo 10.5281/zenodo.10344086, Lamb, Taylor & van Haasteren
# 2023) -- unsere Mock-Likelihood hat also strukturell genau die Form, die
# ceffyl/PTArcade auch fuer echte PTA-Daten benutzen.

import numpy as np

# --- Beobachtungszeit und Frequenz-Gitter -----------------------------------
T_obs = 16.0 * 365.25 * 24 * 3600.0    # 16 Jahre, in Sekunden
f0 = 1.0 / T_obs                       # niedrigste Fourier-Frequenz (Hz)
n_bins = 14
freqs = f0 * np.arange(1, n_bins + 1)  # 14 Frequenz-Bins (Hz)

nHz = 1e-9                             # 1 nHz in Hz, zur Umrechnung


# --- Mock-"Messdaten" --------------------------------------------------------
# Mittelwert und Unsicherheit von log10(h^2 Omega_gw) in jedem der 14
# Frequenz-Bins (aus echten NANOGrav-15yr-Daten abgeleitet). Bin 8
# (f ~ 16 nHz) hat den am wenigsten negativen mu_log10-Wert -- eine "Beule"
# im Spektrum, wie sie auch in der Literatur (z.B. arXiv:2306.16219)
# diskutiert wird.
mu_log10 = np.array([
    -10.1074, -8.9820, -8.9061, -8.9676, -10.0959, -10.3628, -10.6969,
     -8.1848, -10.1441, -9.1966, -9.2726, -9.6692, -9.6163, -9.3616,
])

sigma_log10 = np.array([
    0.3978, 0.1909, 0.2624, 0.5017, 1.7334, 1.9396, 1.3870, 1.2318,
    1.4141, 1.6662, 1.6218, 1.3132, 1.2425, 1.2703,
])


# --- Likelihood ---------------------------------------------------------------
def log_likelihood(model_log10_vals):
    """
    log-Likelihood fuer ein Modell, das an den 14 Frequenzen `freqs` die
    Werte `model_log10_vals` (log10 h^2 Omega_gw) vorhersagt.

    L_k(model_k) = Gauss(model_k; mu_log10[k], sigma_log10[k]) fuer jeden der
    14 Bins; Gesamt-log-Likelihood = Summe der 14 einzelnen log L_k.
    """
    model = np.asarray(model_log10_vals)
    log_likes_per_bin = (
        -0.5 * ((model - mu_log10) / sigma_log10) ** 2
        - 0.5 * np.log(2.0 * np.pi * sigma_log10**2)
    )
    return np.sum(log_likes_per_bin)


# -----------------------------------------------------------------------------
if __name__ == "__main__":
    import matplotlib.pyplot as plt

    plt.errorbar(freqs / nHz, mu_log10, yerr=sigma_log10, fmt="o")
    plt.xscale("log")
    plt.xlabel(r"$f$ [nHz]")
    plt.ylabel(r"$\log_{10}(h^2 \Omega_{\rm gw})$")
    plt.tight_layout()
    plt.savefig("mock_data.pdf", dpi=150, bbox_inches="tight")
    print("Plot gespeichert als mock_data.pdf")