#!/usr/bin/env python
# Copyright (C) 2010  Nickolas Fotopoulos, Stephen Fairhurst
#
# 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.
"""
Given an XML file with a SimInspiralTable, align each binary so that the total
angular momentum has the recorded inclination instead of the orbital
angular momentum. This will alter inclination, spins, and effective distance.

Usage: %prog --output-file OUTFILE siminsp_fname
"""
from __future__ import division

__author__ = "Nickolas Fotopoulos <nickolas.fotopoulos@ligo.org>"

import math
import optparse
import sys

import numpy as np

from glue.ligolw import ligolw
from glue.ligolw import lsctables
from glue.ligolw import table
from glue.ligolw import utils
from glue.ligolw.utils import process as ligolw_process
from pylal import git_version
import lal
from lalsimulation import SimInspiralTransformPrecessingNewInitialConditions

#
# Utility functions
#

def eff_distance(detector, sim):
    """
    Return the effective distance.

    Ref: Duncan's PhD thesis, eq. (4.3) on page 57, implemented in
         LALInspiralSiteTimeAndDist in SimInspiralUtils.c:594.
    """
    f_plus, f_cross = lal.ComputeDetAMResponse(detector.response,
        sim.longitude, sim.latitude, sim.polarization, sim.end_time_gmst)
    ci = math.cos(sim.inclination)
    s_plus = -(1 + ci * ci)
    s_cross = -2 * ci
    return 2 * sim.distance / math.sqrt(f_plus * f_plus * s_plus * s_plus + \
        f_cross * f_cross * s_cross * s_cross)

def get_spins(sim):
    """
    Return the 3-vectors s1, s2 from a sim_inspiral row.
    """
    return np.array([sim.spin1x, sim.spin1y, sim.spin1z]), \
           np.array([sim.spin2x, sim.spin2y, sim.spin2z])

def set_spins(sim, s1, s2):
    """
    Set the components of the spin in a sim_inspiral row from 3-vectors.
    """
    sim.spin1x = s1[0]
    sim.spin1y = s1[1]
    sim.spin1z = s1[2]
    sim.spin2x = s2[0]
    sim.spin2y = s2[1]
    sim.spin2z = s2[2]

def alignTotalSpin(iota, s1, s2, m1, m2, f_lower):
    """
    Rigidly rotate binary so that the total angular momentum has the given
    inclination (iota) instead of the orbital angular momentum. Return
    the new inclination, s1, and s2. s1 and s2 are dimensionless spin.
    Note: the spins are assumed to be given in the frame defined by the orbital
    angular momentum.  
    """

    # Calculate the quantities required by SimInspiralTransformPrecessingNewInitialConditions
    if np.linalg.norm(s1[:2]) != 0 or np.linalg.norm(s1[:2]) != 0:
        thetaJN = iota
        phiJL = np.pi
        l = np.array([0,0,1.])
        theta1 = np.arccos(np.dot(s1, l)/np.linalg.norm(s1))
        theta2 = np.arccos(np.dot(s2, l)/np.linalg.norm(s2))
        phi12 = np.arccos(np.dot(s1[:2], s2[:2]) / np.linalg.norm(s1[:2]) / \
                          np.linalg.norm(s2[:2]))

        news1 = np.zeros(3)
        news2 = np.zeros(3)
        newIota, news1[0], news1[1], news1[2], news2[0], news2[1], news2[2] = \
            SimInspiralTransformPrecessingNewInitialConditions(thetaJN, phiJL,
                    theta1, theta2, phi12, np.linalg.norm(s1),
                    np.linalg.norm(s2), m1 * lal.MSUN_SI, m2 * lal.MSUN_SI,
                    f_lower)
        return newIota, news1, news2
    else:
        return iota, s1, s2

#
# Parse commandline
#

parser = optparse.OptionParser(usage=__doc__, version=git_version.verbose_msg)
parser.add_option("--output-file", metavar="OUTFILE",
    help="name of output XML file")
opts, args = parser.parse_args()

# everything is required
if (opts.output_file is None) or \
   (len(args) != 1):
    parser.print_usage()
    sys.exit(2)

siminsp_fname = args[0]

#
# Read inputs
#

siminsp_doc = utils.load_filename(siminsp_fname,
    gz=siminsp_fname.endswith(".gz"), contenthandler = lsctables.use_in(ligolw.LIGOLWContentHandler))

# Prepare process table with information about the current program
process = ligolw_process.register_to_xmldoc(siminsp_doc,
    "ligolw_cbc_align_total_spin",
    opts.__dict__, version=git_version.tag or git_version.id,
    cvs_repository="lalsuite", cvs_entry_time=git_version.date)


#
# Do the work
#
site_location_list = [\
    ("h", lal.CachedDetectors[lal.LALDetectorIndexLHODIFF]),
    ("l", lal.CachedDetectors[lal.LALDetectorIndexLLODIFF]),
    ("g", lal.CachedDetectors[lal.LALDetectorIndexGEO600DIFF]),
    ("l", lal.CachedDetectors[lal.LALDetectorIndexTAMA300DIFF]),
    ("l", lal.CachedDetectors[lal.LALDetectorIndexVIRGODIFF])]
for sim in table.get_table(siminsp_doc, lsctables.SimInspiralTable.tableName):
    # rotate
    s1, s2 = get_spins(sim)
    newIota, news1, news2 = alignTotalSpin(sim.inclination, s1, s2,
        sim.mass1, sim.mass2, sim.f_lower)
    sim.inclination = newIota
    set_spins(sim, news1, news2)

    # update effective distances at the sites
    for site, detector in site_location_list:
        setattr(sim, "eff_dist_" + site, eff_distance(detector, sim))


#
# Write output
#

ligolw_process.set_process_end_time(process)
utils.write_filename(siminsp_doc, opts.output_file)
