# Copyright (C) 2020  Ssohrab Borhanian
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.


from numpy import pi, euler_gamma

import gwbench.basic_relations as brs
from gwbench.utils import MTsun, strain_fac

wf_symbs_string = 'f Mc eta chi1z chi2z DL tc phic iota lam_t delta_lam_t'

def hfpc_raw(f, Mc, eta, chi1z, chi2z, DL, tc, phic, iota, lam_t, delta_lam_t, cos, sin, exp, log, is_lam12=0):
     '''
     Implementation of TaylorF2 waveform model for BNS systems that depend on the following parameters:
     - f: frequency array
     - Mc: chirp mass
     - eta: symmetric mass ratio
     - chi1x, chi1y, chi1z: dimensionless spin components of the primary BH
     - chi2x, chi2y, chi2z: dimensionless spin components of the secondary BH
     - DL: luminosity distance
     - tc: coalescence time
     - phic: coalescence phase
     - iota: inclination angle
     - lam_t: effective tidal deformability
     - delta_lam_t: delta effictive tidal deformability

     If is_lam12 is set to True, the function assumes that lam_t and delta_lam_t are the individual tidal deformabilities
     of the two neutron stars.

     Parameters
     ----------
     f: np.ndarray or jnp.ndarray or sympy.Symbol
          Frequency array
     Mc: float or sympy.Symbol
          Chirp mass [solar mass]
     eta: float or sympy.Symbol
          Symmetric mass ratio
     chi1x: float or sympy.Symbol
          Dimensionless spin component of the primary BH along the x-axis
     chi1y: float or sympy.Symbol
          Dimensionless spin component of the primary BH along the y-axis
     chi1z: float or sympy.Symbol
          Dimensionless spin component of the primary BH along the z-axis
     chi2x : float or sympy.Symbol
          Dimensionless spin component of the secondary BH along the x-axis
     chi2y: float or sympy.Symbol
          Dimensionless spin component of the secondary BH along the y-axis
     chi2z: float or sympy.Symbol
          Dimensionless spin component of the secondary BH along the z-axis
     DL: float or sympy.Symbol
          Luminosity distance [Mpc]
     tc: float or sympy.Symbol
          Coalescence time [s]
     phic: float or sympy.Symbol
          Coalescence phase [rad]
     iota: float or sympy.Symbol
          Inclination angle [rad]
     lam_t: float or sympy.Symbol
          Effective tidal deformability
     delta_lam_t: float or sympy.Symbol
          Delta effective tidal deformability
     cos: np.cos or jnp.cos or sympy.cos
          cosine function
     sin: np.sin or jnp.sin or sympy.sin
          sine function
     exp: np.exp or jnp.exp or sympy.exp
          exponential function
     log: np.log or jnp.log or sympy.log
          natural logarithm function

     Returns
     -------
     hfp: np.ndarray or jnp.ndarray or sympy expression
          Plus polarization waveform
     hfc: np.ndarray or jnp.ndarray or sympy expression
          Cross polarization waveform
     '''

     # Mc ... in solar mass
     # DL ... in mega parsec
     # convert to sec
     Mc = Mc * MTsun
     DL = DL * MTsun/strain_fac

     # get sym and asym chi combinations
     chi_s = brs.chi_s(chi1z,chi2z)
     chi_a = brs.chi_a(chi1z,chi2z)

     if is_lam12:
          # in this case: lam_t = lam1, delta_lam_t = lam2
          lam1 = lam_t
          lam2 = delta_lam_t
          delta = sqrt(1 - 4 * eta)
          lam_t = 8./13. * ( (1. + 7. * eta - 31. * eta**2) * (lam1 + lam2) +
                         delta * (1. + 9. * eta - 11. * eta**2) * (lam1 - lam2) )
          delta_lam_t = 0.5 * ( delta * (1319. - 13272. * eta + 8944. * eta**2) / 1319. * (lam1 + lam2)+
                         (1319. - 15910. * eta + 32850. * eta**2 + 3380. * eta**3) / 1319. * (lam1 - lam2) )

     # Mc is in sec, e.g., Mc = 10*MTSUN_SI (for 10 solar mass)
     # DL is in sec, e.g., DL = 100*1e6*PC_SI/C_SI (for 100 Mpc)
     M = Mc/eta**(3./5.)
     delta = (1.-4.*eta)**0.5
     v  = (pi*M*f)**(1./3.)
     flso = brs.f_isco(M)
     vlso = (pi*M*flso)**(1./3.)
     A =((5./24.)**0.5/pi**(2./3.))*(Mc**(5./6.)/DL)

     # 3.5PN phasing (point particle limit)
     p0 = 1.

     p1 = 0

     p2 = (3715./756. + (55.*eta)/9.)

     p3 = (-16.*pi + (113.*delta*chi_a)/3. + (113./3. - (76.*eta)/3.)*chi_s)

     p4 = (15293365./508032. + (27145.*eta)/504.+ (3085.*eta**2)/72. + (-405./8. + 200.*eta)*chi_a**2 - (405.*delta*chi_a*chi_s)/4. + (-405./8. + (5.*eta)/2.)*chi_s**2)

     gamma = (732985./2268. - 24260.*eta/81. - 340.*eta**2/9.)*chi_s + (732985./2268. + 140.*eta/9.)*delta*chi_a

     p5 = (38645.*pi/756. - 65.*pi*eta/9. - gamma)

     p5L = (38645.*pi/756. - 65.*pi*eta/9. - gamma)*3*log(v/vlso)

     p6 = (11583231236531./4694215680. - 640./3.*pi**2 - 6848./21.*euler_gamma + eta*(-15737765635./3048192. + 2255./12.*pi**2) + eta*eta*76055./1728. - eta*eta*eta*127825./1296. \
          - (6848./21.)*log(4.) + pi*(2270.*delta*chi_a/3. + (2270./3. - 520.*eta)*chi_s) + (75515./144. - 8225.*eta/18.)*delta*chi_a*chi_s \
          + (75515./288. - 263245.*eta/252. - 480.*eta**2)*chi_a**2 + (75515./288. - 232415.*eta/504. + 1255.*eta**2/9.)*chi_s**2)

     p6L = -(6848./21.)*log(v)

     p7 = (((77096675.*pi)/254016. + (378515.*pi*eta)/1512.- (74045.*pi*eta**2)/756. + (-25150083775./3048192. + (10566655595.*eta)/762048. - (1042165.*eta**2)/3024. + (5345.*eta**3)/36.
          + (14585./8. - 7270.*eta + 80.*eta**2)*chi_a**2)*chi_s + (14585./24. - (475.*eta)/6. + (100.*eta**2)/3.)*chi_s**3 + delta*((-25150083775./3048192.
          + (26804935.*eta)/6048. - (1985.*eta**2)/48.)*chi_a + (14585./24. - 2380.*eta)*chi_a**3 + (14585./8. - (215.*eta)/2.)*chi_a*chi_s**2)))

     phase = 2*f*pi*tc - phic - pi/4. + (3./(128.*v**5*eta))*(p0 + v**2*p2 + v**3*p3+ v**4*p4 + v**5*(p5+p5L) + v**6*(p6+p6L) + v**7*p7)

     # From Lesli Wade's paper ... adding the tidal deformability phase
     pT = (3./128./eta/v**5.)*(-39.*lam_t*v**10/2. + (-3115.*lam_t/64. + 6595.*(1-4*eta)**0.5*delta_lam_t/364.)*v**12)
     phase += pT

     hpc   = A * f**(-7./6.) * exp(1j * phase)

     return 0.5 * (1 + cos(iota)**2) * hpc, 1j * cos(iota) * hpc
