import sys, math
import scipy.integrate as integrate
import numpy as num

# becker@astro.washington.edu

# Many of these adapted from
# astro-ph/9905116

# README :
# Version 1.1 fixes bug in self.Tfunc (lookback time); thanks to Neil Crighton for catching it!

VERSION = 1.1

class Cosmology:
    # Note, all the distance units come from Dh, so this returns everything in meters
    
    def __init__(self):
        self.Om = 0.27
        self.Ol = 0.73
        self.w  = -1
        self.h  = 0.71

        # CONSTANTS
        self.c       = 2.9979E5    # km/s
        self.G       = 6.67259E-11 # m^3 / kg / s^2
        self.Msun    = 1.98892E30  # kg
        self.pc      = 3.085677E16 # m

        # Functions for integration
        # Eqn 14 from Hogg, adding in 'w' to Ol.
        self.Efunc = lambda z, self=self : num.sqrt( (self.Om   * (1. + z)**3 +                 \
                                                       self.Ok() * (1. + z)**2 +                 \
                                                       self.Ol   * (1. + z)**(3 * (1. + self.w)) \
                                                       )**-1 )
        # Eqn 30
        self.Tfunc = lambda z, self=self : self.Efunc(z) / (1. + z)
                                           
                                           
                                           


    # Omega total
    def Otot(self):
        return self.Om + self.Ol

    # Curvature
    def Ok(self):
        return 1. - self.Om - self.Ol

    # Hubble constant, km / s / Mpc
    def H0(self):
        return 100 * self.h

    # Hubble constant at a particular epoch
    # Not sure if this is correct
    #def Hz(self, z):
    #    return self.H0() * (1. + z) * num.sqrt(1 + self.Otot() * z)

    # Got this from Jake
    def Hz(self, z):
        return self.H0 / self.Efunc(z)

    # Scale factor
    def a(self, z):
        return 1. / (1. + z)

    # Hubble distance, c / H0
    # Returns meters
    def Dh(self):
        d  = self.c / self.H0()  # km / s / (km / s / Mpc) = Mpc
        d *= self.pc * 1e6       # m
        return d

    # Hubble time, 1 / H0
    # Returns seconds
    def Th(self):
        t  = 1. / self.H0()  # Mpc s / km
        t *= self.pc * 1e3
        return t

    # Lookback time
    # Difference between the age of the Universe now and the age at z
    def Tl(self, z):
        return self.Th() * integrate.romberg(self.Tfunc, 0, z)

    # Line of sight comoving distance
    # Remains constant with epoch if objects are in the Hubble flow
    def Dc(self, z):
        return self.Dh() * integrate.romberg(self.Efunc, 0, z)

    # Transverse comoving distance
    # At same redshift but separated by angle dtheta; Dm * dtheta is transverse comoving distance
    def Dm(self, z):
        Ok  = self.Ok()
        sOk = num.sqrt(num.abs(Ok))
        Dc  = self.Dc(z)
        Dh  = self.Dh()
        
        if Ok > 0:
            return Dh / sOk * num.sinh(sOk * Dc / Dh)
        elif Ok == 0:
            return Dc
        else:
            return Dh / sOk * num.sin(sOk * Dc / Dh)

    # Angular diameter distance
    # Ratio of an objects physical transvserse size to its angular size in radians
    def Da(self, z):
        return self.Dm(z) / (1. + z)

    # Angular diameter distance between objects at 2 redshifts
    # Useful for gravitational lensing
    def Da2(self, z1, z2):
        # does not work for negative curvature
        assert(self.Ok()) >= 0

        # z1 < z2
        if (z2 < z1):
            foo = z1
            z1  = z2
            z2  = foo
        assert(z1 <= z2)

        Dm1 = self.Dm(z1)
        Dm2 = self.Dm(z2)
        Ok  = self.Ok()
        Dh  = self.Dh()

        return 1. / (1 + z2) * ( Dm2 * num.sqrt(1. + Ok * Dm1**2 / Dh**2) - Dm1 * num.sqrt(1. + Ok * Dm2**2 / Dh**2) )


    # Luminosity distance
    # Relationship between bolometric flux and bolometric luminosity
    def Dl(self, z):
        return (1. + z) * self.Dm(z)

    # Distance modulus
    # Recall that Dl is in m
    def DistMod(self, z):
        return 5. * num.log10(self.Dl(z) / self.pc / 10)

if __name__ == '__main__':
    
    c = Cosmology()

    if len(sys.argv) < 2:
        print 'Usage : cosmology.py z1 [z2]'
    z1 = float(sys.argv[1])

    print 'Cosmology : H0           =', c.H0()
    print 'Cosmology : Omega Matter =', c.Om
    print 'Cosmology : Omega Lambda =', c.Ol
    print ''
    
    print 'Hubble distance                %.2f Mpc' % (c.Dh() / c.pc / 1e6)
    print 'Hubble time                    %.2f Gyr' % (c.Th() / 3600 / 24 / 365.25 / 1e9)
    print ''
    
    print 'For z = %.2f:' % (z1)
    print 'Lookback time                  %.2f Gyr' % (c.Tl(z1) / 3600 / 24 / 365.25 / 1e9)
    print 'Scale Factor a                 %.2f'     % (c.a(z1))
    print 'Comoving L.O.S. Distance (w)   %.2f Mpc' % (c.Dc(z1) / c.pc / 1e6)
    print 'Angular diameter distance      %.2f Mpc' % (c.Da(z1) / c.pc / 1e6)
    print 'Luminosity distance            %.2f Mpc' % (c.Dl(z1) / c.pc / 1e6)
    print 'Distance modulus               %.2f mag' % (c.DistMod(z1))