#!/usr/bin/env python3
"""
plot_js_sensitivity.py

by Peter De Ceuster - 2025
Creative Commons Attribution 4.0 International 

Usage examples:
  python plot_js_sensitivity.py            # runs default toy examples and writes PNGs
  python plot_js_sensitivity.py --json input_params.json

Produces:
  - delta_phi_vs_Lambda.png
  - required_qs_vs_Lambda.png

Notes:
 - Units: GeV for energy scales (Lambda), eV for photon_energy input (converted).
 - Uses the toy normalization choices from Appendix K. Adjust `norm_factor` to match
   the explicit internal overlap integrals you compute for a chosen compactification.
"""

import json
import argparse
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path

# ---------- constants ----------
# conversion: 1 m ~= 5.06773065e15 GeV^-1  (hbar=c=1)
METER_TO_GEV_INV = 5.06773065e15
# eV -> GeV
EV_TO_GEV = 1e-9

# Experimental sensitivity (phase) references from your Appendix I
SENSITIVITIES = {
    "tabletop_MZ (typ.)": 1e-8,
    "LIGO (optimistic)": 1e-10,
    "LIGO (pessimistic)": 1e-12,
    "high-finesse cavity": 1e-14,
    "conservative publisher target": 1e-6,
}

# ---------- physics helper functions ----------
def compute_js_from_compactification(N_over_VK, Lambda_GeV, p):
    """
    Toy ansatz: ||Js|| = (N / V_K) * Lambda^p
    N_over_VK : float
    Lambda_GeV : array or scalar
    p : exponent
    returns ||Js|| in GeV^p (formal units consistent with paper's toy example)
    """
    return N_over_VK * (Lambda_GeV ** p)


def iota_k_js_eff(Js_norm, E_gamma_eV, Lambda_norm_GeV, norm_factor=1.0):
    """
    Toy contraction:
      (i_k J_s)_eff ~ ||Js|| * (E_gamma / Lambda_norm) * norm_factor
    E_gamma_eV : photon energy in eV
    Lambda_norm_GeV : normalization scale in GeV (model dependent; Appendix K used Lambda)
    norm_factor : additional geometric/overlap suppression (<<1 typically)
    """
    E_gamma_GeV = E_gamma_eV * EV_TO_GEV
    return Js_norm * (E_gamma_GeV / (Lambda_norm_GeV if Lambda_norm_GeV > 0 else 1.0)) * norm_factor


def delta_phi_from_iota(iota_eff, L_m, qs=1.0):
    """
    delta_phi = q_s * (i_k J_s)_eff * L
    L_m is arm length in meters -> convert to GeV^-1
    """
    L_gev_inv = L_m * METER_TO_GEV_INV
    return qs * iota_eff * L_gev_inv


# ---------- plotting helpers ----------
def plot_delta_phi_vs_Lambda(Lambda_range, delta_phi_curves, labels, outpath):
    plt.figure(figsize=(8,6))
    for phi, lbl in zip(delta_phi_curves, labels):
        plt.loglog(Lambda_range, np.abs(phi), label=lbl, lw=2)
    # overlay sensitivities
    for name, val in SENSITIVITIES.items():
        plt.hlines(val, Lambda_range[0], Lambda_range[-1], linestyles='--', alpha=0.7,
                   label=f"sensitivity: {name} ({val:.1g})")
    plt.xlabel("Compactification scale Λ (GeV)")
    plt.ylabel("Predicted |δφ| (radians)")
    plt.title("Predicted δφ vs Λ (toy model) — overlayed with sensitivity levels")
    plt.legend(loc='lower left', fontsize='small')
    plt.grid(which='both', ls=':', alpha=0.4)
    plt.tight_layout()
    plt.savefig(outpath, dpi=200)
    plt.close()
    print(f"Wrote {outpath}")


def plot_required_qs_vs_Lambda(Lambda_range, base_factor, E_gamma_eV, L_m, sens_keys, outpath):
    """
    Computes required q_s to reach each sensitivity level as a function of Lambda:
      q_s_required = sensitivity / ( (i_kJs)_eff * L )
    base_factor is (N_over_VK * Lambda^(p-1) * E_gamma * L_gev_inv * norm_factor)
    But we'll compute directly for clarity.
    """
    plt.figure(figsize=(8,6))
    L_gev_inv = L_m * METER_TO_GEV_INV
    for name in sens_keys:
        sens = SENSITIVITIES[name]
        # Compute required qs for each Lambda in the range
        required_qs = []
        for Lambda in Lambda_range:
            Js = base_factor['N_over_VK'] * (Lambda ** base_factor['p'])
            iota = iota_k_js_eff(Js, E_gamma_eV, base_factor['Lambda_norm'], norm_factor=base_factor['norm_factor'])
            denom = iota * L_gev_inv
            # avoid division by zero
            if denom == 0:
                required_qs.append(np.nan)
            else:
                required_qs.append(sens / denom)
        required_qs = np.array(required_qs)
        plt.loglog(Lambda_range, np.abs(required_qs), label=f"q_s required for {name} ({sens:.1g})")
    plt.xlabel("Compactification scale Λ (GeV)")
    plt.ylabel("Required q_s to reach sensitivity")
    plt.title("Coupling q_s required vs Λ")
    plt.legend(fontsize='small', loc='lower left')
    plt.grid(which='both', ls=':', alpha=0.4)
    plt.tight_layout()
    plt.savefig(outpath, dpi=200)
    plt.close()
    print(f"Wrote {outpath}")


# ---------- main routine ----------
def main(params):
    # Unpack params (use defaults if not present)
    N_over_VK = params.get("N_over_VK", 1.0)        # toy flux density
    p = params.get("p", 1.0)                       # exponent in ||Js|| ~ Lambda^p
    qs = params.get("qs", 1.0)
    E_gamma_eV = params.get("E_gamma_eV", 1.0)     # optical photon ~1 eV
    L_m = params.get("L_m", 1.0)                   # arm length in meters
    Lambda_norm = params.get("Lambda_norm_GeV", None)  # if None, set to Lambda in loop
    norm_factor = params.get("norm_factor", 1e-12)  # typical suppression: tune for realism
    Lambda_min = params.get("Lambda_min_GeV", 1e2)
    Lambda_max = params.get("Lambda_max_GeV", 1e15)
    npoints = params.get("npoints", 200)

    Lambda_range = np.logspace(np.log10(Lambda_min), np.log10(Lambda_max), npoints)

    # produce curves for a few choices: default denom Lambda_norm = Lambda
    delta_phi_curves = []
    labels = []
    # Example sets: 3 normalization choices
    for Ln_factor in [1.0, 1e-2, 1e-6]:
        # Ln_factor multiplies N_over_VK (simulates different flux densities / overlaps)
        Nfactor = N_over_VK * Ln_factor
        phi_vals = []
        for Lambda in Lambda_range:
            Js = compute_js_from_compactification(Nfactor, Lambda, p)  # ||Js||
            if Lambda_norm is None:
                Lnorm = Lambda
            else:
                Lnorm = Lambda_norm
            iota = iota_k_js_eff(Js, E_gamma_eV, Lnorm, norm_factor=norm_factor)
            phi = delta_phi_from_iota(iota, L_m, qs=qs)
            phi_vals.append(np.abs(phi))
        delta_phi_curves.append(np.array(phi_vals))
        labels.append(f"N/VK scaled by {Ln_factor:.0e}")

    # Plot predictions
    out1 = Path(params.get("out_prefix", "delta_phi")) / "delta_phi_vs_Lambda.png"
    out1.parent.mkdir(parents=True, exist_ok=True)
    plot_delta_phi_vs_Lambda(Lambda_range, delta_phi_curves, labels, out1)

    # Produce required q_s plot
    base_factor = {
        'N_over_VK': N_over_VK,
        'p': p,
        'Lambda_norm': Lambda_norm if Lambda_norm is not None else 1.0,  # placeholder, used in function
        'norm_factor': norm_factor
    }
    out2 = Path(params.get("out_prefix", "delta_phi")) / "required_qs_vs_Lambda.png"
    plot_required_qs_vs_Lambda(Lambda_range, base_factor, E_gamma_eV, L_m, list(SENSITIVITIES.keys()), out2)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Plot Js -> delta phi predictions and sensitivities.")
    parser.add_argument("--json", type=str, default=None, help="JSON file with model parameters.")
    parser.add_argument("--out-prefix", type=str, default="supplementary_figures", help="Output folder prefix")
    args = parser.parse_args()

    if args.json:
        with open(args.json, 'r') as f:
            params = json.load(f)
        params['out_prefix'] = args.out_prefix
    else:
        # default toy params (you can replace these or provide JSON)
        params = {
            "N_over_VK": 1.0,
            "p": 1.0,
            "qs": 1.0,
            "E_gamma_eV": 1.0,
            "L_m": 1.0,
            "Lambda_norm_GeV": None,   # if None script uses Lambda itself
            "norm_factor": 1e-12,      # tuning knob; Appendix K warns this can be very small
            "Lambda_min_GeV": 1e2,
            "Lambda_max_GeV": 1e15,
            "npoints": 300,
            "out_prefix": args.out_prefix
        }
    main(params)
