import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
from scipy.optimize import minimize
from astropy.cosmology import FlatLambdaCDM
import sympy as sp

# --- Symbolic constants & helpers ---

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

# Fibonacci sequence (closed form / Binet's formula)
def fib(n):
    return (phi**n - (-1/phi)**n) / np.sqrt(5)

# Prime numbers list (extend as needed)
PRIMES = [2,3,5,7,11,13,17,19,23,29,31,37,41,43,47]

# Dimensional DNA operator D_{n,β}(r)
# Here r and k are placeholders (k exponent can vary)
def D(n, beta, r=1.0, k=1.0, Omega=1.0, base=2):
    # F_n,b = fib(n+beta) for fractional n+β
    Fn_beta = fib(n + beta)
    Pn_beta = PRIMES[int((n + beta) % len(PRIMES))]  # Simplified prime injection
    dyadic = base ** (n + beta)
    return np.sqrt(phi * Fn_beta * dyadic * Pn_beta * Omega) * (r ** k)

# Inverse function to approximate (n, β) given numeric val — heuristic only
def invert_D(value, r=1.0, k=1.0, Omega=1.0, base=2, max_n=10, steps=1000):
    candidates = []
    for n in np.linspace(-max_n, 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))
    return min(candidates, key=lambda x: x[0])

# --- Symbolic parameter reconstruction (example from your fitted data) ---

# Fitted numeric params (from prior run or your symbolic compression)
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,          # We will reconstruct symbolically below
    'c0':    299792.458,    # Speed of light (km/s), to emerge symbolically
    'M':    -19.3,          # Absolute magnitude offset, observational
}

print("Symbolic decomposition of fitted parameters:")

symbolic_indices = {}
for name, val in fitted_params.items():
    if name in ['M']:  # leave M as is for now
        print(f"  {name:10} : {val} (fixed observational)")
        continue
    # invert numeric val to symbolic (n, beta)
    err, n, beta = invert_D(val)
    symbolic_indices[name] = (n, beta)
    val_recon = D(n, beta)
    print(f"  {name:10} : approx D({n:.3f}, {beta:.3f}) = {val_recon:.6f} (orig: {val:.6f}, err={err:.3e})")

# Now define cosmological model using full symbolic emergence:

# Cosmological functions (using your prior scaling logic)
def a_of_z(z):
    return 1 / (1 + z)

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

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

def G_z(z, k, r0, Omega0, s0, alpha, beta):
    return Omega_z(z, Omega0, alpha) * k**2 * r0 / s_z(z, s0, beta)

def H_z(z, k, r0, Omega0, s0, alpha, beta, H0):
    Om_m = 0.3
    Om_de = 0.7
    Gz = G_z(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, c0):
    return c0 * (Omega_z(z, Omega0, alpha) / Omega0) ** gamma

def compute_luminosity_distance_grid(z_max, params, n=500):
    k, r0, Omega0, s0, alpha, beta, gamma, H0, c0 = params
    z_grid = np.linspace(0, z_max, n)
    c_z = emergent_c(z_grid, Omega0, alpha, gamma, c0)
    H_grid = H_z(z_grid, k, r0, Omega0, s0, alpha, beta, H0)
    integrand = c_z / H_grid
    integral = np.cumsum((integrand[:-1] + integrand[1:]) / 2 * np.diff(z_grid))
    integral = np.insert(integral, 0, 0)
    d_c = interp1d(z_grid, integral, kind='cubic', fill_value="extrapolate")
    return lambda z: (1 + z) * d_c(z)

def model_mu(z_arr, params, M):
    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 + M

# 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']

# Build fully symbolic parameters from reconstructed indices:
def param_from_index(name, n_beta, r=1.0, k=1.0, Omega=1.0):
    n, beta = n_beta
    return D(n, beta, r=r, k=k, Omega=Omega)

# Compose full parameter list for cosmology model:
param_list = []
for pname in ['k', 'r0', 'Omega0', 's0', 'alpha', 'beta', 'gamma']:
    param_list.append(param_from_index(pname, symbolic_indices[pname]))
# Add H0 and c0 emergent parameters:
param_list.append(param_from_index('H0', (0,0), r=1.0, k=0, Omega=1.0) * 70.0/70.0)  # placeholder, refine as needed
param_list.append(param_from_index('c0', (0,0), r=1.0, k=0, Omega=1.0) * 299792.458/299792.458)

# Fixed M for now:
M = fitted_params['M']

# Calculate fit and residuals
mu_fit = model_mu(z, param_list, M)
residuals = mb - mu_fit  # mb is raw apparent mag, no subtraction of M here since included in model_mu

# Plot
plt.figure(figsize=(10, 6))
plt.errorbar(z, mb, yerr=dmb, fmt='.', alpha=0.5, label='Data')
plt.plot(z, mu_fit, 'r-', label='Fully Symbolic Emergent Model')
plt.xlabel('Redshift (z)')
plt.ylabel('Distance Modulus (μ)')
plt.title('Supernova Distance Modulus: Fully Symbolic Emergent Constants')
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: Fully Symbolic Model')
plt.grid(True)
plt.tight_layout()
plt.show()
