import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import requests
import io

# URL for Pantheon data (hosted on GitHub as CSV)
pantheon_url = 'https://raw.githubusercontent.com/dscolnic/Pantheon/master/pantheon_clean.csv'

response = requests.get(pantheon_url)
response.raise_for_status()  # ensure we got a good response
pantheon_data = pd.read_csv(io.StringIO(response.text))


# Extract redshift, distance modulus, and error
z_pantheon = pantheon_data['zHD'].values
mu_pantheon = pantheon_data['mb'].values  # m_b is the corrected magnitude, proportional to distance modulus
mu_err = pantheon_data['dmb'].values

# Note: m_b from Pantheon is apparent magnitude. To compare with distance modulus μ,
# we need absolute magnitude M to convert:
# μ = m_b - M
# For comparison, shift your model curve by an additive constant to match roughly

# Define a function to shift model curve to data (fit constant offset)
def best_fit_offset(mu_model, z_model, z_data, mu_data):
    # Interpolate model at data redshifts
    from scipy.interpolate import interp1d
    f = interp1d(z_model, mu_model, kind='cubic', fill_value='extrapolate')
    mu_interp = f(z_data)
    offset = np.mean(mu_data - mu_interp)
    return offset

# (Reuse your model arrays here or compute again)
# For example:
# z_array, mu_model from previous model calculation

# -- For demonstration, dummy example --
# Replace this with your actual arrays from the previous model
# For now, let's create dummy smooth curve resembling ΛCDM
z_array = np.linspace(0.01, 2, 100)
mu_model = 5 * np.log10((1 + z_array) * z_array * 3000) + 25  # simplistic approx

# Find offset between model and Pantheon data for fair comparison
offset = best_fit_offset(mu_model, z_array, z_pantheon, mu_pantheon)
mu_model_shifted = mu_model + offset

# Plotting
plt.errorbar(z_pantheon, mu_pantheon, yerr=mu_err, fmt='.', label='Pantheon SNe', alpha=0.5)
plt.plot(z_array, mu_model_shifted, label='Emergent Gravity Model', color='red')
plt.xlabel('Redshift z')
plt.ylabel('Distance Modulus μ')
plt.title('Pantheon Supernovae vs Emergent Gravity Model')
plt.legend()
plt.grid(True)
plt.show()
