#!/usr/bin/python3.13 -s

"""For a given external trigger (GRB, FRB, neutrino, etc...), generate a sky grid covering its localization error region.

This sky grid will be used by `pycbc_multi_inspiral` to find multi-detector gravitational wave triggers and calculate the
coherent SNRs and related statistics. 

The grid is constructed following the method described in Section V of https://arxiv.org/abs/1410.6042"""

import numpy as np
import argparse
import itertools
from scipy.spatial.transform import Rotation as R

import pycbc
from pycbc.detector import Detector
from pycbc.io.hdf import HFile


def spher_to_cart(sky_points):
    """Convert spherical coordinates to cartesian coordinates.
    """
    cart = np.zeros((len(sky_points), 3))
    cart[:,0] = np.cos(sky_points[:,0]) * np.cos(sky_points[:,1])
    cart[:,1] = np.sin(sky_points[:,0]) * np.cos(sky_points[:,1])
    cart[:,2] = np.sin(sky_points[:,1])
    return cart

def cart_to_spher(sky_points):
    """Convert cartesian coordinates to spherical coordinates.
    """
    spher = np.zeros((len(sky_points), 2))
    spher[:,0] = np.arctan2(sky_points[:,1], sky_points[:,0])
    spher[:,1] = np.arcsin(sky_points[:,2])
    return spher

parser = argparse.ArgumentParser(description=__doc__)
pycbc.add_common_pycbc_options(parser)
parser.add_argument('--ra', type=float,
    help="Right ascension (in rad) of the center of the external trigger "
         "error box")
parser.add_argument('--dec', type=float,
    help="Declination (in rad) of the center of the external trigger "
         "error box")
parser.add_argument('--instruments', nargs="+", type=str, required=True,
    help="List of instruments to analyze.")
parser.add_argument('--sky-error', type=float, required=True,
    help="3-sigma confidence radius (in rad) of the external trigger error "
         "box")
parser.add_argument('--trigger-time', type=int, required=True,
    help="Time (in s) of the external trigger")
parser.add_argument('--timing-uncertainty', type=float, default=0.0001,
    help="Timing uncertainty (in s) we are willing to accept")
parser.add_argument('--output', type=str, required=True,
        help="Name of the sky grid")

args = parser.parse_args()

pycbc.init_logging(args.verbose)

if len(args.instruments) == 1:
    parser.error('Can not make a sky grid for only one detector.')

args.instruments.sort() # Put the ifos in alphabetical order
detectors = args.instruments
detectors = [Detector(d) for d in detectors]
detector_pairs = list(itertools.combinations(detectors, 2))

# Calculate the time delay for each detector pair
tds = [detector_pairs[i][0].time_delay_from_detector(detector_pairs[i][1], args.ra, args.dec, args.trigger_time) for i in range(len(detector_pairs))]

# Calculate the light travel time between the detector pairs
light_travel_times = [detector_pairs[i][0].light_travel_time_to_detector(detector_pairs[i][1]) for i in range(len(detector_pairs))]

# Calculate the required angular spacing between the sky points
ang_spacings = [(2*args.timing_uncertainty) / np.sqrt(light_travel_times[i]**2 - tds[i]**2) for i in range(len(detector_pairs))]
angular_spacing = min(ang_spacings)

sky_points = np.zeros((1, 2))

number_of_rings = int(args.sky_error / angular_spacing)

# Generate the sky grid centered at the North pole
for i in range(number_of_rings+1):
    if i == 0:
        sky_points[0][0] = 0
        sky_points[0][1] = np.pi/2
    else:
        number_of_points = int(2*np.pi*i)
        for j in range(number_of_points):
            sky_points = np.row_stack((sky_points, np.array([j/i, np.pi/2 - i*angular_spacing])))

# Convert spherical coordinates to cartesian coordinates
cart = spher_to_cart(sky_points)

grb = np.zeros((1, 2))
grb[0] = args.ra, args.dec
grb_cart = spher_to_cart(grb)

north_pole = [0, 0, 1]

ort = np.cross(grb_cart, north_pole)
norm = np.linalg.norm(ort)
ort /= norm
n = -np.arccos(np.dot(grb_cart, north_pole))
u = ort*n

# Rotate the sky grid to the center of the external trigger error box
r = R.from_rotvec(u)
rota = r.apply(cart)

# Convert cartesian coordinates back to spherical coordinates
spher = cart_to_spher(rota)

# Calculate the time delays between the Earth center and each detector for each sky point
time_delays = [[detectors[i].time_delay_from_earth_center(
               spher[j][0], spher[j][1], args.trigger_time) for j in range(len(spher))] for i in range(len(detectors))]

with HFile(args.output, 'w') as hf:
    hf['ra'] = spher[:,0]
    hf['dec'] = spher[:,1]
    hf['trigger_ra'] = [args.ra]
    hf['trigger_dec'] = [args.dec]
    hf['sky_error'] = [args.sky_error]
    hf['trigger_time'] = [args.trigger_time]
    hf['timing_uncertainty'] = [args.timing_uncertainty]
    hf['instruments'] = [d for d in args.instruments]
    hf['time_delays'] = time_delays
