import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad

# Your data loading part (unchanged)
filename = 'hlsp_ps1cosmo_panstarrs_gpc1_all_model_v1_lcparam-full.txt'

lc_data = np.genfromtxt(
    filename,
    delimiter=' ',            # space-separated columns
    names=True,
    comments='#',
    dtype=None,
    encoding=None
)

z = lc_data['zcmb']
mb = lc_data['mb']
dmb = lc_data['dmb']

M = -19.3
mu = mb - M

# --- Emergent Gravity Model Setup ---

# Constants and parameters (tune these)
k = 1.0
r0 = 1.0
Omega0 = 1.0
s0 = 1.0
H0 = 70.0  # Hubble constant in km/s/Mpc
c = 299792.458  # speed of light km/s

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

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

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

def G(z):
    return Omega(z) * k**2 * r0 / s(z)

def H(z):
    Om_m = 0.3
    Om_de = 0.7
    Gz = G(z)
    Hz_sq = (H0 ** 2) * (Om_m * Gz * (1 + z) ** 3 + Om_de)
    return np.sqrt(Hz_sq)

def integrand(z_prime):
    return c / H(z_prime)

def luminosity_distance(z_val):
    integral, _ = quad(integrand, 0, z_val)
    return (1 + z_val) * integral

def distance_modulus(z_val):
    d_L = luminosity_distance(z_val)  # in Mpc
    return 5 * np.log10(d_L) + 25

# Compute model mu(z) for all data redshifts
mu_model = np.array([distance_modulus(zi) for zi in z])

# --- Plotting both data and model ---
plt.errorbar(z, mu, yerr=dmb, fmt='.', alpha=0.5, label='Pan-STARRS1 SNe')
plt.plot(z, mu_model, color='red', label='Emergent Gravity Model')
plt.xlabel('Redshift (z)')
plt.ylabel('Distance Modulus (μ)')
plt.title('Supernova Data vs. Emergent Gravity Model')
plt.legend()
plt.grid(True)
plt.show()
