# Setup ipython environment
%load_ext autoreload
%autoreload 2
%matplotlib inline

# Import useful things from kerr
from kerr.formula.ksm2_cw import CW as cwfit
from kerr.formula.ksm2_sc import SC as scfit
from kerr.pttools import leaver_workfunction as lvrwork
from kerr import leaver,rgb
from kerr.models import mmrdns 

# Setup plotting backend
import matplotlib as mpl
mpl.rcParams['lines.linewidth'] = 0.8
mpl.rcParams['font.family'] = 'serif'
mpl.rcParams['font.size'] = 12
mpl.rcParams['axes.labelsize'] = 20
from matplotlib.pyplot import *

#
from nrutils.core.nrsc import scsearch,gwylm

# Other useful things
from numpy import sin,cos,log,exp,pi,linspace,zeros,linalg,array

"""
Explanation: <center>Review Python Implementation of MMRDNS
End of explanation
"""


# Define a mass ratio
eta = mmrdns.q2eta(1)
# Define the QNM indeces 
l,m,n = 2,2,0
#
ll,mm = 3,2
# Generate the waveform and plot
hlm = mmrdns.meval_spherical_mode(ll,mm,eta,plot=True)


"""
Explanation: <center>Generate a single QNM ringdown Multipole moment
End of explanation
"""


# Define a mass ratio
eta = mmrdns.q2eta(1)
# Define Opbservation angle 
theta,phi = 0,0
# Generate the waveform and plot
h = mmrdns.meval(theta,phi,eta,plot=True)

"""
Explanation: Notes

The model looks basicall sane, and shows the expect beating behavior for the $(3,2)$ spherical mode
A detailed comparison of all spherical modes to select calibration NR runs is performed in a separate notebook

<center>Get model for waveform, and plot internally
End of explanation
"""


#
A = scsearch(keyword='sxs',q=2,nonspinning=True,verbose=True,unique=True)[0]

#
y = gwylm(A,lmax=3,verbose=True,clean=True)

"""
Explanation: Notes

Here we have evaluated to inherently spheroidal model:
The model contains fits for the true spheroidal amplitudes for nonspinning BBH mergers
Spherical represenations only contain a portion of this spheroidal information (though surely most of it)
We may optionally compare the model's deption of how spherical and spheroidal depictions match up.


Question --- How does MMRDNS' spheroidal represenation comapare to the NR represenation from partial spherical recomposition?

<center>Load NR waveform
End of explanation
"""


# Define the mode for comparison
l,m=3,2;n=0
# Define a mass ratio
eta = mmrdns.q2eta( y.m1/y.m2 )
# Get the ringdown data to compare
nr = y.ringdown().lm[(l,m)]['strain'] # NOTE that the default reference time for the ringdown() mthod is the same one used in mmrdns: 10M after the peak luminosity as defined by the l=m=2 mode
ar = mmrdns.meval_spherical_mode(l,m,eta,gwfout=True)( nr.t )

ax,fig = ar.plot(show=False)
from kerr.basics import lim
from numpy import array,diff

sca(ax[2])
w0,w1 = mmrdns.cwfit(l,m,n,eta).real, mmrdns.cwfit(2,m,n,eta).real

clr = 'k'

axhline( w0, linestyle=':', color=clr, alpha=0.8 )
text( diff(lim(ar.t))*0.95, w0-0.05, '(%i,%i)'%(l,m), color=clr )

axhline( w1, linestyle='--', color=clr, alpha=0.8 )
text( diff(lim(ar.t))*0.95, w1-0.05, '(%i,%i)'%(2,m), color=clr )

ylim( min(ylim()), max( max(ylim()),1.1*max(w0,w1) ) )
show()

nr.plot()

# Plot the amplitudes
fig = figure( figsize=[16,10] )
plot( nr.t, nr.amp ) 
plot( ar.t, ar.amp,'r' ) 
gca().set_yscale("log", nonposy='clip')
xlim([10,100]); ylim(1e-5,1);
title('There is pretty good agreement near the start. Note that there is overtone information in the NR data that contributes to the difference.')

"""
Explanation: Compare the l=m=2 sherical multipole with MMRDNS' shpheroidal multipole (They should be very similar!)
End of explanation
"""


theta=0;phi=0
nr_ = y.ringdown().recompose(theta,phi,kind='strain')
ar_ = mmrdns.meval(theta,phi,eta,gwfout=True)( nr_.t )

# 
from numpy import sqrt
figure( figsize=8*array((2,1)) )
plot( nr_.t, nr_.amp, 'k', lw=4, alpha=0.25 )
plot( nr_.t, ar_.amp, 'r'  )
ax = gca()
xlim([10,75])
ylim([max(ylim())/200,max(ylim())])
# ax.set_yscale("log", nonposy='clip')

"""
Explanation: Notes

Here, it is seen that the NR sherical mode and the spherical mode from MMRDNS agree very well
But how do the recomposed data comapre?

Now compare the recomposed data
End of explanation
"""


