#!/usr/bin/env python
# coding: utf-8

# <script async src="https://www.googletagmanager.com/gtag/js?id=UA-59152712-8"></script>
# <script>
#   window.dataLayer = window.dataLayer || [];
#   function gtag(){dataLayer.push(arguments);}
#   gtag('js', new Date());
# 
#   gtag('config', 'UA-59152712-8');
# </script>
# 
# # The Spinning Effective One-Body Factorized Modes
# 
# ## Author: Tyler Knowles
# 
# ## This module documents the reduced spinning effective one-body factorized modes as numerically implemented in LALSuite's SEOBNRv3 gravitational waveform approximant.  That is, we follow the appendix of of [Taracchini, et. al. (2012)](https://arxiv.org/abs/1202.0790).
# 
# 
# **Notebook Status:** <font color='red'><b> In progress </b></font>
# 
# **Validation Notes:** This module is under active development -- do ***not*** use the resulting code for scientific applications.  In the future, this module will be validated against the LALSuite [SEOBNRv3/SEOBNRv3_opt code]( https://git.ligo.org/lscsoft/lalsuite.) that was reviewed and approved for LIGO parameter estimation by the LIGO Scientific Collaboration.
# 
# 
# ## Introduction
# ### The Physical System of Interest
# 
# Consider two compact objects (e.g. black holes or neutron stars) with masses $m_{1}$, $m_{2}$ (in solar masses) and spin angular momenta ${\bf S}_{1}$, ${\bf S}_{2}$ in a binary system.  The spinning effective one-body ("SEOB") Hamiltonian $H_{\rm real}$ (see [BB2010](https://arxiv.org/abs/0912.3517) Equation (5.69)) describes the dynamics of this system.  We seek to computed the factorized modes for this system, as described in the appendix of [T2012](https://arxiv.org/abs/1202.0790).
# 
# This routine takes as input
# 1. the masses $m_{1}$, $m_{2}$,
# 1. the symmetic mass ratio $\eta$ and Kerr spin parameter $a$, and
# 1. dimensionless spin parameters $\chi_{A}$, $\chi_{S}$.
# 
# ### Citations
# Throughout this module, we will refer to
# * [Taracchini, et. al. (2012)](https://arxiv.org/abs/1202.0790) as T2012.

# <a id='toc'></a>
# 
# # Table of Contents
# $$\label{toc}$$
# 
# This notebook is organized as follows:
# 
# 1. [Step 0](#outputcreation): Creating the output directory for SEOBNR
# 1. [Step 1](#commonterms): Common terms
# 1. [Step 2](#fmodes): Factorized modes
# 1. [Step 3](#validation): Validation
# 1. [Step 4](#latex_pdf_output): Output this notebook to $\LaTeX$-formatted PDF file

# <a id='outputcreation'></a>
# 
# # Step 0: Creating the output directory for SEOBNR \[Back to [top](#toc)\]
# $$\label{outputcreation}$$
# 
# First we create the output directory for SEOBNR (if it does not already exist):

# In[1]:


# Get cmdline_helper; remove when this notebook is finalized and removed from in_progress
import sys
# insert at 1, 0 is the script path (or '' in REPL)
sys.path.insert(1, '../')

import cmdline_helper as cmd     # NRPy+: Multi-platform Python command-line interface

# Create C code output directory:
Ccodesdir = "SEOBNR"
# Then create an output directory in case it does not exist
cmd.mkdir(Ccodesdir)


# <a id='commonterms'></a>
# 
# # Step 1: Common terms \[Back to [top](#toc)\]
# $$\label{commonterms}$$
# 
# We define the following terms that are used repeatedly in the factorized modes.  From [T2012](https://arxiv.org/abs/1202.0790) Equation (A7a):
# 
# \begin{equation*}
#     \delta m = \frac{ m_{1} - m_{2} }{ m_{1} + m_{2} }.
# \end{equation*}

# In[2]:


get_ipython().run_cell_magic('writefile', '$Ccodesdir/Factorized_modes-expressions.txt', 'deltam = (m1 - m2)/(m1 + m2)\ndeltam2 = deltam*deltam\n\na2 = a*a\na3 = a2*a\na4 = a3*a\na5 = a4*a\na6 = a5*a\na7 = a6*a\n\neta2 = eta*eta\neta3 = eta2*eta\neta4 = eta3*eta\n\nm1p3eta = -1 + 3*eta\n')


# <a id='fmodes'></a>
# 
# # Step 2: The factorized modes \[Back to [top](#toc)\]
# $$\label{fmodes}$$
# 
# We will write the coefficient (on each power of $v$) separately, following the formulation in [T2012](https://arxiv.org/abs/1202.0790) Equation (A8a):
# 
# <font color='red'>Need to find a reference for the 1./2. * (chiS + chiA * deltam) term in rho22v4, the term rho22v5, the (89. * a * a)/252. term of rho22v6 (and why we don't include 428 * eulerlog2(v2)/105 as in T2012)... these terms are not properly documented in the LALSuite code.</font>

# In[3]:


get_ipython().run_cell_magic('writefile', '-a $Ccodesdir/Factorized_modes-expressions.txt', '\nrho22v2 = 55.*eta/84. - 43./42\nrho22v3 = -2.*(chiS*(1 - eta) + chiA*deltam)/3.\nrho22v4 = 1./2.*(chiS + chiA*deltam)*(chiS + chiA*deltam) + 19583.*eta2/42336. - 33025.*eta/21168. - 20555./10584.\nrho22v5 = -34.*a/21.\nrho22v6 = (1556919113./122245200. + (89.*a2)/252. - (48993925.*eta)/9779616. - (6292061.*eta2)/3259872.\n           + (10620745.*eta3)/39118464. + (41.*eta*sp.pi*sp.pi)/192.)\nrho22v6l = -428./105.\nrho22v7 = (18733.*a)/15876. + a3/3.\nrho22v8 = -387216563023./160190110080. + (18353.*a2)/21168. - a*4/8.\nrho22v8l = 9202./2205.\nrho22v10 = -16094530514677./533967033600.\nrho22v10l = 439877./55566.\n\nrho21v1 = 0.\nrho21v2 = -59./56. + (23.*eta)/84.\nrho21v3 = 0.\nrho21v4 = -47009./56448. - (865.*a2)/1792. - (405.*a4)/2048. - (10993.*eta)/14112. + (617.*eta2)/4704.\nrho21v5 = (-98635.*a)/75264. + (2031.*a3)/7168. - (1701.*a5)/8192.\nrho21v6 = 7613184941./2607897600. + (9032393.*a2)/1806336. + (3897.*a4)/16384. - (15309.*a6)/65536.\nrho21v6l = -107./105.\nrho21v7 = (-3859374457.*a)/1159065600. - (55169.*a3)/16384. + (18603.*a5)/65536. - (72171.*a7)/262144.\nrho21v7l = 107.*a/140.\nrho21v8 = -1168617463883./911303737344.\nrho21v8l = 6313./5880.\nrho21v10 = -63735873771463./16569158860800.\nrho21v10l = 5029963./5927040.\n\nf21v1 = (-3.*(chiS + chiA/deltam))/2.\nf21v3 = (chiS*deltam*(427. + 79.*eta) + chiA*(147. + 280.*deltam2 + 1251.*eta))/84./deltam\n\nrho33v2 = -7./6. + (2.*eta)/3.\nrho33v3 = 0.\nrho33v4 = -6719./3960. + a2/2. - (1861.*eta)/990. + (149.*eta2)/330.\nrho33v5 = (-4.*a)/3.\nrho33v6 = 3203101567./227026800. + (5.*a2)/36.\nrho33v6l = -26./7.\nrho33v7 = (5297.*a)/2970. + a3/3.\nrho33v8 = -57566572157./8562153600.\nrho33v8l = 13./3.\n\nf33v3 = (chiS*deltam*(-4. + 5.*eta) + chiA*(-4. + 19.*eta))/(2.*deltam)\n\nrho32v = (4.*chiS*eta)/(-3.*m1p3eta)\nrho32v2 = (-4.*a2*eta2)/(9.*m1p3eta*m1p3eta) + (328. - 1115.*eta + 320.*eta2)/(270.*m1p3eta)\nrho32v3 = 2./9.*a\nrho32v4 = a2/3. + (-1444528. + 8050045.*eta - 4725605.*eta2 - 20338960.*eta3\n                   + 3085640.*eta4)/(1603800.*m1p3eta*m1p3eta)\nrho32v5 = (-2788.*a)/1215.\nrho32v6 = 5849948554./940355325. + (488.*a2)/405.\nrho32v6l = -104./63.\nrho32v8 = -10607269449358./3072140846775.\nrho32v8l = 17056./8505.\n\nrho31v2 = -13./18. - (2.*eta)/9.\nrho31v3 = 0.0\nrho31v4 = 101./7128. - (5.*a2)/6. - (1685.*eta)/1782. - (829.*eta2)/1782.\nrho31v5 = (4.*a)/9.\nrho31v6 = 11706720301./6129723600. - (49.*a2)/108.\nrho31v6l = -26./63.\nrho31v7 = (-2579.*a)/5346. + a3/9.\nrho31v8 = 2606097992581./4854741091200.\nrho31v8l = 169./567.\n\nf31v3 = (chiA * (-4. + 11.*eta) + chiS*deltam*(-4. + 13.*eta))/(2.*deltam)\n\nrho44v2 = (1614. - 5870.*eta + 2625.*eta2)/(1320.*m1p3eta)\nrho44v3 = (chiA * (10. - 39.*eta)*deltam + chiS * (10. - 41.*eta + 42.*eta2)) / (15.*m1p3eta)\nrho44v4 = a2/2. + (-511573572. + 2338945704.*eta - 313857376.*eta2 - 6733146000. * eta3 +\n                   1252563795.*eta4)/(317116800.*np.power(m1p3eta,2))\nrho44v5 = (-69.*a)/55.\nrho44v6 = 16600939332793./1098809712000. + (217.*a2)/3960.\nrho44v6l = -12568./3465.\n\nrho43v = 0.\nrho43v2 = (222. - 547.*eta + 160.*eta2)/(176.*(-1. + 2.*eta))\nrho43v4 = -6894273./7047040. + (3.*a2)/8.\nrho43v5 = (-12113.*a)/6160.\nrho43v6 = 1664224207351./195343948800.\nrho43v6l = -1571./770.\n\nf43v = (5.*(chiA - chiS*deltam)*eta)/(2.*deltam*(-1. + 2.*eta))\n\nrho42v2 = (1146. - 3530.*eta + 285.*eta2)/(1320.*m1p3eta)\nrho42v3 = (chiA * (10. - 21.*eta)*deltam + chiS*(10. - 59.*eta + 78.*eta2))/(15.*m1p3eta)\nrho42v4 = a2/2. + (-114859044. + 295834536.*eta + 1204388696.*eta2 - 3047981160.*eta3\n                   - 379526805.*eta4)/(317116800.*m1p3eta*m1p3eta)\nrho42v5 = (-7.*a)/110.\nrho42v6 = 848238724511./219761942400. + (2323.*a2)/3960.\nrho42v6l = -3142. / 3465.\n\nrho41v = 0.0\nrho41v2 = (602. - 1385.*eta + 288.*eta2)/(528.*(-1. + 2.*eta))\nrho41v4 = -7775491./21141120. + (3.* a2)/8.\nrho41v5 = (-20033.*a)/55440. - (5*a3)/6.\nrho41v6 = 1227423222031./1758095539200.\nrho41v6l = -1571. / 6930.\n\nf41v = (5.*(chiA - chiS*deltam)*eta)/(2.*deltam*(-1. + 2.*eta))\n\nrho55v2 = (487. - 1298.*eta + 512.*eta2)/(390.*(-1. + 2.*eta))\nrho55v3 = (-2.*a)/3.\nrho55v4 = -3353747./2129400. + a2/2.\nrho55v5 = -241.*a/195.\n\nrho54v2 = (-17448. + 96019.*eta - 127610.*eta2 + 33320.*eta3)/(13650.*(1. - 5.*eta + 5.*eta2))\nrho54v3 = (-2.*a)/15.\nrho54v4 = -16213384./15526875. + (2.*a2)/5.\n\nrho53v2 = (375. - 850.*eta + 176.*eta2)/(390.*(-1. + 2.*eta))\nrho53v3 = (-2.*a)/3.\nrho53v4 = -410833./709800. + a2/2.\nrho53v5 = -103.*a/325.\n\nrho52v2 = (-15828. + 84679.*eta - 104930.*eta2 + 21980.*eta3)/(13650.*(1. - 5.*eta + 5.*eta2))\nrho52v3 = (-2.*a)/15.\nrho52v4 = -7187914./15526875. + (2.*a2)/5.\n\nrho51v2 = (319. - 626.*eta + 8.*eta2)/(390.*(-1. + 2.*eta))\nrho51v3 = (-2.*a)/3.\nrho51v4 = -31877./304200. + a2/2.\nrho51v5 = 139.*a/975.\n\nrho66v2 = (-106. + 602.*eta - 861.*eta2 + 273.*eta3)/(84.*(1. - 5.*eta + 5.*eta2))\nrho66v3 = (-2.*a)/3.\nrho66v4 = -1025435./659736. + a2/2.\n\nrho65v2 = (-185. + 838.*eta - 910.*eta2 + 220.*eta3)/(144.*(deltam2 + 3.*eta2))\nrho65v3 = -2.*a/9.\n\nrho64v2 = (-86. + 462.*eta - 581.*eta2 + 133.*eta3)/(84.*(1. - 5.*eta + 5.*eta2))\nrho64v3 = (-2.*a)/3.\nrho64v4 = -476887./659736. + a2/2.\n\nrho63v2 = (-169. + 742.*eta - 750.*eta2 + 156.*eta3)/(144.*(deltam2 + 3.*eta2))\nrho63v3 = -2.*a/9.\n\nrho62v2 = (-74. + 378.*eta - 413.*eta2 + 49.*eta3)/(84.*(1. - 5.*eta + 5.*eta2))\nrho62v3 = (-2.*a)/3.\nrho62v4 = -817991./3298680. + a2/2.\n\nrho61v2 = (-161. + 694.*eta - 670.*eta2 + 124.*eta3)/(144.*(deltam2 + 3.*eta2))\nrho61v3 = -2.*a/9.\n\nrho77v2 = (-906. + 4246.*eta - 4963.*eta2 + 1380.*eta3)/(714.*(deltam2 + 3.*eta2))\nrho77v3 = -2.*a/3.\n\nrho76v2 = (2144. - 16185.*eta + 37828.*eta2 - 29351.*eta3 + 6104.*eta4)/(1666.*(-1 + 7*eta - 14*eta2 + 7 * eta3))\n\nrho75v2 = (-762. + 3382.*eta - 3523.*eta2 + 804.*eta3)/(714.*(deltam2 + 3.*eta2))\nrho75v3 = -2.*a/3.\n\nrho74v2 = (17756. - 131805.*eta + 298872.*eta2 - 217959.*eta3\n           + 41076.*eta4)/(14994.*(-1. + 7.*eta - 14.*eta2 + 7. * eta3))\n\nrho73v2 = (-666. + 2806.*eta - 2563.*eta2 + 420.*eta3)/(714.*(deltam2 + 3.*eta2))\nrho73v3 = -2.*a/3.\n\nrho72v2 = (16832. - 123489.*eta + 273924.*eta2 - 190239.*eta3\n           + 32760.*eta4)/(14994.*(-1. + 7.*eta - 14.*eta2 + 7.*eta3))\n\nrho71v2 = (-618. + 2518.*eta - 2083.*eta2 + 228.*eta3)/(714.*(deltam*deltam + 3.*eta2))\nrho71v3 = -2.*a/3.\n\nrho88v2 = (3482. - 26778.*eta + 64659.*eta2 - 53445.*eta3 + 12243.*eta4)/(2736.*(-1. + 7.*eta - 14.*eta2 + 7.*eta3))\n\nrho87v2 = (23478. - 154099.*eta + 309498.*eta2 - 207550.*eta3 + 38920*eta4)/(18240.*(-1 + 6*eta - 10*eta2 + 4*eta3))\n\nrho86v2 = (1002. - 7498.*eta + 17269.*eta2 - 13055.*eta3 + 2653.*eta4)/(912.*(-1. + 7.*eta - 14.*eta2 + 7.*eta3))\n\nrho85v2 = (4350. - 28055.*eta + 54642.*eta2 - 34598.*eta3 + 6056.*eta4)/(3648.*(-1. + 6.*eta - 10.*eta2 + 4.*eta3))\n\nrho84v2 = (2666. - 19434.*eta + 42627.*eta2 - 28965.*eta3 + 4899.*eta4)/(2736.*(-1. + 7.*eta - 14.*eta2 + 7.*eta3))\n\nrho83v2 = (20598. - 131059.*eta + 249018.*eta2 - 149950.*eta3 + 24520.*eta4)/(18240.*(-1. + 6.*eta - 10.*eta2\n                                                                                      + 4.*eta3))\n\nrho82v2 = (2462. - 17598.*eta + 37119.*eta2 - 22845.*eta3 + 3063.*eta4)/(2736.*(-1. + 7.*eta - 14.*eta2 + 7.*eta3))\n\nrho81v2 = (20022.-126451.*eta+236922.*eta2 - 138430.*eta3 + 21640.*eta4)/(18240.*(-1. + 6.*eta - 10.*eta2 + 4.*eta3))\n')


# <a id='validation'></a>
# 
# # Step 3: Validation \[Back to [top](#toc)\]
# $$\label{validation}$$
# 
# The following code cell reverses the order of the expressions output to SEOBNR/Hamiltonian_on_top.txt and creates a Python function to validate the value of $H_{\rm real}$ against the SEOBNRv3 Hamiltonian value computed in LALSuite git commit bba40f21e9 for command-line input parameters
# 
# -f 20 -M 23 -m 10 -f 20 -X 0.01 -Y 0.02 -Z -0.03 -x 0.04 -y -0.05 -z 0.06.

# In[4]:


import numpy as np
import difflib, os

# The modes are sometimes written on more than one line for readability in this Jupyter notebook.
# We first create a file of one-line expressions, SEOBNR-one_line_factorized_modes.txt.
with open(os.path.join(Ccodesdir,"Factorized_modes-one_line_expressions.txt-VALIDATION"), "w") as output:
    count = 0
    # Read output of this notebook
    for line in list(open(os.path.join(Ccodesdir,"Factorized_modes-expressions.txt"))):
        # Read the first line
        if count == 0:
            prevline=line
        #Check if prevline is a complete expression
        elif "=" in prevline and "=" in line:
            output.write("%s\n" % prevline.strip('\n'))
            prevline=line
        # Check if line needs to be adjoined to prevline
        elif "=" in prevline and not "=" in line:
            prevline = prevline.strip('\n')
            prevline = (prevline+line).replace(" ","")
        # Be sure to print the last line.
        if count == len(list(open(os.path.join(Ccodesdir,"Factorized_modes-expressions.txt"))))-1:
            if not "=" in line:
                print("ERROR. Algorithm not robust if there is no equals sign on the final line. Sorry.")
                sys.exit(1)
            else:
                output.write("%s" % line)
        count = count + 1

# Now write the expressions in a function
# We want to return each mode, so we will store them in a dictionary
with open(os.path.join(Ccodesdir,"Factorized_modes.py"), "w") as output:
    output.write("import numpy as np\ndef compute_modes(m1=23., m2=10., a=-9.086823027135883e-03, eta=2.112029384756657e-01, chiA=-4.516044029838846e-02, chiS=1.506880542362110e-02):\n\tmodes={}\n")
    for line in list(open(os.path.join(Ccodesdir,"Factorized_modes-one_line_expressions.txt"))):
        # To prevent storing common terms in the dictionary, sort expressions by those containing 'v'
        if 'v' in line:
            # Split the line so that the mode name can be used to create a dictionary entry
            name, expression = line.split('=') # Should not be more than one '=' per line!
            output.write("\tmodes['%s'] = %s\n" % (name.strip(),expression.rstrip().replace("sp.pi","np.pi")))
        else:
            # For expressions not in the dictionary, just print the expression as-is
            output.write("\t%s\n" % line.rstrip().replace("sp.pi","np.pi"))
    output.write("\treturn modes")
    # For testing
    output.write("\ntest=compute_modes()\n")
    output.write("for key, value in test.items():\n\tprint(\"%s = %.15e\" % (key,value))")


# Check that expressions have not been altered
print("Printing difference between notebook output and a trusted list of expressions...")
# Open the files to compare
file = "Factorized_modes-one_line_expressions.txt"
outfile = "Factorized_modes-one_line_expressions.txt-VALIDATION"

print("Checking file " + outfile)
with open(os.path.join(Ccodesdir,file), "r") as file1, open(os.path.join(Ccodesdir,outfile), "r") as file2:
    # Read the lines of each file
    file1_lines=[]
    file2_lines=[]
    for line in file1.readlines():
        file1_lines.append(line.replace(" ", ""))
    for line in file2.readlines():
        file2_lines.append(line.replace(" ", ""))
    num_diffs = 0
    for line in difflib.unified_diff(file1_lines, file2_lines, fromfile=os.path.join(Ccodesdir,file), tofile=os.path.join(Ccodesdir,outfile)):
        sys.stdout.writelines(line)
        num_diffs = num_diffs + 1
    if num_diffs == 0:
        print("No difference. TEST PASSED!")
    else:
        print("ERROR: Disagreement found with the trusted file. See differences above.")
        sys.exit(1)


# <a id='latex_pdf_output'></a>
# 
# # Output this notebook to $\LaTeX$-formatted PDF file \[Back to [top](#toc)\]
# $$\label{latex_pdf_output}$$
# 
# The following code cell converts this Jupyter notebook into a proper, clickable $\LaTeX$-formatted PDF file. After the cell is successfully run, the generated PDF may be found in the root NRPy+ tutorial directory, with filename
# [Tutorial-SEOBNR_Factorized_Modes.pdf](Tutorial-SEOBNR_Factorized_Modes.pdf) (Note that clicking on this link may not work; you may need to open the PDF file through another means.)

# In[5]:


cmd.output_Jupyter_notebook_to_LaTeXed_PDF("Tutorial-SEOBNR_Factorized_Modes")