import pylab as pl

Nit      = int(1e6) # no. of iterations
H0       = 1.#2.269e-18 # /s


H_f          = lambda a, Omg_m0: H0*( 1. + Omg_m0/(a**3.) - Omg_m0 )**0.5

delta_ddt_f  = lambda a, H, Omg_m0, delta, delta_dt:                \
                      1.5*(H**2.)*(1. + Omg_m0 / (a**3.) - Omg_m0 ) \
                    - 2.*H*delta_dt

def sse2():
    """
    second sub exercise:
    """
    constant = 1.
    dt       = 1e12                           # 1. / (Nit-1.) # for example?
    counter  = pl.zeros(3, dtype=int)
    Omg_0s   = pl.array([ [1., 0.], [0.3, 0.7], [0.8, 0.2] ]) # 3 cases

    a_list     = []
    delta_list = []

    for u in pl.arange(0, 3):
        " for every case of the densities "

        Omg_m0 = Omg_0s[u,0]
        H_A           = pl.zeros( Nit )

        a_A           = pl.zeros( Nit )
        adot_A        = pl.zeros( Nit )
        a_A[0]        = 1.e-3
        adot_A[0]     = a_A[0] * H_f(a_A[0], Omg_m0)

        delta_A       = pl.zeros( Nit )
        delta_dt_A    = pl.zeros( Nit )
        delta_dt_A[0] = constant*a_A[0]

        for i in pl.arange(0, Nit-1):
            " integrating in time "

            if a_A[i] <= 1:
                " not reached today yet "
                H_A[i+1]    = H_f(a_A[i], Omg_m0)
                a_A[i+1]    = a_A[i] \
                                + dt*adot_A[i] # Forward-Euler
                adot_A[i+1] = a_A[i+1] * H_A[i+1]

                delta_dt_A[i+1] = delta_dt_A[i] \
                + dt*delta_ddt_f(a_A[i], H_A[i], Omg_m0, delta_A[i], delta_dt_A[i])

                delta_A[i+1] = delta_A[i] \
                                + dt*delta_dt_A[i+1]  # Euler-Cromer
                pass
            else:
                " a > 1 "
                a_list.append(     a_A [:i]    )
                delta_list.append( delta_A[:i] )
                break # breaks out of innermost for-loop

            continue

        continue

    return a_list, delta_list