############ CONSTANTS AND CHOICE OF QNMs ######################
# Physical units
TSUN = 4.9169e-6 #Mass of the sun converted in seconds GM/c**3
clight=2.99792458e8 #m/s
Mpc=3.0857e22 #m

#epsilon parameter for conditioning fisher matrix
epsiloncond=1E-9

#Frequency bounds for LISA
fmin = 1E-5
fmax = 1.

# Which modes to include in the waveform and over which to compute the Fisher matrix. 
# Possible choices are 220, 221, 210, 211, 330, 331, 320, 440 and the retrograde modes r220, r210.
# If we include them all, the dimension of the Fisher is 20*20 (4 params for each mode).
linearModes = ['220', '221', '210', '330', '440', '211', '320', '331', 'r220', 'r210'] 
# linearModes = ['220', '221', '210', '330', '440'] 
# linearModes = ['220', '210', '330', '440'] 

# Parameter controlling the inclusion of quadratic modes
# 0 for not including nonlinear modes at all in the model, 
# 1: include the dominant QQNM both in the template and in the parameter estimation from Fisher,
# 2: include the dominant QQNM both in the template but not in the parameter estimation from Fisher (i.e. we assume we know this quadratic mode)
# The current implementation for 2 uses a fit from NR for the amplitude of QQNM
# 3: All nonlinear modes coming from interactions among all linear modes will be included in the waveform template, but not in parameter estimation from Fisher (i.e. we assume we know all quadratic modes). 
# The current implementation of 3 uses QQNMs for Schwarzschild BH computed in perturbation theory in 2405.06012, 2406.14611
includeNL = 1
nbLinearModes = len(linearModes)



#################### SYMPY DEFINITIONS ################

### First we need to define all params of the fisher: amplitude, phase, real and imaginary part of omega for each mode. We use sympy for symbolic derivatives.
# cf notes for the definition of angles:
# theta, phi: angles of the GW event as seen from Earth
# psi: polarization angle
# iota, phiSH: angles in the BBH system of reference, used to define spherical harmonics
M, theta, phi, psi, phiSH, iota, dl, f = symbols('M, theta, phi, psi, phiSH, iota, dl, f', real=True)

paramList_modes = [] # used to create hplus, hcross
paramList_all, paramList_quad_only = [], [] #used to transform sympy symbolic expression into actual functions of these parameters
paramList_derivatives = [] # Used to take derivatives over these params

# Define all parameters by looping over modes
for mode in linearModes:
    if mode == '220':
        A220, phi220, rew220, imw220 = symbols('A220, phi220, rew220, imw220', real=True)
        paramList_modes += [(2, 2, 0, False, A220, phi220, rew220, imw220)] # (l, m, n, mirror, ampl, phase, real(w), imag(w))
        paramList_all += [A220, phi220, rew220, imw220]
        paramList_derivatives += [A220, phi220, rew220, imw220]
    elif mode == '221':
        A221, phi221, rew221, imw221 = symbols('A221, phi221, rew221, imw221', real=True)
        paramList_modes += [(2, 2, 1, False, A221, phi221, rew221, imw221)]
        paramList_all += [A221, phi221, rew221, imw221]
        paramList_derivatives += [A221, phi221, rew221, imw221]
    elif mode == '210':
        A210, phi210, rew210, imw210 = symbols('A210, phi210, rew210, imw210', real=True)
        paramList_modes += [(2, 1, 0, False, A210, phi210, rew210, imw210)]
        paramList_all += [A210, phi210, rew210, imw210]
        paramList_derivatives += [A210, phi210, rew210, imw210]
    elif mode == '330':
        A330, phi330, rew330, imw330 = symbols('A330, phi330, rew330, imw330', real=True)
        paramList_modes += [(3, 3, 0, False, A330, phi330, rew330, imw330)]
        paramList_all += [A330, phi330, rew330, imw330]
        paramList_derivatives += [A330, phi330, rew330, imw330]
    elif mode == '440':
        A440, phi440, rew440, imw440 = symbols('A440, phi440, rew440, imw440', real=True)
        paramList_modes += [(4, 4, 0, False, A440, phi440, rew440, imw440)]
        paramList_all += [A440, phi440, rew440, imw440]
        paramList_derivatives += [A440, phi440, rew440, imw440]
    elif mode == '211':
        A211, phi211, rew211, imw211 = symbols('A211, phi211, rew211, imw211', real=True)
        paramList_modes += [(2, 1, 1, False, A211, phi211, rew211, imw211)]
        paramList_all += [A211, phi211, rew211, imw211]
        paramList_derivatives += [A211, phi211, rew211, imw211]
    elif mode == '331':
        A331, phi331, rew331, imw331 = symbols('A331, phi331, rew331, imw331', real=True)
        paramList_modes += [(3, 3, 1, False, A331, phi331, rew331, imw331)]
        paramList_all += [A331, phi331, rew331, imw331]
        paramList_derivatives += [A331, phi331, rew331, imw331]
    elif mode == '320':
        A320, phi320, rew320, imw320 = symbols('A320, phi320, rew320, imw320', real=True)
        paramList_modes += [(3, 2, 0, False, A320, phi320, rew320, imw320)]
        paramList_all += [A320, phi320, rew320, imw320]
        paramList_derivatives += [A320, phi320, rew320, imw320]
    elif mode == 'r220':
        Ar220, phir220, rewr220, imwr220 = symbols('Ar220, phir220, rewr220, imwr220', real=True)
        paramList_modes += [(2, 2, 0, True, Ar220, phir220, rewr220, imwr220)]
        paramList_all += [Ar220, phir220, rewr220, imwr220]
        paramList_derivatives += [Ar220, phir220, rewr220, imwr220]
    elif mode == 'r210':
        Ar210, phir210, rewr210, imwr210 = symbols('Ar210, phir210, rewr210, imwr210', real=True)
        paramList_modes += [(2, 1, 0, True, Ar210, phir210, rewr210, imwr210)]
        paramList_all += [Ar210, phir210, rewr210, imwr210]
        paramList_derivatives += [Ar210, phir210, rewr210, imwr210]
    else:
        raise Exception('Mode not implemented yet')

if (includeNL==1):
    ANL, phiNL, rewNL, imwNL = symbols('ANL, phiNL, rewNL, imwNL', real=True)
    paramList_modes += [(4, 4, 0, False, ANL, phiNL, rewNL, imwNL)]
    paramList_all += [ANL, phiNL, rewNL, imwNL]
    paramList_quad_only += [ANL, phiNL, rewNL, imwNL] # for computing only the quadratic mode SNR
    paramList_derivatives += [ANL, phiNL, rewNL, imwNL]
elif (includeNL==2):
    ANL, phiNL, rewNL, imwNL = symbols('ANL, phiNL, rewNL, imwNL', real=True)
    paramList_modes += [(4, 4, 0, False, ANL, phiNL, rewNL, imwNL)]
    paramList_all += [ANL, phiNL, rewNL, imwNL]
    paramList_quad_only += [ANL, phiNL, rewNL, imwNL] # for computing only the quadratic mode SNR

paramList_all += [M, theta, phi, psi, phiSH, iota, dl]
paramList_quad_only += [M, theta, phi, psi, phiSH, iota, dl]

nbparams_all = len(paramList_all)
fisher_dim = len(paramList_derivatives)
sig = nb.complex128(nb.float64[:], nb.float64) #Signature for numba functions




###################### DEFINING WAVEFORM TEMPLATE ################

# Spherical harmonics
def SWSH(s, l, m, phi, theta):
    # Normalization coefficient
    coeff = (-1)**(l + m - s) * sqrt((factorial(l + m) * factorial(l - m) * (2*l + 1)) / (4 * pi * factorial(l + s) * factorial(l - s)))   
    # trigonometric part
    trig_part = sin(theta / 2)**(2*l) * exp(I * m * phi)    
    # sum
    sum_part = Add(*[(-1)**r * binomial(l - s, r) * binomial(l + s, r + s - m) * 
                         cot(theta / 2)**(2 * r + s - m) for r in range(l - s + 1)])
    Y_lm = coeff * trig_part * sum_part
    return Y_lm

def Yplus(l,m,iota):
    Yp = SWSH(-2, l, m, 0, iota) + (-1)**l * SWSH(-2, l, -m, 0, iota)
    return Yp

def Ycross(l,m,iota):
    Yc = SWSH(-2, l, m, 0, iota) - (-1)**l * SWSH(-2, l, -m, 0, iota)
    return Yc



#Define hplus and hcross by summing over modes
hplus = 0
hcross = 0
for mode in paramList_modes:
    (l, m, n, mirror, Am, phim, rew, imw) = mode
    hplus += M*TSUN/(2*dl*Mpc/clight)*Am*Yplus(l, m, iota)*(
        exp(I*phim+I*m*phiSH)/(imw-I*(2*np.pi*f-rew)) 
       +exp(-I*phim-I*m*phiSH)/(imw-I*(2*np.pi*f+rew)))
    hcross += M*TSUN/(2*dl*Mpc/clight)*I*Am*Ycross(l, m, iota)*(
        exp(I*phim+I*m*phiSH)/(imw-I*(2*np.pi*f-rew)) 
       -exp(-I*phim-I*m*phiSH)/(imw-I*(2*np.pi*f+rew)))

# for nonlinear SNR only
hplus_quad = 0
hcross_quad = 0

ANL, phiNL, rewNL, imwNL = symbols('ANL, phiNL, rewNL, imwNL', real=True)
(l, m, n, mirror, Am, phim, rew, imw) = (4, 4, 0, False, ANL, phiNL, rewNL, imwNL)
hplus_quad += M*TSUN/(2*dl*Mpc/clight)*Am*Yplus(l, m, iota)*(
    exp(I*phim+I*m*phiSH)/(imw-I*(2*np.pi*f-rew)) 
    +exp(-I*phim-I*m*phiSH)/(imw-I*(2*np.pi*f+rew)))
hcross_quad += M*TSUN/(2*dl*Mpc/clight)*I*Am*Ycross(l, m, iota)*(
    exp(I*phim+I*m*phiSH)/(imw-I*(2*np.pi*f-rew)) 
    -exp(-I*phim-I*m*phiSH)/(imw-I*(2*np.pi*f+rew)))

# include all QQNMs if needed
# if (includeNL==3):
#     for mode1 in paramList_modes:
#         (l1, m1, n1, mirror1, Am1, phim1, rew1, imw1) = mode1
#         for mode2 in paramList_modes:
#             (l2, m2, n2, mirror2, Am2, phim2, rew2, imw2) = mode2
#             m = m1 + m2
#             for l in range(max(abs(l1-l2), m), l1+l2):
#                 try:
#                     ratio = QuadraticQNM_equatorialSym(l, l1, l2, m1, m2, n1, n2, mirror1, mirror2).NLRatio()
#                     #print(ratio)
#                     #assert (not math.isnan(ratio))
#                     Am = abs(ratio)*Am1*Am2
#                     phim = np.angle(ratio) + phim1 + phim2
#                     rew = rew1 + rew2
#                     imw = imw1 + imw2
#                     hplus += M*TSUN/(2*dl*Mpc/clight)*Am*Yplus(l, m, iota)*(
#                         exp(I*phim+I*m*phiSH)/(imw-I*(2*np.pi*f-rew)) 
#                        +exp(-I*phim-I*m*phiSH)/(imw-I*(2*np.pi*f+rew)))
#                     hcross += M*TSUN/(2*dl*Mpc/clight)*I*Am*Ycross(l, m, iota)*(
#                         exp(I*phim+I*m*phiSH)/(imw-I*(2*np.pi*f-rew)) 
#                        -exp(-I*phim-I*m*phiSH)/(imw-I*(2*np.pi*f+rew)))
#                 except Exception as e:
#                     print(e)


# Project hplus and hcross on detector antenna pattern functions
Fplus=np.sqrt(3)/2*(0.5*cos(2*psi)*(1+cos(theta)**2)*cos(2*phi) - sin(2*psi)*cos(theta)*sin(2*phi))
Fcross=np.sqrt(3)/2*(0.5*sin(2*psi)*(1+cos(theta)**2)*cos(2*phi) + cos(2*psi)*cos(theta)*sin(2*phi))
h=Fplus*hplus+Fcross*hcross
h_quad = Fplus*hplus_quad+Fcross*hcross_quad

hfun = njit(sig)(lambdify([paramList_all, f], h)) #Function generated from scipy, optimized by numba. frequency is a separate argument for the function because it will be useful later on to define an integrand.
hfun_quad_only = njit(sig)(lambdify([paramList_quad_only, f], h_quad))

# Define derivatives of wavefunction wrt parameters
# CAREFUL: if I want to introduce derivatives wrt spin or M, the QNM freqs and amplitudes depend on it but it's not obvious from this code
derivatives=[]
for param in paramList_derivatives:
    derivatives += [njit(sig)(lambdify([paramList_all, f], diff(h, param)))]



####################### DEFINING INTEGRANDS FOR FISHER AND SNR ################


### LISA is like 2 detectors with phi rotated by Pi/4, cf https://arxiv.org/abs/gr-qc/9703068

def integrandFisher(f, data): #data is an array containing all parameters
    i = int(data[0])
    j = int(data[1])
    args = data[2:]

    # Detector rotated by pi/4, the -5th argument corresponds to phi
    argsRotated = args.copy()
    argsRotated[-5] = argsRotated[-5] - np.pi/4 
    
    return 4/Sn(f,4)*np.real(derivatives[i](args, f)*np.conjugate(derivatives[j](args, f))
          +derivatives[i](argsRotated, f)*np.conjugate(derivatives[j](argsRotated, f)))


def integrandSNR(f, data): #data is an array containing all parameters

    # Detector rotated by pi/4, the -5th argument corresponds to phi
    argsRotated = data.copy()
    argsRotated[-5] = argsRotated[-5] - np.pi/4 
    
    return 4/Sn(f,4)*(np.abs(hfun(data, f))**2 + np.abs(hfun(argsRotated, f))**2)

def integrandSNR_quad(f, data): #data is an array containing all parameters

    # Detector rotated by pi/4, the -5th argument corresponds to phi
    argsRotated = data.copy()
    argsRotated[-5] = argsRotated[-5] - np.pi/4 
    
    return 4/Sn(f,4)*(np.abs(hfun_quad_only(data, f))**2 + np.abs(hfun_quad_only(argsRotated, f))**2)



# This function build a list of parameters from the modes amplitudes and phases to feed integrands
def get_params(A220, phi220, omega220, A221, phi221, omega221, A210, phi210, omega210, A211, phi211, omega211, A330, phi330, omega330, A331, phi331, omega331, A320, phi320, omega320, A440, phi440, omega440, Ar220, phir220, omegar220, Ar210, phir210, omegar210, ANL, phiNL, M, theta, phi, psi, phiSH, iota, dl):
    params = []
    for mode in linearModes:
        if mode == '220':
            params += [A220, phi220, np.real(omega220), np.imag(omega220)]
        elif mode == '221':
            params += [A221, phi221, np.real(omega221), np.imag(omega221)]
        elif mode == '210':
            params += [A210, phi210, np.real(omega210), np.imag(omega210)]
        elif mode == '330':
            params += [A330, phi330, np.real(omega330), np.imag(omega330)]
        elif mode == '440':
            params += [A440, phi440, np.real(omega440), np.imag(omega440)]
        elif mode == '211':
            params += [A211, phi211, np.real(omega211), np.imag(omega211)]
        elif mode == '331':
            params += [A331, phi331, np.real(omega331), np.imag(omega331)]
        elif mode == '320':
            params += [A320, phi320, np.real(omega320), np.imag(omega320)]
        elif mode == 'r220':
            params += [Ar220, phir220, np.real(omegar220), np.imag(omegar220)]
        elif mode == 'r210':
            params += [Ar210, phir210, np.real(omegar210), np.imag(omegar210)]
        else:
            raise Exception('Mode not implemented yet')
    if (includeNL==1 or includeNL==2):
        params += [ANL, phiNL, 2*np.real(omega220), 2*np.imag(omega220)]
    params += [M, theta, phi, psi, phiSH, iota, dl]
    return params


################## COMPUTING FISHER AND SNR ####################



# Compute Fisher and SNR from initial parameters of the BBH (masses, spins, angles, distance)
# alpha, beta, gamma refer to projections of the spins of the initial BHs, cf e.g. 1206.3803 and 1605.01938
# !!! CAREFUL here m1, m2 are the detector-frame masses (=source frame mass*(1+z))

def compute_fisher_SNR(m1, m2, a1, a2, alpha, beta, gamma, theta, phi, psi, phiSH, iota, dl):
    # Get remnant mass and spin from initial conditions
    M, afin = remnantMassSpin(m1, m2, a1, a2, alpha, beta, gamma)

    # Get QNM amplitudes and phases from initial conditions
    (A220, phi220, A221, phi221, A210, phi210, A211, phi211, A330, phi330, A331, phi331, A320, phi320, A440, phi440, Ar220, phir220, Ar210, phir210, ANL, phiNL) = QNMAmpl(m1, m2, a1, a2, alpha, beta, gamma)
    # Get QNM freqs
    omega220 = QNM_freq_highAccuracy(2, 2, 0, afin)/M/TSUN
    omega221 = QNM_freq_highAccuracy(2, 2, 1, afin)/M/TSUN
    omega210 = QNM_freq_highAccuracy(2, 1, 0, afin)/M/TSUN
    omega211 = QNM_freq_highAccuracy(2, 1, 1, afin)/M/TSUN
    omega330 = QNM_freq_highAccuracy(3, 3, 0, afin)/M/TSUN
    omega331 = QNM_freq_highAccuracy(3, 3, 1, afin)/M/TSUN
    omega320 = QNM_freq_highAccuracy(3, 2, 0, afin)/M/TSUN
    omega440 = QNM_freq_highAccuracy(4, 4, 0, afin)/M/TSUN
    omegar220 = -np.conjugate(QNM_freq_highAccuracy(2, -2, 0, afin))/M/TSUN
    omegar210 = -np.conjugate(QNM_freq_highAccuracy(2, -1, 0, afin))/M/TSUN

    # Put all the param in a list to feed to the Fisher functions
    params = get_params(A220, phi220, omega220, A221, phi221, omega221, A210, phi210, omega210, A211, phi211, omega211, A330, phi330, omega330, A331, phi331, omega331, A320, phi320, omega320, A440, phi440, omega440, Ar220, phir220, omegar220, Ar210, phir210, omegar210, ANL, phiNL, M, theta, phi, psi, phiSH, iota, dl)
    params_quad = [ANL, phiNL, 2*np.real(omega220), 2*np.imag(omega220), M, theta, phi, psi, phiSH, iota, dl] # for quadratic mode only
    params_fisher = np.array([0,0]+params)

    # Initialize fisher and cov as zero matrices
    fisher = np.zeros((fisher_dim, fisher_dim))
    cov = np.zeros((fisher_dim, fisher_dim))

    # First compute the SNR with an integral
    sol, err = integrate.quad(integrandSNR, fmin, fmax, args=(np.array(params),))
    SNR = np.sqrt(sol)

    # NEW: compute quadratic mode SNR 
    sol_quad, err_quad = integrate.quad(integrandSNR_quad, fmin, fmax, args=(np.array(params_quad),)) # added definitions of integrandSNR_quad and params_quad above
    SNR_quad = np.sqrt(sol_quad)

    #If SNR is too small or integration unsuccessful it's not even worth computing Fisher
    if (SNR_quad < 8.) or (abs(err_quad/(sol+epsiloncond))>0.1):
        #print(SNR_quad)
        fisher = np.nan * np.ones((fisher_dim, fisher_dim))
        cov = np.nan * np.ones((fisher_dim, fisher_dim))
        return (SNR, SNR_quad, fisher, cov)

    # Now we loop over all params to compute Fisher
    totSuccess=True # Useful to break the loop in case integration fails
    for i in range(fisher_dim):
        for j in range(i, fisher_dim):
            params_fisher[0] = i
            params_fisher[1] = j
            sol, err = integrate.quad(integrandFisher, fmin, fmax, args=(params_fisher,))
            if (abs(err/(sol+epsiloncond))>0.1):
                print('Integration unsuccessful')
                totSuccess = False
                break
            fisher[i,j] = sol
            fisher[j,i] = sol

    #Compute cov matrix; eliminate cases where there has been numerical integration errors
    if not totSuccess:
        fisher = np.nan * np.ones((fisher_dim, fisher_dim))
        cov = np.nan * np.ones((fisher_dim, fisher_dim))
    else:
        fishercond = fisher + epsiloncond*np.identity(fisher_dim)
        cov = np.linalg.inv(fishercond)

    return (SNR, fisher, cov)