#!/usr/bin/env python

"""Followup utility to optimize the SNR of a PyCBC Live trigger."""

import os
import shutil
import argparse, numpy, h5py
import time
import types
import logging
import tempfile
from scipy.optimize import differential_evolution, shgo
import pycbc
from pycbc import (
    DYN_RANGE_FAC, fft, scheme, version, waveform
)
from pycbc.types import MultiDetOptionAction, zeros, load_frequencyseries
from pycbc.filter import matched_filter_core
import pycbc.waveform.bank
from pycbc.conversions import (
    mchirp_from_mass1_mass2, mtotal_from_mchirp_eta,
    mass1_from_mtotal_eta, mass2_from_mtotal_eta
)
from pycbc.io import live
from pycbc.io.hdf import load_hdf5_to_dict
from pycbc.detector import Detector
from pycbc.psd import interpolate

try:
    import pyswarms as ps
except:
    ps = None


Nfeval = 0
start_time = time.time()

def callback_func(Xi, convergence=0):
    global Nfeval
    logging.info("Currently at %d %s", Nfeval, convergence)
    if (time.time() - start_time) > 360:
        return True
    Nfeval += 1


def compute_network_snr_core(v, *argv, raise_err=False):
    """
    Compute network SNR as a function over mchirp, eta and two aligned
    spin components, stored in that order in the sequence v
    """    
    data = argv[0]
    coinc_times = argv[1]
    ifos = argv[2]
    flen = argv[3]
    approximant = argv[4]
    flow = argv[5]
    f_end = argv[6]
    delta_f = argv[7]
    sample_rate = argv[8]
    distance = 1.0 / DYN_RANGE_FAC
    mtotal = mtotal_from_mchirp_eta(v[0], v[1])
    mass1 = mass1_from_mtotal_eta(mtotal, v[1])
    mass2 = mass2_from_mtotal_eta(mtotal, v[1])

    # enforce broadly accepted search space boundaries
    if mass1 < 1 or mass2 < 1 or mtotal > 500:
        return -numpy.inf, {}

    try:
        htilde = waveform.get_waveform_filter(
                zeros(flen, dtype=numpy.complex64),
                approximant=approximant,
                mass1=mass1, mass2=mass2, spin1z=v[2], spin2z=v[3],
                f_lower=flow, f_final=f_end, delta_f=delta_f,
                delta_t=1.0/sample_rate, distance=distance)
    except RuntimeError:
        if raise_err:
            raise
        # assume a failure in the waveform approximant
        # due to the choice of parameters
        else:
            return -numpy.inf, {}

    if not hasattr(htilde, 'params'):
        htilde.params = dict(mass1=mass1, mass2=mass2,
                             spin1z=v[2], spin2z=v[3])
    if not hasattr(htilde, 'end_idx'):
        htilde.end_idx = int(f_end / htilde.delta_f)
    htilde.approximant = approximant
    htilde.sigmasq = types.MethodType(pycbc.waveform.bank.sigma_cached,
                                      htilde)
    htilde.min_f_lower = flow
    htilde.end_frequency = f_end
    htilde.f_lower = flow
    network_snrsq = 0
    snr_series_dict = {}
    for ifo in ifos:
        sigmasq = htilde.sigmasq(data[ifo].psd)
        snr, _, norm = matched_filter_core(htilde, data[ifo],
                                           h_norm=sigmasq)
        duration = 0.095
        half_dur_samples = int(snr.sample_rate * duration / 2)
        onsource_idx = float(coinc_times[ifo] - snr.start_time) * snr.sample_rate
        onsource_idx = int(round(onsource_idx))
        onsource_slice = slice(onsource_idx - half_dur_samples,
                               onsource_idx + half_dur_samples + 1)
        snr_series = snr[onsource_slice] * norm
        snr_series_dict[ifo] = snr * norm
        snr_series_dict['sigmasq_' + ifo] = sigmasq
        network_snrsq += max(abs(snr_series._data)) ** 2.

    return network_snrsq ** 0.5, snr_series_dict


def compute_network_snr(v, *argv):
    if len(argv) == 1:
        argv = argv[0]
    nsnr, _ = compute_network_snr_core(v, *argv)
    return -nsnr


def compute_network_snr_pso(v, *argv, **kwargs):
    argv = kwargs['args']
    nsnr_array = numpy.array([-compute_network_snr_core(v_i, *argv)[0] for v_i in v])
    return nsnr_array


def optimize_di(bounds, cli_args, extra_args):
    bounds = [
        bounds['mchirp'],
        bounds['eta'],
        bounds['spin1z'],
        bounds['spin2z']
    ]
    results = differential_evolution(
        compute_network_snr,
        bounds,
        maxiter=cli_args.di_maxiter,
        workers=(cli_args.cores or -1),
        popsize=cli_args.di_popsize,
        mutation=(0.5, 1),
        recombination=0.7,
        callback=callback_func,
        args=extra_args
    )
    return results.x


def optimize_shgo(bounds, cli_args, extra_args):
    bounds = [
        bounds['mchirp'],
        bounds['eta'],
        bounds['spin1z'],
        bounds['spin2z']
    ]
    results = shgo(
        compute_network_snr,
        bounds=bounds,
        args=extra_args,
        iters=args.shgo_iters,
        n=args.shgo_samples,
        sampling_method="sobol"
    )
    return results.x


def optimize_pso(bounds, cli_args, extra_args):
    options = {
        'c1': cli_args.pso_c1,
        'c2': cli_args.pso_c2,
        'w': cli_args.pso_w
    }
    min_bounds = [
        bounds['mchirp'][0],
        bounds['eta'][0],
        bounds['spin1z'][0],
        bounds['spin2z'][0]
    ]
    max_bounds = [
        bounds['mchirp'][1],
        bounds['eta'][1],
        bounds['spin1z'][1],
        bounds['spin2z'][1]
    ]
    optimizer = ps.single.GlobalBestPSO(
        n_particles=cli_args.pso_particles,
        dimensions=4,
        options=options,
        bounds=(min_bounds, max_bounds)
    )
    _, results = optimizer.optimize(
        compute_network_snr_pso,
        iters=cli_args.pso_iters,
        n_processes=cli_args.cores,
        args=extra_args
    )
    return results


optimize_funcs = {
    'differential_evolution': optimize_di,
    'shgo': optimize_shgo,
    'pso': optimize_pso
}


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--version', action='version',
                    version=version.git_verbose_msg)
parser.add_argument('--verbose', action='store_true')
parser.add_argument('--params-file', required=True,
                    help='Location of the attributes file created by PyCBC '
                         'Live')
parser.add_argument('--data-files', type=str, nargs='+',
                    action=MultiDetOptionAction,
                    metavar='IFO:DATA_FILE',
                    help='Locations of the overwhitened data files produced '
                         'by PyCBC Live.')
parser.add_argument('--psd-files', type=str, nargs='+',
                    action=MultiDetOptionAction,
                    metavar='IFO:PSD_FILE',
                    help='Locations of the PSD files produced '
                         'by PyCBC Live.')
parser.add_argument('--approximant', required=True,
                    help='Waveform approximant string.')
parser.add_argument('--snr-threshold', default=4.0,
                    help='If the SNR in ifo X is below this threshold do not '
                         'consider it part of the coincidence. Not implemented')
parser.add_argument('--gracedb-server', metavar='URL',
                    help='URL of GraceDB server API for uploading events. '
                         'If not provided, the default URL is used.')
parser.add_argument('--gracedb-search', type=str, default='AllSky',
                    help='String going into the "search" field of the GraceDB '
                         'events')
parser.add_argument('--production', action='store_true',
                    help='Upload a production event rather than a test event')
parser.add_argument('--enable-gracedb-upload', action='store_true', default=False,
                    help='Upload triggers to GraceDB')
parser.add_argument('--output-path', required=True,
                    help='Path to a directory to store results in')
parser.add_argument('--cores', type=int,
                    help='Restrict calculation to given number of CPU cores')
parser.add_argument('--optimizer', type=str, default='differential_evolution',
                    choices=sorted(optimize_funcs),
                    help='The optimizer to use, differential_evolution, shgo or pso')
parser.add_argument('--di-maxiter', type=int, default=50,
                    help='Only relevant for --optimizer differential_evolution: '
                    'The maximum number of generations over which the entire '
                    'population is evolved.')
parser.add_argument('--di-popsize', type=int, default=100,
                    help='Only relevant for --optimizer differential_evolution: '
                         'A multiplier for setting the total population size.')
parser.add_argument('--shgo-samples', type=int, default=76,
                    help='Only relevant for --optimizer shgo: '
                    'Number of sampling points used in the construction of the '
                    'simplicial complex.')
parser.add_argument('--shgo-iters', type=int, default=3,
                    help='Only relevant for --optimizer shgo: '
                    'Number of iterations used in the construction of the simplicial complex.')
parser.add_argument('--pso-iters', type=int, default=5,
                    help='Only relevant for --optimizer pso: '
                    'Number of iterations used in the particle swarm optimization.')
parser.add_argument('--pso-particles', type=int, default=250,
                    help='Only relevant for --optimizer pso: '
                    'Number of particles used in the swarm.')
parser.add_argument('--pso-c1', type=float, default=0.5,
                    help='Only relevant for --optimizer pso: '
                    'The hyperparameter c1: the cognitive parameter.')
parser.add_argument('--pso-c2', type=float, default=2.0,
                    help='Only relevant for --optimizer pso: '
                    'The hyperparameter c2: the social parameter.')
parser.add_argument('--pso-w', type=float, default=0.01,
                    help='Only relevant for --optimizer pso: '
                    'The hyperparameter w: the intertia parameter.')

scheme.insert_processing_option_group(parser)
fft.insert_fft_option_group(parser)

# Input checking
args = parser.parse_args()
if args.snr_threshold != 4:
    parser.error("Sorry, the SNR threshold option doesn't work yet")
if args.optimizer == 'pso' and ps == None:
    parser.error('You need to install pyswarms to use the pso optimizer.')
pycbc.init_logging(args.verbose)

scheme.verify_processing_options(args, parser)
fft.verify_fft_options(args, parser)

scheme_context = scheme.from_cli(args)
fft.from_cli(args)

logging.info('Starting optimize SNR')

data = {}
ifos = list(args.data_files.keys())
for ifo in ifos:
    data[ifo] = load_frequencyseries(args.data_files[ifo])
    data[ifo].psd = load_frequencyseries(args.psd_files[ifo])

fp = h5py.File(args.params_file, 'r')
original_gid = None
if 'gid' in fp:
    original_gid = fp['gid'].asstr()[()]
if args.enable_gracedb_upload and not original_gid:
    raise RuntimeError('Params must include original gracedb ID in order '
                       'to upload followup!')
if original_gid:
    logging.info('Following up GID %s', original_gid)

coinc_times = {}
for ifo in ifos:
    try:
        coinc_times[ifo] = fp['coinc_times'][ifo][()]
    except KeyError:
        pass
coinc_ifos = list(coinc_times.keys())

flen = fp['flen'][()]
approximant = args.approximant
flow = float(fp['flow'][()])
f_end = float(fp['f_end'][()])
delta_f = fp['delta_f'][()]
sample_rate = fp['sample_rate'][()]

extra_args = [data, coinc_times, coinc_ifos, flen,
              approximant, flow, f_end, delta_f, sample_rate]

mchirp = mchirp_from_mass1_mass2(
    fp['mass1'][()],
    fp['mass2'][()]
)
minchirp = mchirp * (1 - mchirp / 50.0)
maxchirp = mchirp * (1 + mchirp / 50.0)
minchirp = 1 if minchirp < 1 else minchirp
maxchirp = 80 if maxchirp > 80 else maxchirp

# boundary of the optimization space (dict of (min, max) tuples)
bounds = {
    'mchirp': (minchirp, maxchirp),
    'eta': (0.01, 0.2499),
    'spin1z': (-0.9, 0.9),
    'spin2z': (-0.9, 0.9)
}

with scheme_context:
    logging.info('Starting optimization')

    optimize_func = optimize_funcs[args.optimizer]
    opt_params = optimize_func(bounds, args, extra_args)

    logging.info('Optimization complete')

    fup_ifos = set(ifos) - set(coinc_ifos)
    for ifo in fup_ifos:
        coinc_times[ifo] = coinc_times[coinc_ifos[0]]

    extra_args[2] = ifos

    _, snr_series_dict = compute_network_snr_core(opt_params, *extra_args)

mtotal = mtotal_from_mchirp_eta(opt_params[0], opt_params[1])
mass1 = mass1_from_mtotal_eta(mtotal, opt_params[1])
mass2 = mass2_from_mtotal_eta(mtotal, opt_params[1])
spin1z = opt_params[2]
spin2z = opt_params[3]

# Prepare for GraceDB upload
coinc_results = {}
followup_data = {}

loudest_ifo = None
loudest_snr_allifos = 0
loudest_snr_time_allifos = None
loudest_snrs = {}
loudest_snr_times = {}
loudest_snr_idxs = {}

# Determine which ifo has the loudest SNR peak ... We'll use this to determine
# the time window for other ifos (e.g. if one ifo has a loud SNR peak, and the
# other has SNR < 4, we would need to determine the loudest SNR for the quieter
# instrument within light-travel time of the loud SNR peak)
for ifo in ifos:
    duration = 0.095
    half_dur_samples = int(sample_rate * duration / 2)
    onsource_idx = float(coinc_times[ifo] - snr_series_dict[ifo].start_time) \
        * snr_series_dict[ifo].sample_rate
    onsource_idx = int(round(onsource_idx))
    onsource_slice = slice(onsource_idx - half_dur_samples,
                           onsource_idx + half_dur_samples + 1)
    max_snr_idx = numpy.argmax(abs(snr_series_dict[ifo][onsource_slice]))
    max_snr_idx = max_snr_idx + onsource_idx - half_dur_samples
    max_snr = snr_series_dict[ifo][max_snr_idx]
    max_snr_time = snr_series_dict[ifo].start_time \
        + max_snr_idx*(1./sample_rate)
    loudest_snr_idxs[ifo] = max_snr_idx
    loudest_snr_times[ifo] = float(max_snr_time)
    loudest_snrs[ifo] = max_snr
    if abs(max_snr) > loudest_snr_allifos:
        loudest_snr_allifos = abs(max_snr)
        loudest_ifo = ifo
        loudest_snr_time_allifos = loudest_snr_times[ifo]

# Now create the snr_series_dict
for ifo in ifos:
    # Find loudest SNR within coincidence window of the largest peak
    snr_peak = abs(loudest_snrs[ifo])
    # And check light travel time
    lttbd = Detector(ifo).light_travel_time_to_detector(Detector(loudest_ifo))
    # Add small buffer
    lttbd += 0.005

    time_window = [loudest_snr_time_allifos-lttbd,
                   loudest_snr_time_allifos+lttbd]

    snr_slice = snr_series_dict[ifo].time_slice(time_window[0], time_window[1])
    max_snr_idx = numpy.argmax(abs(snr_slice))
    idx_offset = snr_slice.start_time - snr_series_dict[ifo].start_time
    idx_offset = int(float(idx_offset) * sample_rate + 0.5)
    max_snr_idx = max_snr_idx + idx_offset
    new_snr_slice = slice(max_snr_idx - int(sample_rate/10),
                          max_snr_idx + int(sample_rate/10)+1)
    snr_series_dict[ifo] = snr_series_dict[ifo][new_snr_slice]

netsnr = 0
for idx, ifo in enumerate(ifos):
    coinc_results['foreground/'+ifo+'/mass1'] = mass1
    coinc_results['foreground/'+ifo+'/mass2'] = mass2
    coinc_results['foreground/'+ifo+'/spin1z'] = spin1z
    coinc_results['foreground/'+ifo+'/spin2z'] = spin2z
    coinc_results['foreground/'+ifo+'/f_lower'] = flow
    coinc_results['foreground/'+ifo+'/window'] = 0.1
    coinc_results['foreground/'+ifo+'/sample_rate'] = int(sample_rate)
    coinc_results['foreground/'+ifo+'/template_id'] = 0
    # Apparently GraceDB gets upset if chi-squared is not set
    coinc_results['foreground/'+ifo+'/chisq'] = 1.
    coinc_results['foreground/'+ifo+'/chisq_dof'] = 1
    coinc_results['foreground/'+ifo+'/template_duration'] = \
        fp['template_duration'][()]
    followup_data[ifo] = {}
    followup_data[ifo]['psd'] = interpolate(data[ifo].psd, 0.25)
    followup_data[ifo]['snr_series'] = snr_series_dict[ifo]

    # Find loudest SNR within coincidence window of the largest peak
    snr_peak = abs(loudest_snrs[ifo])
    # And check light travel time
    lttbd = Detector(ifo).light_travel_time_to_detector(Detector(loudest_ifo))
    # Add small buffer
    lttbd += 0.005

    time_window = [loudest_snr_time_allifos-lttbd,
                   loudest_snr_time_allifos+lttbd]

    snr_slice = snr_series_dict[ifo].time_slice(time_window[0], time_window[1])
    max_snr_idx = numpy.argmax(abs(snr_slice))
    loudest_snr_time = snr_slice.sample_times[max_snr_idx]
    loudest_snr = snr_slice[max_snr_idx]

    coinc_results['foreground/'+ifo+'/end_time'] = loudest_snr_time
    coinc_results['foreground/'+ifo+'/snr_series'] = snr_series_dict[ifo]
    coinc_results['foreground/'+ifo+'/psd_series'] = data[ifo].psd
    coinc_results['foreground/'+ifo+'/delta_f'] = delta_f
    coinc_results['foreground/'+ifo+'/snr'] = abs(loudest_snr)
    netsnr += abs(loudest_snr)**2
    coinc_results['foreground/'+ifo+'/sigmasq'] \
        = snr_series_dict['sigmasq_' + ifo]
    coinc_results['foreground/'+ifo+'/coa_phase'] \
        = numpy.angle(loudest_snr)

coinc_results['foreground/stat'] = numpy.sqrt(netsnr)
coinc_results['foreground/ifar'] = fp['ifar'][()]

channel_names = {}
for ifo in fp['channel_names'].keys():
    channel_names[ifo] = fp['channel_names'][ifo].asstr()[()]

mc_area_args = load_hdf5_to_dict(fp, 'mc_area_args/')

kwargs = {'psds': {ifo: followup_data[ifo]['psd'] for ifo in ifos},
          'low_frequency_cutoff': flow,
          'followup_data': followup_data,
          'channel_names': channel_names,
          'mc_area_args': mc_area_args}

doc = live.SingleCoincForGraceDB(ifos, coinc_results,
                                 upload_snr_series=True, **kwargs)

if args.enable_gracedb_upload:
    comment = ('Automatic PyCBC followup of trigger '
               '<a href="/events/{0}/view">{0}</a> to find the template '
               'parameters that maximize the SNR. The FAR of this trigger is '
               'copied from {0} and does not reflect this triggers\' template '
               'parameters.')
    comment = comment.format(original_gid)

    tmpdir = tempfile.mkdtemp()
    xml_path = os.path.join(tmpdir,
                            '{:.3f}.xml.gz'.format(loudest_snr_time_allifos))
    gid = doc.upload(xml_path, gracedb_server=args.gracedb_server,
                     testing=(not args.production), extra_strings=[comment],
                     search=args.gracedb_search)
    if gid is not None:
        logging.info('Event uploaded as %s', gid)
        shutil.rmtree(tmpdir)

        # add a note to the original G event pointing to the optimized one
        from ligo.gracedb.rest import GraceDb

        gracedb = GraceDb(args.gracedb_server) \
                if args.gracedb_server is not None else GraceDb()
        comment = ('Result of SNR maximization uploaded as '
                   '<a href="/events/{0}/view">{0}</a>').format(gid)
        gracedb.writeLog(original_gid, comment, tag_name=['analyst_comments'])
else:
    fname = 'coinc-{:.3f}-{}.xml.gz'.format(loudest_snr_time_allifos,
                                            pycbc.random_string(6))
    fname = os.path.join(args.output_path, fname)
    doc.save(fname)