# STEP 1: Install CAMB
!pip install camb

# STEP 2: Imports
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy.interpolate import interp1d
import astropy.constants as const
import astropy.units as u
import camb
from camb import model, initialpower

print("🔬 Starting detailed FST cosmological simulation with dynamic w(a)")

# STEP 3: FST Parameters from your paper
FST_PARAMS = {
    'c1': 0.51,
    'c2': -0.07,
    'c3': 0.32,
    'mV_eV': 3.2e-30,
    'lambda': 1.2e14
}

# Convert units and define constants
G = const.G.value
H0_km_s_Mpc = 72.1
H0_SI = H0_km_s_Mpc * 1000 / (3.086e22)
mV_kg = (FST_PARAMS['mV_eV'] * u.eV).to(u.kg, equivalencies=u.mass_energy()).value

# Cosmological parameters for FST
Omega_b_FST = 0.0486
Omega_r_FST = 9.2e-5

# STEP 4: FST cosmological equations system (derived from your Lagrangian - Corrected)
def fst_cosmology_system(t, y):
    """
    Corrected FST cosmological evolution equations based on your paper.
    y = [a, adot, V, Vdot]
    """
    a, adot, V, Vdot = y
    a_safe = max(a, 1e-30)
    H = adot / a_safe
    rho_V = 0.5 * FST_PARAMS['c1'] * Vdot**2 + 0.5 * mV_kg**2 * V**2 + 0.25 * FST_PARAMS['lambda'] * V**4
    p_V = 0.5 * FST_PARAMS['c1'] * Vdot**2 - 0.5 * mV_kg**2 * V**2 - 0.25 * FST_PARAMS['lambda'] * V**4
    rho_b = Omega_b_FST * (3 * H0_SI**2) / (8 * np.pi * G) / a_safe**3
    rho_r = Omega_r_FST * (3 * H0_SI**2) / (8 * np.pi * G) / a_safe**4
    addot = -(4 * np.pi * G / 3) * a_safe * (rho_b + rho_r + rho_V + 3 * p_V)
    Vddot = -3 * H * Vdot - (mV_kg**2 / FST_PARAMS['c1']) * V - (FST_PARAMS['lambda'] / FST_PARAMS['c1']) * V**3
    return [adot, addot, Vdot, Vddot]

# STEP 5: Solve the FST system with error handling
def solve_fst_cosmology_intelligent():
    methods_to_try = ['BDF', 'LSODA', 'Radau', 'RK45']
    for method in methods_to_try:
        try:
            print(f"🔄 Trying integration method: {method}")
            a0 = 1e-4
            H_initial = H0_SI * np.sqrt(Omega_b_FST/a0**3 + Omega_r_FST/a0**4)
            adot0 = a0 * H_initial
            V0 = 1e-25
            Vdot0 = 0.0
            t_span = (1e-5, 4.3e17)
            t_eval = np.logspace(np.log10(1e-5), np.log10(4.3e17), 2000)
            sol = solve_ivp(fst_cosmology_system, t_span, [a0, adot0, V0, Vdot0],
                           t_eval=t_eval, method=method, rtol=1e-8, atol=1e-10)
            if sol.success:
                print(f"✅ Success with method: {method}")
                return sol
        except Exception as e:
            print(f"❌ Method {method} error: {e}")
            continue
    print("❌ All numerical methods failed.")
    return None

sol_detailed = solve_fst_cosmology_intelligent()
if sol_detailed is None:
    print("💥 Simulation failed completely. Cannot proceed.")
    exit()

# STEP 6: Extract solutions and calculate derived quantities
a_t, adot_t, V_t, Vdot_t, t_seconds = sol_detailed.y[0], sol_detailed.y[1], sol_detailed.y[2], sol_detailed.y[3], sol_detailed.t
H_t = np.where(a_t > 1e-30, adot_t / a_t, H0_SI)
time_Gyr = t_seconds / (3.154e16)
rho_V = 0.5 * FST_PARAMS['c1'] * Vdot_t**2 + 0.5 * mV_kg**2 * V_t**2 + 0.25 * FST_PARAMS['lambda'] * V_t**4
p_V = 0.5 * FST_PARAMS['c1'] * Vdot_t**2 - 0.5 * mV_kg**2 * V_t**2 - 0.25 * FST_PARAMS['lambda'] * V_t**4
w_V = np.where(np.abs(rho_V) > 1e-30, p_V / rho_V, -1)
w_V = np.clip(w_V, -2, 1)

print("✅ FST cosmology processing completed")

# STEP 7: Create w(a) table for CAMB
a_interp = np.linspace(min(a_t), 1.0, 1000)
w_interp = interp1d(a_t, w_V, kind='linear', bounds_error=False, fill_value=(w_V[0], w_V[-1]))(a_interp)
w_table_detailed = pd.DataFrame({'a': a_interp, 'w': w_interp})

# STEP 8: Run CAMB simulation with dynamic w(a)
def run_camb_simulation():
    pars = camb.CAMBparams()
    pars.set_cosmology(
        H0=H0_km_s_Mpc,
        ombh2=Omega_b_FST * (H0_km_s_Mpc/100)**2,
        omch2=0,
        omk=0,
        tau=0.054
    )
    pars.DarkEnergy.set_w_a_table(a_interp, w_interp)
    pars.InitPower.set_params(As=2e-9, ns=0.965)
    pars.set_for_lmax(2500)
    results = camb.get_results(pars)
    powers = results.get_cmb_power_spectra(pars, CMB_unit='muK')
    return powers

try:
    powers_detailed = run_camb_simulation()
    totCL = powers_detailed['total']
    ell = np.arange(totCL.shape[0])
    D_ell_tot = totCL[:,0] * ell * (ell + 1) / (2 * np.pi)
    print("✅ CAMB simulation successful.")
except Exception as e:
    print(f"❌ CAMB simulation failed: {e}")
    ell = np.arange(2, 2500)
    D_ell_tot = np.zeros_like(ell)
    print("⚠️ Using empty CMB spectrum for plotting.")

# STEP 9: Plotting
fig = plt.figure(figsize=(15, 10))
plt.subplot(2, 2, 1)
plt.semilogx(ell, D_ell_tot, label='FST CMB Spectrum', color='darkblue', linewidth=2)
plt.xlabel(r'Multipole $\ell$')
plt.ylabel(r'$D_\ell$ [$\mu K^2$]')
plt.title('CMB Temperature Power Spectrum')
plt.grid(True, alpha=0.3)
plt.legend()
plt.subplot(2, 2, 2)
plt.plot(time_Gyr, w_V, label='FST Dynamic w(t)', color='darkred', linewidth=2)
plt.axhline(y=-1, color='navy', linestyle='--', label='ΛCDM (w=-1)', alpha=0.7)
plt.xlabel('Time (Billion years)')
plt.ylabel('Equation of State w(t)')
plt.title('FST: Dynamic Equation of State')
plt.legend()
plt.grid(True, alpha=0.3)
plt.subplot(2, 2, 3)
plt.plot(time_Gyr, a_t, label='FST Scale Factor', color='teal', linewidth=2)
plt.xlabel('Time (Billion years)')
plt.ylabel('Scale Factor a(t)')
plt.title('Cosmic Expansion History')
plt.legend()
plt.grid(True, alpha=0.3)
plt.subplot(2, 2, 4)
plt.plot(time_Gyr, V_t, label='Vector Field V(t)', color='purple', linewidth=2)
plt.xlabel('Time (Billion years)')
plt.ylabel('Vector Field V(t)')
plt.title('Fundamental Vector Field Evolution')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

# STEP 10: Final results summary
print("\n📈 FST MODEL RESULTS SUMMARY:")
print("=" * 50)
print(f"• Hubble constant: {H_t[-1] * 3.086e19:.2f} km/s/Mpc")
print(f"• Final equation of state: w = {w_V[-1]:.4f}")
print(f"• Universe age: {time_Gyr[-1]:.2f} billion years")
print("=" * 50)
print("🚀 INTELLIGENT FST SIMULATION COMPLETED SUCCESSFULLY!")
# STEP 11: Angular Peak Analysis and Comparison with Planck
def find_peak(ell, D_ell, range_min, range_max):
    mask = (ell >= range_min) & (ell <= range_max)
    peak_ell = ell[mask][np.argmax(D_ell[mask])]
    peak_amp = np.max(D_ell[mask])
    return peak_ell, peak_amp

peak1 = find_peak(ell, D_ell_tot, 100, 300)
peak2 = find_peak(ell, D_ell_tot, 400, 600)
peak3 = find_peak(ell, D_ell_tot, 700, 900)

def deviation(predicted, observed):
    return 100 * abs(predicted - observed) / observed

dev1 = deviation(peak1[0], 220)
dev2 = deviation(peak2[0], 540)
dev3 = deviation(peak3[0], 800)

print("\n🔍 Angular Peak Comparison:")
print(f"• First Peak: ℓ ≈ {peak1[0]}, Amplitude ≈ {peak1[1]:.2f} μK², Deviation from Planck: {dev1:.2f}%")
print(f"• Second Peak: ℓ ≈ {peak2[0]}, Amplitude ≈ {peak2[1]:.2f} μK², Deviation from Planck: {dev2:.2f}%")
print(f"• Third Peak: ℓ ≈ {peak3[0]}, Amplitude ≈ {peak3[1]:.2f} μK², Deviation from Planck: {dev3:.2f}%")