import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad
from scipy.optimize import minimize
from scipy.interpolate import interp1d

# 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']
mu_obs = mb  # Use raw mb; M will be fitted

# Emergent cosmological scaling
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 emergent_c(z, Omega0, alpha, gamma):
    return (Omega(z, Omega0, alpha) / Omega0) ** gamma

def H(z, k, r0, Omega0, s0, alpha, beta, gamma):
    Gz = G(z, k, r0, Omega0, s0, alpha, beta)
    return np.sqrt(Gz * (1 + z)**3)  # Emergent matter-like scaling

def compute_luminosity_distance_grid(z_max, params, n=500):
    k, r0, Omega0, s0, alpha, beta, gamma = params[:7]
    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, gamma)
    integrand = c_z / H_z
    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 = params[-1]
    d_L_func = compute_luminosity_distance_grid(np.max(z_arr), params[:-1])
    d_L_vals = d_L_func(z_arr)
    return 5 * np.log10(d_L_vals) + 25 + M

def chi_squared(params, z_arr, mu_obs, mu_err):
    mu_model = model_mu(z_arr, params)
    return np.sum(((mu_obs - mu_model) / mu_err) ** 2)

# Initial guess: [k, r0, Omega0, s0, alpha, beta, gamma, M]
p0 = [1.0, 1.0, 1.0, 1.0, 0.3, 0.3, 1.0, -19.3]
bounds = [(0.01, 10)] * 7 + [(-21, -18)]

# Optimize
result = minimize(chi_squared, p0, args=(z, mu_obs, dmb), bounds=bounds, method='L-BFGS-B')

# Report
print("Fit success:", result.success)
print("Best-fit parameters:")
param_names = ['k', 'r₀', 'Ω₀', 's₀', 'α', 'β', 'γ', 'M']
for name, val in zip(param_names, result.x):
    print(f"  {name:>2} = {val:.6f}")
print(f"Minimum chi²: {result.fun:.2f}")

# Plot
mu_fit = model_mu(z, result.x)
residuals = mu_obs - mu_fit

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