#!/usr/bin/env python3
"""
Test BigG's exact distance modulus calculation for specific redshifts
to compare with C implementation
"""
import numpy as np
from scipy.interpolate import interp1d

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

# Fitted parameters (from BigG)
params_reconstructed = {
    'k':    1.049342,
    'r0':   1.049676,
    'Omega0': 1.049675,
    's0':   0.994533,
    'alpha': 0.340052,
    'beta':  0.360942,
    'gamma': 0.993975,
    'H0':   70.0,
    'c0':   3303.402087,
    'M':    -19.3
}

c0_emergent = params_reconstructed['c0']
lambda_scale = 299792.458 / c0_emergent
H0 = params_reconstructed['H0']

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 H(z, k, r0, Omega0, s0, alpha, beta):
    Om_m = 0.3
    Om_de = 0.7
    Gz = G(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):
    return c0_emergent * (Omega(z, Omega0, alpha) / Omega0) ** gamma * lambda_scale

def compute_luminosity_distance_grid(z_max, params, n=500):
    k, r0, Omega0, s0, alpha, beta, gamma = params
    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)

    integrand_values = c_z / H_z

    # Trapezoidal integral
    integral_grid = np.cumsum((integrand_values[:-1] + integrand_values[1:]) / 2 * np.diff(z_grid))
    integral_grid = np.insert(integral_grid, 0, 0)

    d_c = interp1d(z_grid, integral_grid, kind='cubic', fill_value="extrapolate")

    def d_L(z):
        return (1 + z) * d_c(z)

    return d_L

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

# Test redshifts
test_z = np.array([0.010, 0.050, 0.100, 0.200, 0.300, 0.400, 0.500,
                   0.600, 0.700, 0.800, 0.900, 1.000, 1.200, 1.500])

param_list = [
    params_reconstructed['k'],
    params_reconstructed['r0'],
    params_reconstructed['Omega0'],
    params_reconstructed['s0'],
    params_reconstructed['alpha'],
    params_reconstructed['beta'],
    params_reconstructed['gamma'],
]

print("BigG Python Exact Distance Modulus Values:")
print("=" * 60)
print(f"{'z':>8} {'d_L [Mpc]':>15} {'μ [mag]':>12}")
print("=" * 60)

mu_model = model_mu(test_z, param_list)
d_L_func = compute_luminosity_distance_grid(np.max(test_z), param_list)

for i, z in enumerate(test_z):
    d_L = d_L_func(z)
    mu = mu_model[i]
    print(f"{z:8.3f} {d_L:15.3f} {mu:12.3f}")

print("=" * 60)

# Also print H(z) and c(z) values for verification
print("\nCosmological Evolution:")
print("=" * 60)
print(f"{'z':>8} {'G(z)':>12} {'c(z) [km/s]':>15} {'H(z) [km/s/Mpc]':>18}")
print("=" * 60)

test_z_cosmo = np.array([0.0, 0.5, 1.0, 1.5, 2.0])
for z in test_z_cosmo:
    Gz = G(z, *param_list[:6])
    cz = emergent_c(z, param_list[2], param_list[4], param_list[6])
    Hz = H(z, *param_list[:6])
    print(f"{z:8.3f} {Gz:12.6f} {cz:15.2f} {Hz:18.2f}")

print("=" * 60)
