import numpy as np
import cosmolopy as cpy

class CosmoGrowth:

  def __init__(self, **cosmo):
    self.cosmo = cosmo

  def D(self, a):
    return cpy.perturbation.fgrowth(1/a-1, self.cosmo['omega_M_0'], unnormed=True)

  def compute_E(self, a):
    om = self.cosmo['omega_M_0']
    ol = self.cosmo['omega_lambda_0']
    ok = self.cosmo['omega_k_0']

    E = np.sqrt(om/a**3 + ol + ok/a**2)
  
    H2 = -3*om/a**4  - 2*ok/a**3

    Eprime = 0.5*H2/E

    return E,Eprime

  def Ddot(self, a):
    E,Eprime = self.compute_E(a)
    D = self.D(a)
    Ddot_D = Eprime/E + 2.5 * self.cosmo['omega_M_0']/(a**3*E**2*D)
    Ddot_D *= a
    return Ddot_D

  def compute_velmul(self, a):
    E,_ = self.compute_E(a)
    velmul = self.Ddot(a)
    velmul *= 100 * a * E
    return velmul