#
# Copyright (C) 2013-2020  Leo Singer
#
# 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 3 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, see <https://www.gnu.org/licenses/>.
#
"""
Rapid sky localization with BAYESTAR [1]_.

References
----------
.. [1] Singer & Price, 2016. "Rapid Bayesian position reconstruction for
       gravitational-wave transients." PRD, 93, 024013.
       :doi:`10.1103/PhysRevD.93.024013`

"""
import inspect
import logging
import os
import sys
from textwrap import wrap

from astropy.table import Column, Table
from astropy import units as u
import lal
import lalsimulation
import numpy as np

from .. import distance
from . import filter
from ..io.hdf5 import write_samples
from ..io.fits import metadata_for_version_module
from ..io.events.base import Event
from . import filter  # noqa
from ..kde import Clustered2Plus1DSkyKDE
from .. import moc
from .. import healpix_tree
from .. import version
from .. import core
from ..core import (antenna_factor, signal_amplitude_model,
                    log_posterior_toa_phoa_snr as _log_posterior_toa_phoa_snr)
from ..util.numpy import require_contiguous_aligned
from ..util.stopwatch import Stopwatch
from .ez_emcee import ez_emcee

__all__ = ('derasterize', 'localize', 'rasterize', 'antenna_factor',
           'signal_amplitude_model')

log = logging.getLogger('BAYESTAR')

_RESCALE_LOGLIKELIHOOD = 0.83

antenna_factor = require_contiguous_aligned(antenna_factor)
signal_amplitude_model = require_contiguous_aligned(signal_amplitude_model)
_log_posterior_toa_phoa_snr = require_contiguous_aligned(
    _log_posterior_toa_phoa_snr)


# Wrap so that ufunc parameter names are known
def log_posterior_toa_phoa_snr(
        ra, sin_dec, distance, u, twopsi, t, min_distance, max_distance,
        prior_distance_power, cosmology, gmst, sample_rate, epochs, snrs,
        responses, locations, horizons,
        rescale_loglikelihood=_RESCALE_LOGLIKELIHOOD):
    return _log_posterior_toa_phoa_snr(
        ra, sin_dec, distance, u, twopsi, t, min_distance, max_distance,
        prior_distance_power, cosmology, gmst, sample_rate, epochs, snrs,
        responses, locations, horizons, rescale_loglikelihood)


# Wrap so that fixed parameter values are pulled from keyword arguments
def log_post(params, *args, **kwargs):
    # Names of parameters
    keys = ('ra', 'sin_dec', 'distance', 'u', 'twopsi', 't')

    params = list(params.T)
    kwargs = dict(kwargs)

    return log_posterior_toa_phoa_snr(
        *(kwargs.pop(key) if key in kwargs else params.pop(0) for key in keys),
        *args, **kwargs)


def localize_emcee(args, xmin, xmax, chain_dump=None):
    # Gather posterior samples
    chain = ez_emcee(log_post, xmin, xmax, args=args, vectorize=True)

    # Transform back from sin_dec to dec and cos_inclination to inclination
    chain[:, 1] = np.arcsin(chain[:, 1])
    chain[:, 3] = np.arccos(chain[:, 3])

    # Optionally save posterior sample chain to file.
    if chain_dump:
        _, ndim = chain.shape
        names = 'ra dec distance inclination twopsi time'.split()[:ndim]
        write_samples(Table(rows=chain, names=names, copy=False), chain_dump,
                      path='/bayestar/posterior_samples', overwrite=True)

    # Pass a random subset of 1000 points to the KDE, to save time.
    pts = np.random.permutation(chain)[:1000, :3]
    ckde = Clustered2Plus1DSkyKDE(pts)
    return ckde.as_healpix()


def condition(
        event, waveform='o2-uberbank', f_low=30.0,
        enable_snr_series=True, f_high_truncate=0.95):

    if len(event.singles) == 0:
        raise ValueError('Cannot localize an event with zero detectors.')

    singles = event.singles
    if not enable_snr_series:
        singles = [single for single in singles if single.snr is not None]

    ifos = [single.detector for single in singles]

    # Extract SNRs from table.
    snrs = np.ma.asarray([
        np.ma.masked if single.snr is None else single.snr
        for single in singles])

    # Look up physical parameters for detector.
    detectors = [lalsimulation.DetectorPrefixToLALDetector(str(ifo))
                 for ifo in ifos]
    responses = np.asarray([det.response for det in detectors])
    locations = np.asarray([det.location for det in detectors]) / lal.C_SI

    # Power spectra for each detector.
    psds = [single.psd for single in singles]
    psds = [filter.InterpolatedPSD(filter.abscissa(psd), psd.data.data,
                                   f_high_truncate=f_high_truncate)
            for psd in psds]

    log.debug('calculating templates')
    H = filter.sngl_inspiral_psd(waveform, f_min=f_low, **event.template_args)

    log.debug('calculating noise PSDs')
    HS = [filter.signal_psd_series(H, S) for S in psds]

    # Signal models for each detector.
    log.debug('calculating Fisher matrix elements')
    signal_models = [filter.SignalModel(_) for _ in HS]

    # Get SNR=1 horizon distances for each detector.
    horizons = np.asarray([signal_model.get_horizon_distance()
                           for signal_model in signal_models])

    weights = np.ma.asarray([
        1 / np.square(signal_model.get_crb_toa_uncert(snr))
        for signal_model, snr in zip(signal_models, snrs)])

    # Center detector array.
    locations -= (np.sum(locations * weights.reshape(-1, 1), axis=0) /
                  np.sum(weights))

    if enable_snr_series:
        snr_series = [single.snr_series for single in singles]
        if all(s is None for s in snr_series):
            snr_series = None
    else:
        snr_series = None

    # Maximum barycentered arrival time error:
    # |distance from array barycenter to furthest detector| / c + 5 ms.
    # For LHO+LLO, this is 15.0 ms.
    # For an arbitrary terrestrial detector network, the maximum is 26.3 ms.
    max_abs_t = np.max(
        np.sqrt(np.sum(np.square(locations), axis=1))) + 0.005

    # Calculate a minimum arrival time prior for low-bandwidth signals.
    # This is important for early-warning events, for which the light travel
    # time between detectors may be only a tiny fraction of the template
    # autocorrelation time.
    #
    # The period of the autocorrelation function is 2 pi times the timing
    # uncertainty at SNR=1 (see Eq. (24) of [1]_). We want a quarter period:
    # there is a factor of a half because we want the first zero crossing,
    # and another factor of a half because the SNR time series will be cropped
    # to a two-sided interval.
    max_acor_t = max(0.5 * np.pi * signal_model.get_crb_toa_uncert(1)
                     for signal_model in signal_models)
    max_abs_t = max(max_abs_t, max_acor_t)

    if snr_series is None:
        log.warning("No SNR time series found, so we are creating a "
                    "zero-noise SNR time series from the whitened template's "
                    "autocorrelation sequence. The sky localization "
                    "uncertainty may be underestimated.")

        acors, sample_rates = zip(
            *[filter.autocorrelation(_, max_abs_t) for _ in HS])
        sample_rate = sample_rates[0]
        deltaT = 1 / sample_rate
        nsamples = len(acors[0])
        assert all(sample_rate == _ for _ in sample_rates)
        assert all(nsamples == len(_) for _ in acors)
        nsamples = nsamples * 2 - 1

        snr_series = []
        for acor, single in zip(acors, singles):
            series = lal.CreateCOMPLEX8TimeSeries(
                'fake SNR', 0, 0, deltaT, lal.StrainUnit, nsamples)
            series.epoch = single.time - 0.5 * (nsamples - 1) * deltaT
            acor = np.concatenate((np.conj(acor[:0:-1]), acor))
            series.data.data = single.snr * filter.exp_i(single.phase) * acor
            snr_series.append(series)

    # Ensure that all of the SNR time series have the same sample rate.
    # FIXME: for now, the Python wrapper expects all of the SNR time series to
    # also be the same length.
    deltaT = snr_series[0].deltaT
    sample_rate = 1 / deltaT
    if any(deltaT != series.deltaT for series in snr_series):
        raise ValueError('BAYESTAR does not yet support SNR time series with '
                         'mixed sample rates')

    # Ensure that all of the SNR time series have odd lengths.
    if any(len(series.data.data) % 2 == 0 for series in snr_series):
        raise ValueError('SNR time series must have odd lengths')

    # Trim time series to the desired length.
    max_abs_n = int(np.ceil(max_abs_t * sample_rate))
    desired_length = 2 * max_abs_n - 1
    for i, series in enumerate(snr_series):
        length = len(series.data.data)
        if length > desired_length:
            snr_series[i] = lal.CutCOMPLEX8TimeSeries(
                series, length // 2 + 1 - max_abs_n, desired_length)

    # FIXME: for now, the Python wrapper expects all of the SNR time sries to
    # also be the same length.
    nsamples = len(snr_series[0].data.data)
    if any(nsamples != len(series.data.data) for series in snr_series):
        raise ValueError('BAYESTAR does not yet support SNR time series of '
                         'mixed lengths')

    # Perform sanity checks that the middle sample of the SNR time series match
    # the sngl_inspiral records to the nearest sample (plus the smallest
    # representable LIGOTimeGPS difference of 1 nanosecond).
    for ifo, single, series in zip(ifos, singles, snr_series):
        shift = np.abs(0.5 * (nsamples - 1) * series.deltaT +
                       float(series.epoch - single.time))
        if shift >= deltaT + 1e-8:
            raise ValueError('BAYESTAR expects the SNR time series to be '
                             'centered on the single-detector trigger times, '
                             'but {} was off by {} s'.format(ifo, shift))

    # Extract the TOAs in GPS nanoseconds from the SNR time series, assuming
    # that the trigger happened in the middle.
    toas_ns = [series.epoch.ns() + 1e9 * 0.5 * (len(series.data.data) - 1) *
               series.deltaT for series in snr_series]

    # Collect all of the SNR series in one array.
    snr_series = np.vstack([series.data.data for series in snr_series])

    # Center times of arrival and compute GMST at mean arrival time.
    # Pre-center in integer nanoseconds to preserve precision of
    # initial datatype.
    epoch = sum(toas_ns) // len(toas_ns)
    toas = 1e-9 * (np.asarray(toas_ns) - epoch)
    mean_toa = np.average(toas, weights=weights)
    toas -= mean_toa
    epoch += int(np.round(1e9 * mean_toa))
    epoch = lal.LIGOTimeGPS(0, int(epoch))

    # Translate SNR time series back to time of first sample.
    toas -= 0.5 * (nsamples - 1) * deltaT

    # Convert complex SNRS to amplitude and phase
    snrs_abs = np.abs(snr_series)
    snrs_arg = filter.unwrap(np.angle(snr_series))
    snrs = np.stack((snrs_abs, snrs_arg), axis=-1)

    return epoch, sample_rate, toas, snrs, responses, locations, horizons


def condition_prior(horizons, min_distance=None, max_distance=None,
                    prior_distance_power=None, cosmology=False):
    if cosmology:
        log.warning('Enabling cosmological prior. This feature is UNREVIEWED.')

    # If minimum distance is not specified, then default to 0 Mpc.
    if min_distance is None:
        min_distance = 0

    # If maximum distance is not specified, then default to the SNR=4
    # horizon distance of the most sensitive detector.
    if max_distance is None:
        max_distance = max(horizons) / 4

    # If prior_distance_power is not specified, then default to 2
    # (p(r) ~ r^2, uniform in volume).
    if prior_distance_power is None:
        prior_distance_power = 2

    # Raise an exception if 0 Mpc is the minimum effective distance and the
    # prior is of the form r**k for k<0
    if min_distance == 0 and prior_distance_power < 0:
        raise ValueError(('Prior is a power law r^k with k={}, '
                          'undefined at min_distance=0').format(
                              prior_distance_power))

    return min_distance, max_distance, prior_distance_power, cosmology


def localize(
        event, waveform='o2-uberbank', f_low=30.0,
        min_inclination=0, max_inclination=np.pi / 2,
        min_distance=None, max_distance=None, prior_distance_power=None,
        cosmology=False, mcmc=False, chain_dump=None,
        enable_snr_series=True, f_high_truncate=0.95,
        rescale_loglikelihood=_RESCALE_LOGLIKELIHOOD):
    """Localize a compact binary signal using the BAYESTAR algorithm.

    Parameters
    ----------
    event : `ligo.skymap.io.events.Event`
        The event candidate.
    waveform : str, optional
        The name of the waveform approximant.
    f_low : float, optional
        The low frequency cutoff.
    min_distance, max_distance : float, optional
        The limits of integration over luminosity distance, in Mpc
        (default: determine automatically from detector sensitivity).
    prior_distance_power : int, optional
        The power of distance that appears in the prior
        (default: 2, uniform in volume).
    cosmology: bool, optional
        Set to enable a uniform in comoving volume prior (default: false).
    mcmc : bool, optional
        Set to use MCMC sampling rather than more accurate Gaussian quadrature.
    chain_dump : str, optional
        Save posterior samples to this filename if `mcmc` is set.
    enable_snr_series : bool, optional
        Set to False to disable SNR time series.
    f_high_truncate : float, optional
        Truncate the noise power spectral densities at this factor times the
        highest sampled frequency to suppress artifacts caused by incorrect
        PSD conditioning by some matched filter pipelines.

    Returns
    -------
    skymap : `astropy.table.Table`
        A 3D sky map in multi-order HEALPix format.

    """
    # Hide event parameters, but show all other arguments
    def formatvalue(value):
        if isinstance(value, Event):
            return '=...'
        else:
            return '=' + repr(value)

    frame = inspect.currentframe()
    argstr = inspect.formatargvalues(*inspect.getargvalues(frame),
                                     formatvalue=formatvalue)

    stopwatch = Stopwatch()
    stopwatch.start()

    epoch, sample_rate, toas, snrs, responses, locations, horizons = \
        condition(event, waveform=waveform, f_low=f_low,
                  enable_snr_series=enable_snr_series,
                  f_high_truncate=f_high_truncate)

    min_distance, max_distance, prior_distance_power, cosmology = \
        condition_prior(horizons, min_distance, max_distance,
                        prior_distance_power, cosmology)

    gmst = lal.GreenwichMeanSiderealTime(epoch)

    # Time and run sky localization.
    log.debug('starting computationally-intensive section')
    if mcmc:
        max_abs_t = 2 * snrs.data.shape[1] / sample_rate
        if min_inclination != 0 or max_inclination != np.pi / 2:
            log.warn('inclination limits are not supported for MCMC mode')
        args = (min_distance, max_distance, prior_distance_power, cosmology,
                gmst, sample_rate, toas, snrs, responses, locations, horizons,
                rescale_loglikelihood)
        skymap = localize_emcee(
            args=args,
            xmin=[0, -1, min_distance, -1, 0, 0],
            xmax=[2 * np.pi, 1, max_distance, 1, 2 * np.pi, 2 * max_abs_t],
            chain_dump=chain_dump)
    else:
        args = (min_inclination, max_inclination, min_distance, max_distance,
                prior_distance_power, cosmology, gmst, sample_rate, toas, snrs,
                responses, locations, horizons, rescale_loglikelihood)
        skymap, log_bci, log_bsn = core.toa_phoa_snr(*args)
        skymap = Table(skymap, copy=False)
        skymap.meta['log_bci'] = log_bci
        skymap.meta['log_bsn'] = log_bsn

    # Convert distance moments to parameters
    try:
        distmean = skymap.columns.pop('DISTMEAN')
        diststd = skymap.columns.pop('DISTSTD')
    except KeyError:
        distmean, diststd, _ = distance.parameters_to_moments(
            skymap['DISTMU'], skymap['DISTSIGMA'])
    else:
        skymap['DISTMU'], skymap['DISTSIGMA'], skymap['DISTNORM'] = \
            distance.moments_to_parameters(distmean, diststd)

    # Add marginal distance moments
    good = np.isfinite(distmean) & np.isfinite(diststd)
    prob = (moc.uniq2pixarea(skymap['UNIQ']) * skymap['PROBDENSITY'])[good]
    distmean = distmean[good]
    diststd = diststd[good]
    rbar = (prob * distmean).sum()
    r2bar = (prob * (np.square(diststd) + np.square(distmean))).sum()
    skymap.meta['distmean'] = rbar
    skymap.meta['diststd'] = np.sqrt(r2bar - np.square(rbar))

    stopwatch.stop()
    end_time = lal.GPSTimeNow()
    log.info('finished computationally-intensive section in %s', stopwatch)

    # Fill in metadata and return.
    program, _ = os.path.splitext(os.path.basename(sys.argv[0]))
    skymap.meta.update(metadata_for_version_module(version))
    skymap.meta['creator'] = 'BAYESTAR'
    skymap.meta['origin'] = 'LIGO/Virgo'
    skymap.meta['gps_time'] = float(epoch)
    skymap.meta['runtime'] = stopwatch.real
    skymap.meta['instruments'] = {single.detector for single in event.singles}
    skymap.meta['gps_creation_time'] = end_time
    skymap.meta['history'] = [
        '',
        'Generated by calling the following Python function:',
        *wrap('{}.{}{}'.format(__name__, frame.f_code.co_name, argstr), 72),
        '',
        'This was the command line that started the program:',
        *wrap(' '.join([program] + sys.argv[1:]), 72)]

    return skymap


def rasterize(skymap, order=None):
    orig_order, _ = moc.uniq2nest(skymap['UNIQ'].max())

    # Determine whether we need to do nontrivial downsampling.
    downsampling = (order is not None and 0 <= order < orig_order
                    and 'DISTMU' in skymap.dtype.fields.keys())

    # If we are downsampling, then convert from distance parameters to
    # distance moments times probability density so that the averaging
    # that is automatically done by moc.rasterize() correctly marginalizes
    # the moments.
    if downsampling:
        skymap = Table(skymap, copy=True, meta=skymap.meta)

        probdensity = skymap['PROBDENSITY']
        distmu = skymap.columns.pop('DISTMU')
        distsigma = skymap.columns.pop('DISTSIGMA')

        bad = ~(np.isfinite(distmu) & np.isfinite(distsigma))
        distmean, diststd, _ = distance.parameters_to_moments(
            distmu, distsigma)
        distmean[bad] = np.nan
        diststd[bad] = np.nan
        skymap['DISTMEAN'] = probdensity * distmean
        skymap['DISTVAR'] = probdensity * (
            np.square(diststd) + np.square(distmean))

    skymap = Table(moc.rasterize(skymap, order=order),
                   meta=skymap.meta, copy=False)

    # If we are downsampling, then convert back to distance parameters.
    if downsampling:
        distmean = skymap.columns.pop('DISTMEAN') / skymap['PROBDENSITY']
        diststd = np.sqrt(
            skymap.columns.pop('DISTVAR') / skymap['PROBDENSITY']
            - np.square(distmean))
        skymap['DISTMU'], skymap['DISTSIGMA'], skymap['DISTNORM'] = \
            distance.moments_to_parameters(
                distmean, diststd)

    skymap.rename_column('PROBDENSITY', 'PROB')
    skymap['PROB'] *= 4 * np.pi / len(skymap)
    skymap['PROB'].unit = u.pixel ** -1
    return skymap


def derasterize(skymap):
    skymap.rename_column('PROB', 'PROBDENSITY')
    skymap['PROBDENSITY'] *= len(skymap) / (4 * np.pi)
    skymap['PROBDENSITY'].unit = u.steradian ** -1
    nside, _, ipix, _, _, value = zip(
        *healpix_tree.reconstruct_nested(skymap))
    nside = np.asarray(nside)
    ipix = np.asarray(ipix)
    value = np.stack(value)
    uniq = (4 * np.square(nside) + ipix)
    old_units = [column.unit for column in skymap.columns.values()]
    skymap = Table(value, meta=skymap.meta, copy=False)
    for old_unit, column in zip(old_units, skymap.columns.values()):
        column.unit = old_unit
    skymap.add_column(Column(uniq, name='UNIQ'), 0)
    skymap.sort('UNIQ')
    return skymap


def test():
    """Run BAYESTAR C unit tests.

    Examples
    --------
    >>> test()
    0

    """
    return int(core.test())


del require_contiguous_aligned
