# Piro waveform generator
# Fragmentation of collapsar disks for SN GW radiation
# theta: colatitude
# phi: azimuth

# Version 2, 2013/08/18
# Changelog:
#  - changed mass in angular velocity from M_BH to M_BH + M_f


#!/usr/bin/python
import sys
import numpy
import pylab
from pylab import *


# constants
msun = 1.98892e33
clite = 2.9979e10
ggrav = 6.6726e-8
factor = 1.0 * ggrav / clite**4
pi = 3.14159265358979e0
trd = 1./3.

# Angular velocity courtesy of Kepler.
def Omega(M,dist):
    omega = sqrt(ggrav * M / dist**3.)
    return omega

# Mu: reduced mass
def mu(M1,M2):
    mu = M1 * M2 / (M1 + M2)
    return mu

# Chirp mass
def Mchirp(M1,M2):
    Mchirp = mu(M1,M2)**(3./5.) * (M1+M2)**(2./5.)
    return Mchirp

# Second derivative of the reduced mass quadrupole moment
def Idotdot(M1,M2,dist,omega,time):
    fact = 2. * mu(M1,M2) * dist**2. * omega**2.
    Idotdot = zeros((3,3),float)
    Idotdot[0,0] = - cos(2.0*omega*time)
    Idotdot[0,1] = - sin(2.0*omega*time)
    Idotdot[0,2] = 0.0
    Idotdot[1,0] = - sin(2.0*omega*time)
    Idotdot[1,1] =   cos(2.0*omega*time)
    Idotdot[1,2] = 0.0
    Idotdot[2,0] = 0.0
    Idotdot[2,1] = 0.0
    Idotdot[2,2] = 0.0
    Idotdot = Idotdot * fact    
    return Idotdot

# GW and viscous times
def Tgw(M1,M2,omega):
    tgw = (5./(64.*omega)) 
    tgw = tgw * (ggrav * Mchirp(M1,M2)*omega / clite**3.)**(-5./3.)
    return tgw

def Tv(alpha,eta,omega):
    tv = 1./(alpha * eta**2. * omega)
    return tv

# Evolution of the orbital radius
def Rdot(r,tgw,tv):
    rdot = - r/tgw - r/tv
    return rdot

# Functions for the RK4 integration scheme
# We evolve Risco, J, M simultaneously
def diff(t,dat):
    r = dat

    rdot = Rdot(r,Tgw(Mbh,mch,Omega(Mbh+mch,r)),Tv(alpha,eta,Omega(Mbh+mch,r)))

    ret = rdot
    return ret

def rk4(old,t,dt):
    k1 = diff(t,old)
    k2 = diff(t + 0.5*dt, old + 0.5*dt*k1)
    k3 = diff(t + 0.5*dt, old + 0.5*dt*k2)
    k4 = diff(t + dt, old + dt*k3)

    new = old + dt * (k1 + 2.*k2 + 2.*k3 + k4)/6.

    return new

# Spherical harmonics
def re_ylm_2(l,m,theta,phi):
    if l != 2:
        print "l != 2 not implemented! Good Bye!"
        sys.exit()
    if m < -2 or m > 2:
        print "m must be in [-2,2]! Good Bye!"
        sys.exit()
    if m == 0:
        ret = sqrt(15.0/32.0/pi) * sin(theta)**2
    if m == 1:
        ret = sqrt(5.0/16.0/pi) * (sin(theta) + (1.0+cos(theta))) * cos(phi)
    if m == 2:
        ret = sqrt(5.0/64.0/pi) * (1.0+cos(theta))**2 * cos(2.0*phi)
    if m == -1:
        ret = sqrt(5.0/16.0/pi) * (sin(theta) + (1.0-cos(theta))) * cos(phi)
    if m == -2:
        ret = sqrt(5.0/64.0/pi) * (1.0-cos(theta))**2 * cos(2.0*phi)
    return ret

def im_ylm_2(l,m,theta,phi):
    if l != 2:
        print "l != 2 not implemented! Good Bye!"
        sys.exit()
    if m < -2 or m > 2:
        print "m must be in [-2,2]! Good Bye!"
        sys.exit()
    if m == 0:
        ret = 0.0e0
    if m == 1:
        ret = sqrt(5.0/16.0/pi) * (sin(theta) + (1.0+cos(theta))) * sin(phi)
    if m == 2:
        ret = sqrt(5.0/64.0/pi) * (1.0+cos(theta))**2 * sin(2.0*phi)
    if m == -1:
        ret = - sqrt(5.0/16.0/pi) * (sin(theta) + (1.0-cos(theta))) * sin(phi)
    if m == -2:
        ret = - sqrt(5.0/64.0/pi) * (1.0-cos(theta))**2 * sin(2.0*phi)
    return ret

# Expansion parameters to get the (l,m) modes
def re_Hlm(l,m,Idd):
    if l != 2:
        print "l != 2 not implemented! Good Bye!"
        sys.exit()
    if m < -2 or m > 2:
        print "m must be in [-2,2]! Good Bye!"
        sys.exit()
    if m == 0:
        ret = factor * sqrt(32*pi/15.0) * \
              (Idd[2,2] - 0.5e0*(Idd[0,0] + Idd[1,1]))
    if m == 1:
        ret =  - factor * sqrt(16*pi/5) * \
              Idd[0,2]
    if m == 2:
        ret = factor * sqrt(4*pi/5) * \
              (Idd[0,0] - Idd[1,1])
    if m == -1:
        ret =  factor * sqrt(16*pi/5) * \
              Idd[0,2]
    if m == -2:
        ret = factor * sqrt(4.0*pi/5) * \
               (Idd[0,0] - Idd[1,1])
    return ret

def im_Hlm(l,m,Idd):
    if l != 2:
        print "l != 2 not implemented! Good Bye!"
        sys.exit()
    if m < -2 or m > 2:
        print "m must be in [-2,2]! Good Bye!"
        sys.exit()
    if m == 0:
        ret = 0.0e0
    if m == 1:
        ret =  factor * sqrt(16*pi/5) * \
              Idd[1,2]
    if m == 2:
        ret = factor * sqrt(4*pi/5) * \
              -2 * Idd[0,1]
    if m == -1:
        ret = +factor * sqrt(16*pi/5) * \
              Idd[1,2]
    if m == -2:
        ret =  factor * sqrt(4.0*pi/5) * \
              +2 * Idd[0,1]
    return ret    

# Useful function that gets hp, hx and returns instantaneous frequency
def instfreq(hp,hx,dt):
    # Construct dot vectors (2nd order differencing)
    hpdot = [0.]
    hxdot = [0.]
    for i in range(1,len(hp)-1):
        hpdot.append((hp[i+1]-hp[i-1])*(1./(2.*dt)))
        hxdot.append((hx[i+1]-hx[i-1])*(1./(2.*dt)))
    hpdot.append(0.)
    hxdot.append(0.)
    # Compute frequency using the fact that  
    # h(t) = A(t) e^(i Phi) = Re(h) + i Im(h) 
    freq = [0.] 
    for i in range(1,len(hp)):
        den = (hp[i]**2. + hx[i]**2)
        freq.append((-hxdot[i] * hp[i] + hpdot[i] * hx[i])/den)
    return freq


# Let's have fun
# parameters
totaltime = 8. # seconds, duration
dt = 6.1035e-05 

# physical parameters
Mbh = 5. * msun  # mass of the central BH (3-10 Msun)
#mch = 1. * msun  # mass of the fragmented blob
fac = 0.2         # 0.2-0.5
eta = 0.6         # 0.3-0.6.  H/r where H is the verical scale height of the torus
alpha = 0.1       # viscosity (< 1)
r0 = 100. * ggrav * Mbh /clite**2.   # intial radius in grav. radii G Mbh/c2
D = 3.08568e18*10*1000  # 10 kpc

risco = 6. *  ggrav * Mbh /clite**2.
rlightring = 3. *  ggrav * Mbh /clite**2.

mch = fac * (eta/0.5)**3. * Mbh/3.

# sky position
theta = 0.0
phi = 0.0

# Initial conditions
R0 = r0

olddata = R0

print '.....................................'
print 'Computing Piro waveforms'
print 'mass of BH (Msun) =', Mbh/msun
print 'mass of fragment (Msun) =', mch/msun
print 'disk with alpha =', alpha, 'and eta =', eta
print 'evolving the system for', totaltime, 'seconds'
print '.....................................'

filename = 'piroM' + str(round(Mbh/msun)) + 'eta' + str(eta) + 'fac' + str(fac) + '.dat'

f1 = open('piro.dat','w')  # File for testing and fun
f2 = open(filename,'w')   # File for actual output for xpipeline

time = 0.0

# Variables to store Risco, J, M in case we want to plot them 
t = [time]
rad = [R0]

hplus = [0.]
hcross = [0.]

# M1 = Mbh, M2 = mch
idotdot = Idotdot(Mbh,mch,R0,Omega(Mbh+mch,R0),time)

# Output file 
f1.write('Time \t Radius \t h+ \t hx')
f1.write('\n')

# Write to file initial conditions before we start evolving
f1.write('\t'.join(str(col) for col in [time, R0, hplus[0], hcross[0]]))
f1.write('\n')

f2.write('\t'.join(str(col) for col in [time, hplus[0], hcross[0]]))
f2.write('\n')

# Integrate
while (time < totaltime and R0 > 2.5*risco):
    # Evolve variables using Runge Kutta 4
    newdata = rk4(olddata,time,dt)
    time += dt

    # Store data in variables
    t.append(time)
    rad.append(newdata)

    # Actualize variable to check whether we are too close to the isco
    R0 = newdata

    # Having evolved Risco, Jbh, Mbh, compute the mass quadrupole momentum
    idotdot = Idotdot(Mbh,mch,newdata,Omega(Mbh+mch,newdata),time)
    printstring = str(time) + ' ' + str(idotdot[0,0]) + ' ' + \
    str(idotdot[0,1]) + ' ' + str(idotdot[0,2]) + ' ' + str(idotdot[1,0]) + \
    ' ' + str(idotdot[1,1]) + ' ' + str(idotdot[1,2]) + ' ' + \
    str(idotdot[2,0]) + ' ' + str(idotdot[2,1]) + ' ' + str(idotdot[2,2])
    print printstring

    # Compute gravitational radiation!
    h = 0.0
    for m in [-2,-1,0,1,2]:
        cylm = complex(re_ylm_2(2,m,theta,phi), im_ylm_2(2,m,theta,phi))
        Hlm  = complex(re_Hlm(2,m,idotdot),im_Hlm(2,m,idotdot))
        h = h + cylm * Hlm
    hp = h.real/D
    hx = h.imag/D

    hplus.append(hp)
    hcross.append(hx)

    # Print to file
    f1.write('\t'.join(str(col) for col in [time, newdata,
    newdata/(ggrav * Mbh /clite**2.), hp, hx, Omega(Mbh+mch,newdata),
    Tgw(Mbh,mch,Omega(Mbh+mch,newdata)), Tv(alpha,eta,Omega(Mbh+mch,newdata))]))
    f1.write('\n')

    # Print to xpipeline file
    f2.write('\t'.join(str(col) for col in [time, hp, hx]))
    f2.write('\n')

    olddata = newdata

f1.close()    

# Plotting results   
# Uncomment to see more plots
amp=[]
for i in range(len(hplus)):
    amp.append(sqrt(hplus[i]**2. + hcross[i]**2.))
#pylab.subplot(211)
#pylab.plot(t,hplus,color='red')
#plot(t,hcross,color='green')
#pylab.plot(t,amp,color='black')
#xlabel ('t (s)', size=16)
#ylabel ('h+, hx, |h|', size=16)

#pylab.subplot(212)
pylab.plot(t,instfreq(hplus,hcross,dt),color='black')
xlabel ('t (s)', size=16)
ylabel ('f (Hz)', size=16)

show()

print "Done"