import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d

# Golden ratio constant
phi = (1 + np.sqrt(5)) / 2

PRIMES = [
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29,
    31, 37, 41, 43, 47, 53, 59, 61, 67, 71,
    73, 79, 83, 89, 97, 101, 103, 107, 109, 113,
    127, 131, 137, 139, 149, 151, 157, 163, 167, 173,
    179, 181, 191, 193, 197, 199, 211, 223, 227, 229
]

def fib_real(n):
    from math import cos, pi, sqrt
    phi_inv = 1 / phi
    term1 = phi**n / sqrt(5)
    term2 = (phi_inv**n) * cos(pi * n)
    return term1 - term2

def D(n, beta, r=1.0, k=1.0, Omega=1.0, base=2):
    Fn_beta = fib_real(n + beta)
    idx = int(np.floor(n + beta) + len(PRIMES)) % len(PRIMES)
    Pn_beta = PRIMES[idx]
    dyadic = base ** (n + beta)
    val = phi * Fn_beta * dyadic * Pn_beta * Omega
    val = np.maximum(val, 1e-15)  # Avoid sqrt of zero/negative
    return np.sqrt(val) * (r ** k)

def invert_D(value, r=1.0, k=1.0, Omega=1.0, base=2, max_n=10, steps=100):
    candidates = []
    for n in np.linspace(0, max_n, steps):
        for beta in np.linspace(0, 1, 10):
            val = D(n, beta, r, k, Omega, base)
            candidates.append((abs(val - value), n, beta))
    best = min(candidates, key=lambda x: x[0])
    return best[1], best[2]

# Use your fitted parameters from symbolic inversion (dimensionless scalings)
fitted_params = {
    'k':    1.049342,
    'r0':   1.049676,
    'Omega0': 1.049675,
    's0':   0.994533,
    'alpha': 0.340052,
    'beta':  0.360942,
    'gamma': 0.993975,
    'H0':   70.0,              # Physical Hubble constant (km/s/Mpc)
    'c0':   3303.402087,       # Emergent base speed of light (symbolic scale)
    'M':    -19.3              # Fixed absolute magnitude
}

print("Symbolic decomposition of fitted parameters:")
for name, val in fitted_params.items():
    if name == 'M':
        print(f"  {name:<10}: {val} (fixed observational)")
        continue
    n, beta = invert_D(val)
    approx_val = D(n, beta)
    err = abs(val - approx_val)
    print(f"  {name:<10}: approx D({n:.3f}, {beta:.3f}) = {approx_val:.6f} (orig: {val:.6f}, err={err:.2e})")

params_reconstructed = {}
for name, val in fitted_params.items():
    if name == 'M':
        params_reconstructed[name] = val
        continue
    n, beta = invert_D(val)
    params_reconstructed[name] = D(n, beta)

print("\nReconstructed parameters:")
for name, val in params_reconstructed.items():
    print(f"  {name:<10} = {val:.6f}")

# Load supernova data
filename = 'hlsp_ps1cosmo_panstarrs_gpc1_all_model_v1_lcparam-full.txt'
lc_data = np.genfromtxt(
    filename,
    delimiter=' ',
    names=True,
    comments='#',
    dtype=None,
    encoding=None
)

z = lc_data['zcmb']
mb = lc_data['mb']
dmb = lc_data['dmb']
M = params_reconstructed['M']
mu_obs = mb - M

H0 = params_reconstructed['H0']
c0_emergent = params_reconstructed['c0']  # Emergent base c scale

def a_of_z(z):
    return 1 / (1 + z)

def Omega(z, Omega0, alpha):
    return Omega0 / (a_of_z(z) ** alpha)

def s(z, s0, beta):
    return s0 * (1 + z) ** (-beta)

def G(z, k, r0, Omega0, s0, alpha, beta):
    return Omega(z, Omega0, alpha) * k**2 * r0 / s(z, s0, beta)

def H(z, k, r0, Omega0, s0, alpha, beta):
    Om_m = 0.3
    Om_de = 0.7
    Gz = G(z, k, r0, Omega0, s0, alpha, beta)
    Hz_sq = (H0 ** 2) * (Om_m * Gz * (1 + z) ** 3 + Om_de)
    return np.sqrt(Hz_sq)

def emergent_c(z, Omega0, alpha, gamma):
    # Emergent speed of light, relative to emergent base scale c0_emergent
    return c0_emergent * (Omega(z, Omega0, alpha) / Omega0) ** gamma

def emergent_length_scale(z, Omega0, alpha, delta=1.0):
    """
    Emergent length scale factor as a function of redshift,
    consistent with c(z) scaling and field tension Omega scaling.
    delta is tunable exponent controlling length scaling power.
    """
    # Use inverse scaling of speed of light to define length scale
    # Length scale ~ c(z)^{delta} / H0^{delta} or a similar functional form
    # Here we use Omega scaling to define length scale for dimensional consistency
    return (Omega(z, Omega0, alpha) / Omega0) ** delta

def compute_luminosity_distance_grid(z_max, params, n=500):
    k, r0, Omega0, s0, alpha, beta, gamma = params
    z_grid = np.linspace(0, z_max, n)

    c_z = emergent_c(z_grid, Omega0, alpha, gamma)
    H_z = H(z_grid, k, r0, Omega0, s0, alpha, beta)

    # Integrand for comoving distance: c(z)/H(z)
    integrand_values = c_z / H_z

    # Compute integral numerically with trapezoidal rule
    integral_grid = np.cumsum((integrand_values[:-1] + integrand_values[1:]) / 2 * np.diff(z_grid))
    integral_grid = np.insert(integral_grid, 0, 0)

    # Apply emergent length scale correction
    length_scale = emergent_length_scale(z_grid, Omega0, alpha, delta=1.0)
    scaled_integral = integral_grid * length_scale

    d_c = interp1d(z_grid, scaled_integral, kind='cubic', fill_value="extrapolate")

    def d_L(z):
        return (1 + z) * d_c(z)

    return d_L

def model_mu(z_arr, params):
    d_L_func = compute_luminosity_distance_grid(np.max(z_arr), params)
    d_L_vals = d_L_func(z_arr)
    return 5 * np.log10(d_L_vals) + 25

param_list = [
    params_reconstructed['k'],
    params_reconstructed['r0'],
    params_reconstructed['Omega0'],
    params_reconstructed['s0'],
    params_reconstructed['alpha'],
    params_reconstructed['beta'],
    params_reconstructed['gamma'],
]

mu_fit = model_mu(z, param_list)
residuals = mu_obs - mu_fit

plt.figure(figsize=(10, 6))
plt.errorbar(z, mu_obs, yerr=dmb, fmt='.', alpha=0.5, label='Pan-STARRS1 SNe')
plt.plot(z, mu_fit, 'r-', label='Symbolic Emergent Gravity Model')
plt.xlabel('Redshift (z)')
plt.ylabel('Distance Modulus (μ)')
plt.title('Supernova Distance Modulus with Context-Aware Emergent c(z)')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()

plt.figure(figsize=(10, 4))
plt.errorbar(z, residuals, yerr=dmb, fmt='.', alpha=0.5)
plt.axhline(0, color='red', linestyle='--')
plt.xlabel('Redshift (z)')
plt.ylabel('Residuals (μ_data - μ_model)')
plt.title('Residuals of Symbolic Model with Emergent c(z)')
plt.grid(True)
plt.tight_layout()
plt.show()
