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

# --- Your cosmological model parameters ---
k = 1.0
r0 = 1.0
Omega0 = 1.0
s0 = 1.0

# Assume a simple scale factor relation for Omega(t) and s(t)
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):
    # Use cosmological time scaling approx t ~ 1/H0 * (1+z)^(-3/2)
    # Just a rough inverse power for demonstration
    return s0 * (1+z)**(-beta)

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

# Simplified Hubble parameter for flat universe (matter + effective dark energy)
def H(z):
    # Matter density fraction
    Om_m = 0.3
    Om_de = 0.7
    Gz = G(z)
    H0 = 70  # km/s/Mpc
    
    # Effective Hubble squared: include G(z)
    Hz_sq = H0**2 * (Om_m * Gz * (1+z)**3 + Om_de)
    return np.sqrt(Hz_sq)

# Comoving distance integral
def comoving_distance(z):
    z_array = np.linspace(0, z, 500)
    Hz_array = H(z_array)
    integrand = 1 / Hz_array
    d_c = np.trapz(integrand, z_array)
    c = 299792.458  # speed of light in km/s
    return c * d_c

# Luminosity distance
def luminosity_distance(z):
    return (1 + z) * comoving_distance(z)

# Distance modulus
def distance_modulus(z):
    d_L = luminosity_distance(z)  # in Mpc
    return 5 * np.log10(d_L) + 25  # standard formula

# Calculate model μ(z) on data redshifts
mu_model = np.array([distance_modulus(zi) for zi in z])

# Now plot 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.legend()
plt.grid(True)
plt.title('Data vs. Emergent Gravity Model')
plt.show()
