import numpy as np
from common import omega_phi
from scipy.integrate import solve_ivp


def a_of_N(N, a0=1):
    return a0 * np.exp(N)

def z_of_a(a):
    return 1/a - 1

def Gamma(type, alpha=1):
    if type == 'power law':
        return (alpha+1)/alpha
    elif type == 'exponential':
        return 1
    else:
        raise ValueError("Invalid potential type. Must be 'power law' or 'exponential'.")
    
def ode(u, N, type):
    x1 = u[0]
    x2 = u[1]
    x3 = u[2]
    lmbda = u[3]

    dx1_dN = -3*x1 + np.sqrt(6)/2*lmbda*x2**2 + 1/2*x1*(3 + 3*x1**2 - 3*x2**2 + x3**2)
    dx2_dN = -np.sqrt(6)/2*lmbda*x1*x2 + 1/2*x2*(3 + 3*x1**2 - 3*x2**2 + x3**2)
    dx3_dN = -2*x3 + 1/2*x3*(3 + 3*x1**2 - 3*x2**2 + x3**2)
    dlmbda_dN = -np.sqrt(6)*lmbda**2*(Gamma(type) - 1)*x1
    return [dx1_dN, dx2_dN, dx3_dN, dlmbda_dN]

def Omega_phi(x1, x2):
    return x1**2 + x2**2

def Omega_r(x3):
    return x3**2

def Omega_m(Omega_phi, Omega_r):
    return 1 - Omega_phi - Omega_r

def problem9(N, type):
    # power law potential
    if type == 'power law':
        x1_start = 5e-5
        x2_start = 1e-8
        x3_start = 0.9999
        lmbda_start = 1e9
    elif type == 'exponential':
        zeta = 3/2
        x1_start = 0
        x2_start = 5e-13
        x3_start = 0.9999
        lmbda_start = zeta
    else:
        raise ValueError("Invalid potential type. Must be 'power law' or 'exponential'.")
    u0 = [x1_start, x2_start, x3_start, lmbda_start]

    sol = solve_ivp(lambda N, u: ode(u, N, type), [N[0], N[-1]], u0,
                    t_eval=N, rtol=1e-8, atol=1e-8)
    x1, x2, x3, lmbda = sol.y[0], sol.y[1], sol.y[2], sol.y[3]

    Omega_phi_values = Omega_phi(x1, x2)
    Omega_r_values = Omega_r(x3)
    Omega_m_values = Omega_m(Omega_phi_values, Omega_r_values)
    omega_phi_values = omega_phi(x1, x2)

    return Omega_m_values, Omega_r_values, Omega_phi_values, omega_phi_values