# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""

# Need to put the functions in the path
# Probably not necessary if I understood Python/Git/modules better
import os
import sys
parent_directory = os.getcwd().rsplit(sep='\\',maxsplit=1)[0]
if parent_directory not in sys.path:
    sys.path.insert(1, parent_directory)

import initElements_P3
ed = initElements_P3.initElements()

import numpy as np

def In_doped_GaN():
    #                           N      Ga      In  Da
    # Define possible peaks
    pk_data =   np.array(    [  (1,     0,      0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,      0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,      0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,      0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0]),
                                (0,     0,      1,  ed['In'].isotopes[113][0]/2),
                                (0,     0,      1,  ed['In'].isotopes[115][0]/2),
                                (0,     0,      1,  ed['In'].isotopes[113][0]),
                                (0,     0,      1,  ed['In'].isotopes[115][0]),
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('In','i4'),('m2q','f4')] )
    return pk_data

def Mg_doped_GaN():
    #                            N      Ga      Mg  Da
    # Define possible peaks
    pk_data =   np.array(    [  (1,     0,      0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,      0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,      0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,      0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0]),
                                (0,     0,      1,  ed['Mg'].isotopes[24][0]/2),
                                (0,     0,      1,  ed['Mg'].isotopes[24][0])
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('Mg','i4'),('m2q','f4')] )
    return pk_data

def AlGaN():
    #                           N      Ga      Al  Da
    # Define possible peaks
    pk_data =   np.array(    [  (1,     0,      0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,      0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,      0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,      0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0]),
                                (0,     0,      1,  ed['Al'].isotopes[27][0]/3),
                                (0,     0,      1,  ed['Al'].isotopes[27][0]/2),
                                (1,     0,      1,  (ed['Al'].isotopes[27][0]+ed['N'].isotopes[14][0])/2),
                                (0,     0,      1,  ed['Al'].isotopes[27][0]/1)
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('Al','i4'),('m2q','f4')] )
    return pk_data

def GaN():
    #                           N      Ga  Da
    # Define possible peaks
    pk_data =   np.array(    [  (1,     0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,  ed['Ga'].isotopes[69][0]),
                                (0,     1,  ed['Ga'].isotopes[71][0]),
                                (0,     1,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0])
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('m2q','f4')] )
    return pk_data

def GaN_with_H():
    #                           N      Ga  Da
    # Define possible peaks
    pk_data =   np.array(    [  (0,     0,  ed['H'].isotopes[1][0]),
                                (0,     0,  ed['H'].isotopes[1][0]*2),
                                (0,     0,  ed['H'].isotopes[1][0]*3),
                                (1,     0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,  ed['Ga'].isotopes[69][0]),
                                (0,     1,  ed['Ga'].isotopes[71][0]),
                                (0,     1,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0])
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('m2q','f4')] )
    return pk_data


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#    Project: Azimuthal integration
#             https://github.com/silx-kit/pyFAI
#
#    Copyright (C) 2014-2018 European Synchrotron Radiation Facility, Grenoble, France
#
#    Principal author:       Jérôme Kieffer (Jerome.Kieffer@ESRF.eu)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

__author__ = "Jérôme Kieffer"
__contact__ = "Jerome.Kieffer@ESRF.eu"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
__date__ = "17/10/2018"
__status__ = "production"

import sys
import os
import threading
from math import ceil, sqrt
import logging
logger = logging.getLogger(__name__)
import numpy
import fabio
from scipy.ndimage import label, distance_transform_edt
from scipy.ndimage.filters import median_filter
from .utils.decorators import deprecated
from .ext.bilinear import Bilinear
from .utils import gaussian_filter, binning, unbinning, is_far_from_group
from .third_party import six

if os.name != "nt":
    WindowsError = RuntimeError

class Massif(object):
    """
    A massif is defined as an area around a peak, it is used to find neighboring peaks
    """
    TARGET_SIZE = 1024

    def __init__(self, data=None, mask=None):
        """Constructor of the class...

        :param data: 2D array or filename (discouraged)
        :param mask: array with non zero for invalid data
        """
        if isinstance(data, six.string_types) and os.path.isfile(data):
            self.data = fabio.open(data).data.astype("float32")
        elif isinstance(data, fabio.fabioimage.fabioimage):
            self.data = data.data.astype("float32")
        else:
            try:
                self.data = data.astype("float32")
            except Exception as error:
                logger.error("Unable to understand this type of data %s: %s", data, error)
        self.mask = mask
        self._cleaned_data = None
        self._bilin = Bilinear(self.data)
        self._blurred_data = None
        self._median_data = None
        self._labeled_massif = None
        self._number_massif = None
        self._valley_size = None
        self._binned_data = None
        self._reconstruct_used = None
        self.binning = None  # Binning is 2-list usually
        self._sem = threading.Semaphore()
        self._sem_label = threading.Semaphore()
        self._sem_binning = threading.Semaphore()
        self._sem_median = threading.Semaphore()

    def nearest_peak(self, x):
        """
        :param x: coordinates of the peak
        :returns: the coordinates of the nearest peak
        """
        out = self._bilin.local_maxi(x)
        if isinstance(out, tuple):
            res = out
        elif isinstance(out, numpy.ndarray):
            res = tuple(out)
        else:
            res = [int(i) for idx, i in enumerate(out) if 0 <= i < self.data.shape[idx]]
        if (len(res) != 2) or not((0 <= out[0] < self.data.shape[0]) and (0 <= res[1] < self.data.shape[1])):
            logger.warning("in nearest_peak %s -> %s", x, out)
            return
        elif (self.mask is not None) and self.mask[int(res[0]), int(res[1])]:
            logger.info("Masked pixel %s -> %s", x, out)
            return
        else:
            return res

    def calculate_massif(self, x):
        """
        defines a map of the massif around x and returns the mask
        """
        labeled = self.get_labeled_massif()
        if labeled[x[0], x[1]] > 0:  # without relabeled the background is 0 labeled.max():
            return (labeled == labeled[x[0], x[1]])

    def find_peaks(self, x, nmax=200, annotate=None, massif_contour=None, stdout=sys.stdout):
        """
        All in one function that finds a maximum from the given seed (x)
        then calculates the region extension and extract position of the neighboring peaks.
        :param x: coordinates of the peak, seed for the calculation
        :type x: tuple of integer
        :param nmax: maximum number of peak per region
        :param annotate: call back method taking number of points + coordinate as input.
        :param massif_contour: callback to show the contour of a massif with the given index.
        :param stdout: this is the file where output is written by default.
        :return: list of peaks
        """
        listpeaks = []
        region = self.calculate_massif(x)
        if region is None:
            logger.error("You picked a background point at %s", x)
            return listpeaks
        xinit = self.nearest_peak(x)
        if xinit is None:
            logger.error("Unable to find peak in the vinicy of %s", x)
            return listpeaks
        else:
            if not region[int(xinit[0] + 0.5), int(xinit[1] + 0.5)]:
                logger.error("Nearest peak %s is not in the same region  %s", xinit, x)
                return listpeaks

            if annotate is not None:
                try:
                    annotate(xinit, x)
                except Exception as error:
                    logger.error("Error in annotate %i: %i %i. %s", len(listpeaks), xinit[0], xinit[1], error)

        listpeaks.append(xinit)
        cleaned_data = self.cleaned_data
        mean = cleaned_data[region].mean(dtype=numpy.float64)
        region2 = region * (cleaned_data > mean)
        idx = numpy.vstack(numpy.where(region2)).T
        numpy.random.shuffle(idx)
        nmax = min(nmax, int(ceil(sqrt(idx.shape[0]))))
        if massif_contour is not None:
            try:
                massif_contour(region)
            except (WindowsError, MemoryError) as error:
                logger.error("Error in plotting region: %s", error)
        nbFailure = 0
        for j in idx:
            xopt = self.nearest_peak(j)
            if xopt is None:
                nbFailure += 1
                continue
            if (region2[int(xopt[0] + 0.5), int(xopt[1] + 0.5)]) and not (xopt in listpeaks):
                stdout.write("[ %4i, %4i ] --> [ %5.1f, %5.1f ] after %3i iterations %s" % (tuple(j) + tuple(xopt) + (nbFailure, os.linesep)))
                listpeaks.append(xopt)
                nbFailure = 0
            else:
                nbFailure += 1
            if (len(listpeaks) > nmax) or (nbFailure > 2 * nmax):
                break
        return listpeaks

    def peaks_from_area(self, mask, Imin=numpy.finfo(numpy.float64).min,
                        keep=1000, dmin=0.0, seed=None, **kwarg):
        """
        Return the list of peaks within an area

        :param mask: 2d array with mask.
        :param Imin: minimum of intensity above the background to keep the point
        :param keep: maximum number of points to keep
        :param kwarg: ignored parameters
        :param dmin: minimum distance to another peak
        :param seed: list of good guesses to start with
        :return: list of peaks [y,x], [y,x], ...]
        """
        all_points = numpy.vstack(numpy.where(mask)).T
        res = []
        cnt = 0
        dmin2 = dmin * dmin
        if len(all_points) > 0:
            numpy.random.shuffle(all_points)
        if seed:
            seeds = numpy.array(list(seed))
            if len(seeds) > 0:
                numpy.random.shuffle(seeds)
            all_points = numpy.concatenate((seeds, all_points))
        for idx in all_points:
            out = self.nearest_peak(idx)
            if out is not None:
                msg = "[ %3i, %3i ] -> [ %.1f, %.1f ]"
                logger.debug(msg, idx[1], idx[0], out[1], out[0])
                p0, p1 = int(round(out[0])), int(round(out[1]))
                if mask[p0, p1]:
                    if (self.data[p0, p1] > Imin) and is_far_from_group(out, res, dmin2):
                        res.append(out)
                        cnt = 0
            if len(res) >= keep or cnt > keep:
                break
            else:
                cnt += 1
        return res

    def init_valley_size(self):
        if self._valley_size is None:
            self.valley_size = max(5., max(self.data.shape) / 50.)

    @property
    def valley_size(self):
        "Defines the minimum distance between two massifs"
        if self._valley_size is None:
            self.init_valley_size()
        return self._valley_size

    @valley_size.setter
    def valley_size(self, size):
        new_size = float(size)
        if self._valley_size != new_size:
            self._valley_size = new_size
            t = threading.Thread(target=self.get_labeled_massif)
            t.start()

    @valley_size.deleter
    def valley_size(self):
        self._valley_size = None
        self._blurred_data = None

    @property
    def cleaned_data(self):
        if self.mask is None:
            return self.data
        else:
            if self._cleaned_data is None:
                idx = distance_transform_edt(self.mask,
                                             return_distances=False,
                                             return_indices=True)
                self._cleaned_data = self.data[tuple(idx)]
            return self._cleaned_data

    def get_binned_data(self):
        """
        :return: binned data
        """
        if self._binned_data is None:
            with self._sem_binning:
                if self._binned_data is None:
                    logger.info("Image size is %s", self.data.shape)
                    self.binning = []
                    for i in self.data.shape:
                        if i % self.TARGET_SIZE == 0:
                            self.binning.append(max(1, i // self.TARGET_SIZE))
                        else:
                            for j in range(i // self.TARGET_SIZE - 1, 0, -1):
                                if i % j == 0:
                                    self.binning.append(max(1, j))
                                    break
                            else:
                                self.binning.append(1)
#                    self.binning = max([max(1, i // self.TARGET_SIZE) for i in self.data.shape])
                    logger.info("Binning size is %s", self.binning)
                    self._binned_data = binning(self.cleaned_data, self.binning)
        return self._binned_data

    def get_median_data(self):
        """
        :return: a spatial median filtered image 3x3
        """
        if self._median_data is None:
            with self._sem_median:
                if self._median_data is None:
                    self._median_data = median_filter(self.cleaned_data, 3)
        return self._median_data

    def get_blurred_data(self):
        """
        :return: a blurred image
        """
        if self._blurred_data is None:
            with self._sem:
                if self._blurred_data is None:
                    logger.debug("Blurring image with kernel size: %s", self.valley_size)
                    self._blurred_data = gaussian_filter(self.get_binned_data(),
                                                         [self.valley_size / i for i in self.binning],
                                                         mode="reflect")
        return self._blurred_data

    def get_labeled_massif(self, pattern=None, reconstruct=True):
        """
        :param pattern: 3x3 matrix
        :param reconstruct: if False, split massif at masked position, else reconstruct missing part.
        :return: an image composed of int with a different value for each massif
        """
        if self._labeled_massif is None:
            with self._sem_label:
                if self._labeled_massif is None:
                    if pattern is None:
                        pattern = numpy.ones((3, 3), dtype=numpy.int8)
                    logger.debug("Labeling all massifs. This takes some time !!!")
                    massif_binarization = (self.get_binned_data() > self.get_blurred_data())
                    if (self.mask is not None) and (not reconstruct):
                            binned_mask = binning(self.mask.astype(int), self.binning, norm=False)
                            massif_binarization = numpy.logical_and(massif_binarization, binned_mask == 0)
                    self._reconstruct_used = reconstruct
                    labeled_massif, self._number_massif = label(massif_binarization,
                                                                pattern)
                    # TODO: investigate why relabel fails
                    # relabeled = relabel(labeled_massif, self.get_binned_data(), self.get_blurred_data())
                    relabeled = labeled_massif
                    self._labeled_massif = unbinning(relabeled, self.binning, False)
                    logger.info("Labeling found %s massifs.", self._number_massif)
        return self._labeled_massif

    @deprecated(reason="switch to pep8 style", replacement="init_valley_size", since_version="0.16.0")
    def initValleySize(self):
        self.init_valley_size()

    @deprecated(reason="switch to PEP8 style", replacement="get_median_data", since_version="0.16.0")
    def getMedianData(self):
        return self.get_median_data()

    @deprecated(reason="switch to PEP8 style", replacement="get_binned_data", since_version="0.16.0")
    def getBinnedData(self):
        return self.get_binned_data()

    @deprecated(reason="switch to PEP8 style", replacement="get_blurred_data", since_version="0.16.0")
    def getBluredData(self):
        return self.get_blurred_data()

    @deprecated(reason="switch to PEP8 style", replacement="get_labeled_massif", since_version="0.16.0")
    def getLabeledMassif(self, pattern=None):
        return self.get_labeled_massif(pattern)


import numpy as np
from astropy.io import fits
from astropy.table import Table, vstack
from astropy.wcs import WCS
import os
import argparse
import logging, traceback
import pandas as pd
from copy import copy

from bkg_rate_estimation import get_avg_lin_cub_rate_quad_obs
from config import quad_dicts, EBINS0, EBINS1,\
                solid_angle_dpi_fname, bright_source_table_fname
from sqlite_funcs import write_rate_fits_from_obj, get_conn
from dbread_funcs import get_info_tab, guess_dbfname, get_files_tab
from event2dpi_funcs import filter_evdata
from models import Bkg_Model_wFlatA, CompoundModel, Point_Source_Model_Binned_Rates
from LLH import LLH_webins
from minimizers import NLLH_ScipyMinimize, NLLH_ScipyMinimize_Wjacob
from ray_trace_funcs import RayTraces
from coord_conv_funcs import convert_radec2imxy
from gti_funcs import add_bti2gti, bti2gti, gti2bti, union_gtis
from wcs_funcs import world2val

def cli():
    parser = argparse.ArgumentParser()
    parser.add_argument('--evfname', type=str,\
            help="Event data file",
            default=None)
    parser.add_argument('--dmask', type=str,\
            help="Detmask fname",
            default=None)
    parser.add_argument('--dbfname', type=str,\
            help="Name to save the database to",\
            default=None)
    parser.add_argument('--job_id', type=int,\
            help="Job ID",\
            default=0)
    parser.add_argument('--Njobs', type=int,\
            help="Number of jobs",\
            default=1)
    parser.add_argument('--twind', type=float,\
            help="Number of seconds to go +/- from the trigtime",\
            default=20)
    parser.add_argument('--bkg_dur', type=float,\
            help="bkg duration",\
            default=40.0)
    parser.add_argument('--bkg_nopost',\
            help="Don't use time after signal window for bkg",\
            action='store_true')
    parser.add_argument('--bkg_nopre',\
            help="Don't use time before signal window for bkg",\
            action='store_true')
    parser.add_argument('--pcfname', type=str,\
            help="partial coding file name",\
            default='pc_2.img')

    args = parser.parse_args()
    return args

def ang_sep(ra0, dec0, ra1, dec1):

    dcos = np.cos(np.radians(np.abs(ra0 - ra1)))
    angsep = np.arccos(np.cos(np.radians(90-dec0))*np.cos(np.radians(90-dec1)) +\
                        np.sin(np.radians(90-dec0))*np.sin(np.radians(90-dec1))*dcos)
    return np.rad2deg(angsep)

def im_dist(imx0, imy0, imx1, imy1):
    return np.hypot((imx1 - imx0), (imy1 - imy0))

def add_imxy2src_tab(src_tab, attfile, t0):

    att_ind = np.argmin(np.abs(attfile['TIME'] - t0))
    att_quat = attfile['QPARAM'][att_ind]
    pnt_ra, pnt_dec = attfile['POINTING'][att_ind,:2]
    imxs = np.zeros(len(src_tab))
    imys = np.zeros(len(src_tab))
    src_tab['PntSep'] = ang_sep(pnt_ra, pnt_dec, src_tab['RAJ2000'], src_tab['DEJ2000'])
    for i in range(len(imxs)):
        if src_tab['PntSep'][i] > 80.0:
            imxs[i], imys[i] = np.nan, np.nan
            continue
        imxs[i], imys[i] = convert_radec2imxy(src_tab['RAJ2000'][i],\
                                              src_tab['DEJ2000'][i],\
                                              att_quat)
    src_tab['imx'] = imxs
    src_tab['imy'] = imys
    return src_tab

def get_srcs_infov(attfile, t0, pcfname=None, pcmin=5e-2):

    brt_src_tab = Table.read(bright_source_table_fname)
    add_imxy2src_tab(brt_src_tab, attfile, t0)
    bl_infov = (np.abs(brt_src_tab['imy'])<.95)&(np.abs(brt_src_tab['imx'])<1.75)
    if pcfname is not None:
        PC = fits.open(pcfname)[0]
        pc = PC.data
        w_t = WCS(PC.header, key='T')
        pcvals = world2val(w_t, pc, brt_src_tab['imx'],\
                            brt_src_tab['imy'])
        bl_infov = bl_infov&(pcvals>=pcmin)
    N_infov = np.sum(bl_infov)
    return brt_src_tab[bl_infov]

def bkg_withPS_fit(PS_tab, model, llh_obj, t0s, t1s,\
                   dimxy=2e-3, im_steps=5, test_null=False):

    Nps = len(PS_tab)
    imax = np.linspace(-dimxy, dimxy, im_steps)
    if im_steps == 3:
        imax = np.linspace(-dimxy/2., dimxy/2., im_steps)
    elif im_steps == 2:
        imax = np.linspace(-dimxy/2., dimxy/2., im_steps)
    elif im_steps == 1:
        imax = np.array([0.0])

    imlist = []
    for i in range(Nps):
        imlist += [imax,imax]
    imgs = np.meshgrid(*imlist)
    Npnts = imgs[0].size

    bkg_miner = NLLH_ScipyMinimize_Wjacob('')
    bkg_miner.set_llh(llh_obj)
    llh_obj.set_time(t0s,t1s)

    bf_params_list = []
    bkg_nllhs = np.zeros(Npnts)

    nebins = model.nebins

    for i in range(Npnts):
        bf_params = {}
        im_names = []
        for j in range(Nps):
            row = PS_tab[j]
            psname = row['Name']
            bf_params[psname+'_imx'] = imgs[2*j].ravel()[i] + row['imx']
            bf_params[psname+'_imy'] = imgs[2*j+1].ravel()[i] + row['imy']
            im_names = [psname+'_imx', psname+'_imy']
        im_vals = [bf_params[nm] for nm in im_names]

        for e0 in range(nebins):
            bkg_miner.set_fixed_params(bkg_miner.param_names)
            bkg_miner.set_fixed_params(im_names, im_vals)
            e0_pnames = []
            for pname in bkg_miner.param_names:
                try:
                    if int(pname[-1])==e0:
                        e0_pnames.append(pname)
                except:
                    pass
            bkg_miner.set_fixed_params(e0_pnames, fixed=False)
            llh_obj.set_ebin(e0)

            bf_vals, bkg_nllh, res = bkg_miner.minimize()
            bkg_nllhs[i] += bkg_nllh[0]
            for ii, pname in enumerate(e0_pnames):
                bf_params[pname] = bf_vals[0][ii]
        bf_params_list.append(bf_params)

    bf_ind = np.argmin(bkg_nllhs)
    bf_params = bf_params_list[bf_ind]
    bf_nllh = bkg_nllhs[bf_ind]

    if test_null:
        TS_nulls = {}
        for i in range(Nps):
            params_ = copy(bf_params)
            row = PS_tab[i]
            psname = row['Name']
            for j in range(nebins):
                params_[psname+'_rate_'+str(j)] = 0.0
            llh_obj.set_ebin(-1)
            nllh_null = -llh_obj.get_logprob(params_)
            TS_nulls[psname] = np.sqrt(2.*(nllh_null - bf_nllh))
            if np.isnan(TS_nulls[psname]):
                TS_nulls[psname] = 0.0
        return bf_nllh, bf_params, TS_nulls

    return bf_nllh, bf_params

def do_init_bkg_wPSs(bkg_mod, llh_obj, src_tab, rt_obj, GTI, sig_twind):

    gti_bkg = add_bti2gti(sig_twind, GTI)
    bkg_t0s = gti_bkg['START']
    bkg_t1s = gti_bkg['STOP']

    Nsrcs = len(src_tab)
    nebins = bkg_mod.nebins

    for ii in range(Nsrcs):
        mod_list = [bkg_mod]
        im_steps = 5
        if Nsrcs >= 3:
            im_steps = 3
        if Nsrcs >= 5:
            im_steps = 2
        if Nsrcs >= 7:
            im_steps = 1

        ps_mods = []
        for i in range(Nsrcs):
            row = src_tab[i]
            mod = Point_Source_Model_Binned_Rates(row['imx'], row['imy'], 0.1,\
                                                  [llh_obj.ebins0,llh_obj.ebins1],\
                                                  rt_obj, llh_obj.bl_dmask,\
                                                  use_deriv=True, name=row['Name'])
            ps_mods.append(mod)

        mod_list += ps_mods
        comp_mod = CompoundModel(mod_list)

        llh_obj.set_model(comp_mod)

        bf_nllh, bf_params, TS_nulls = bkg_withPS_fit(src_tab, comp_mod,\
                                      llh_obj, bkg_t0s, bkg_t1s,\
                                      test_null=True, im_steps=im_steps)

        logging.debug("TS_nulls: ")
        logging.debug(TS_nulls)

        bkg_rates = np.array([bf_params['Background'+'_bkg_rate_'+str(j)]\
                            for j in range(nebins)])
        min_rate = 1e-1*bkg_rates
        logging.debug("min_rate: ")
        logging.debug(min_rate)
        PSs2keep = []
        for name, TS in TS_nulls.items():
            if TS < 8.0:
                ps_rates = np.array([bf_params[name+'_rate_'+str(j)] for j in range(nebins)])
                logging.debug(name + " rates: ")
                logging.debug(ps_rates)
                # print ps_rates
                if np.all(ps_rates<min_rate):
                    continue
            PSs2keep.append(name)

        if len(PSs2keep) == len(src_tab):
            break
        if len(PSs2keep) == 0:
            Nsrcs = 0
            src_tab = src_tab[np.zeros(len(src_tab), dtype=np.bool)]
            break
        bl = np.array([src_tab['Name'][i] in PSs2keep for i in range(Nsrcs)])
        src_tab = src_tab[bl]
        Nsrcs = len(src_tab)
        logging.debug("src_tab: ")
        logging.debug(src_tab)

    return bf_params, src_tab

def get_info_mat_around_min(llh_obj, mod, params_, ebin):

    params = copy(params_)
    dt = llh_obj.dt
    mod_cnts = llh_obj.model.get_rate_dpi(params, ebin)*dt
    data_cnts = llh_obj.data_dpis[ebin]

    dR_dparams = mod.get_dr_dp(params, ebin)

    cov_ndim = len(dR_dparams)

    info_mat = np.zeros((cov_ndim,cov_ndim))

    for i in range(cov_ndim):
        for j in range(cov_ndim):
            info_mat[i,j] = np.sum(((dR_dparams[j]*dt)*(dR_dparams[i]*dt)*\
                                    data_cnts)/np.square(mod_cnts))
#             cov_mat[i,j] = np.sum(np.square(mod_cnts)/\
#                         ((dR_dparams[j]*dt)*(dR_dparams[i]*dt)*data_cnts))

    return info_mat

def get_errs_corrs(llh_obj, model, params, e0, pnames2skip=[]):

    imat = get_info_mat_around_min(llh_obj, model, copy(params), e0)
    cov_mat = np.linalg.inv(imat)
    e0_pnames = []
    for pname in model.param_names:
        try:
            if int(pname[-1])==e0 and not (pname in pnames2skip):
                e0_pnames.append(pname)
        except:
            pass

    err_dict = {}
    corr_dict = {}
    errs = np.sqrt(np.diag(cov_mat))
    for i, pname in enumerate(e0_pnames):
        k = 'err_' + pname
        err_dict[k] = errs[i]
    Npars = len(e0_pnames)
    for i in range(Npars-1):
        pname0 = e0_pnames[i]
        for j in range(i+1,Npars):
            pname1 = e0_pnames[j]
            k = 'corr_' + pname0 + '_' + pname1
            corr_dict[k] = cov_mat[i,j]/(errs[i]*errs[j])

    return err_dict, corr_dict

def bkg_withPS_fit_fiximxy(PS_tab, model, llh_obj, t0s, t1s, params_,\
                           fixed_pnames=None):

    Nps = len(PS_tab)
    params = copy(params_)

    llh_obj.set_model(model)
    bkg_miner = NLLH_ScipyMinimize_Wjacob('')
    bkg_miner.set_llh(llh_obj)
    llh_obj.set_time(t0s,t1s)

    nllh = 0.0
#     bf_params = {fixed_pars[i]:fixed_vals[i] for i in range(len(fixed_pars))}
    bf_params = copy(params)
    fixed_vals = [bf_params[pname] for pname in fixed_pnames]
    errs_dict = {}
    corrs_dict = {}

    for e0 in range(llh_obj.nebins):
        bkg_miner.set_fixed_params(bkg_miner.param_names)
        if fixed_pnames is not None:
            bkg_miner.set_fixed_params(fixed_pnames, values=fixed_vals)
        e0_pnames = []
        for pname in bkg_miner.param_names:
            try:
                if int(pname[-1])==e0 and not (pname in fixed_pnames):
                    e0_pnames.append(pname)
            except:
                pass
        bkg_miner.set_fixed_params(e0_pnames, fixed=False)
        llh_obj.set_ebin(e0)

        bf_vals, bkg_nllh, res = bkg_miner.minimize()
        nllh += bkg_nllh[0]
        for ii, pname in enumerate(e0_pnames):
            bf_params[pname] = bf_vals[0][ii]

        print(bf_params)

        err_dict, corr_dict = get_errs_corrs(llh_obj, model, copy(bf_params), e0, pnames2skip=fixed_pnames)
        for k, val in err_dict.items():
            errs_dict[k] = val
        for k, val in corr_dict.items():
            corrs_dict[k] = val

    return nllh, bf_params, errs_dict, corrs_dict

def main(args):

    logging.basicConfig(filename='bkg_rate_estimation_wPSs.log', level=logging.DEBUG,\
                    format='%(asctime)s-' '%(levelname)s- %(message)s')

    if args.bkg_nopost and args.bkg_nopre:
        raise Exception("Can't have no pre and no post")

    if args.dbfname is None:
        db_fname = guess_dbfname()
        if isinstance(db_fname, list):
            db_fname = db_fname[0]
    else:
        db_fname = args.dbfname

    logging.info('Connecting to DB')
    conn = get_conn(db_fname)

    info_tab = get_info_tab(conn)
    logging.info('Got info table')

    files_tab = get_files_tab(conn)
    logging.info('Got files table')

    trigtime = info_tab['trigtimeMET'][0]
    tstart = trigtime - args.twind
    tstop = trigtime + args.twind

    evfname = files_tab['evfname'][0]
    dmfname = files_tab['detmask'][0]
    attfname = files_tab['attfname'][0]
    ev_data = fits.open(evfname)[1].data
    try:
        GTI = Table.read(evfname, hdu='GTI_POINTING')
    except:
        GTI = Table.read(evfname, hdu='GTI')
    dmask = fits.open(dmfname)[0].data
    bl_dmask = (dmask==0.)
    attfile = fits.open(attfname)[1].data
    logging.debug('Opened up event, detmask, and att files')

    ebins0 = np.array(EBINS0)
    ebins1 = np.array(EBINS1)
    nebins = len(ebins0)
    logging.debug("ebins0")
    logging.debug(ebins0)
    logging.debug("ebins1")
    logging.debug(ebins1)

    solid_angle_dpi = np.load(solid_angle_dpi_fname)

    src_tab = get_srcs_infov(attfile, trigtime, pcfname=args.pcfname)
    Nsrcs = len(src_tab)
    logging.info("src_tab: ")
    logging.info(src_tab)

    bkg_mod = Bkg_Model_wFlatA(bl_dmask, solid_angle_dpi, nebins, use_deriv=True)

    llh_obj = LLH_webins(ev_data, ebins0, ebins1, bl_dmask)

    # add in stuff later for if there's no srcs

    if Nsrcs > 0:

        rt_dir = files_tab['rtDir'][0]
        rt_obj = RayTraces(rt_dir)

        sig_dtwind = (-10*1.024, 20*1.024)
        sig_twind = (trigtime + sig_dtwind[0], trigtime + sig_dtwind[1])

        init_bf_params, src_tab = do_init_bkg_wPSs(bkg_mod, llh_obj, src_tab, rt_obj, GTI, sig_twind)

        Nsrcs = len(src_tab)

        logging.debug("Final src_tab:")
        logging.debug(src_tab)

        # Now need to do each time, with these PSs and these imxys

    else:
        init_bf_params = {k:bkg_mod.param_dict[k]['val'] for k in bkg_mod.param_names}

    if Nsrcs > 0:
        fixed_pars = [pname for pname in list(init_bf_params.keys()) if '_flat_' in pname\
                    or '_imx' in pname or '_imy' in pname]

        mod_list = [bkg_mod]
        ps_mods = []
        for i in range(Nsrcs):
            row = src_tab[i]
            mod = Point_Source_Model_Binned_Rates(row['imx'], row['imy'], 0.1,\
                                                  [ebins0,ebins1], rt_obj, bl_dmask,\
                                                  use_deriv=True, name=row['Name'])
            ps_mods.append(mod)

        mod_list += ps_mods
        mod = CompoundModel(mod_list)

    else:
        init_bf_params = {k:bkg_mod.param_dict[k]['val'] for k in bkg_mod.param_names}
        mod = bkg_mod
        fixed_pars = []

    bkg_dur = args.bkg_dur*1.024
    twind = args.twind*1.024
    bkg_tstep = 2*1.024
    dt_ax = np.arange(-twind, twind+1, bkg_tstep)
    t_ax = dt_ax + trigtime
    Ntpnts = len(dt_ax)
    logging.info('Ntpnts: %d'%(Ntpnts))
    logging.info('min(dt_ax): %.3f'%(np.min(dt_ax)))
    logging.info('max(dt_ax): %.3f'%(np.max(dt_ax)))
    sig_wind = 10*1.024
    # sig_twind = (trigger_time + sig_dtwind[0], trigger_time + sig_dtwind[1])

    bkg_bf_dicts = []

    for i in range(Ntpnts):

        tmid = t_ax[i]
        sig_twind = (-sig_wind/2. + tmid, sig_wind/2. + tmid)
        gti_ = add_bti2gti(sig_twind, GTI)
        bkg_t0 = tmid - sig_wind/2. - bkg_dur/2.
        bkg_t1 = tmid + sig_wind/2. + bkg_dur/2.
        bkg_bti = Table(data=([-np.inf, bkg_t1], [bkg_t0, np.inf]), names=('START', 'STOP'))
        gti_ = add_bti2gti(bkg_bti, gti_)
        print(tmid - trigtime)
        print(gti_)
        t0s = gti_['START']
        t1s = gti_['STOP']

        nllh, params, errs_dict, corrs_dict = bkg_withPS_fit_fiximxy(src_tab, mod, llh_obj, t0s, t1s,\
                                                                     copy(init_bf_params),\
                                                                     fixed_pnames=fixed_pars)

        params.update(errs_dict)
        params.update(corrs_dict)
        params['nllh'] = nllh
        params['time'] = tmid
        params['dt'] = tmid - trigtime
        bkg_bf_dicts.append(params)
        params['exp'] = llh_obj.dt

    bkg_df = pd.DataFrame(bkg_bf_dicts)

    save_fname = 'bkg_estimation.csv'
    logging.info("Saving results in a DataFrame to file: ")
    logging.info(save_fname)
    bkg_df.to_csv(save_fname, index=False)

if __name__ == "__main__":

    args = cli()

    main(args)


# Copyright 2019 Pascal Audet
#
# This file is part of RfPy.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Functions to calculate piercing points from a velocity model and slowness
values, bin them, and produce CCP stacks using receiver functions.

"""

import os
import sys
import pickle
import numpy as np
import scipy as sp
from scipy.signal import hilbert
from rfpy import binning
import matplotlib.pyplot as plt
from matplotlib import cm

class CCPimage(object):
    """
    A CCPimage object contains attributes and methods to produce
    Common Conversion Point (CCP) stacks for each of the main
    three main phases (Ps, Pps and Pss) using radial-component
    receiver functions. The object is used to project the stacks along
    a linear profile, specified by start and end geographical coordinate
    locations, which are subsequently averaged to produce a final
    CCP image. The averaging can be done using a linear weighted sum,
    or a phase-weighted sum. Methods should be used in the appropriate
    sequence (see ``rfpy_ccp.py`` for details).

    Note
    ----
    By default, the object initializes with un-defined coordinate locations for
    the profile. If not specified during initialization, make sure
    they are specified later, before the other methods are used, e.g.
    ``ccpimage.xs_lat1 = 10.; ccpimage.xs_lon1 = 110.``, etc. Note also that the
    default 1D velocity model may not be applicable to your region of
    interest and a different model can be implemented during initialization
    or later during processing.

    Parameters
    ----------
    coord_start : list
        List of two floats corresponding to the (latitude, longitude)
        pair for the start point of the profile
    coord_end : list
        List of two floats corresponding to the (latitude, longitude)
        pair for the end point of the profile
    weights : list
        List of three floats with corresponding weights for the Ps, Pps
        and Pss phases used during linear, weighted averaging
    dep : :class:`~numpy.ndarray`
        Array of depth values defining the 1D background seismic velocity model.
        Note that the maximum depth defined here sets the maximum depth
        in each of the CCP stacks and the final CCP image.
    vp : :class:`~numpy.ndarray`
        Array of Vp values defining the 1D background seismic velocity model
    vpvs : float
        Constant Vp/Vs ratio for the 1D model.

    Other Parameters
    ----------------
    radialRF : list
        List of :class:`~obspy.core.Stream` objects containing the radial
        receiver functions along the line. Each item in the list contains the
        streams for one single station.
    vs : :class:`~numpy.ndarray`
        Array of Vp values defining the 1D background seismic velocity model
    xs_lat1 : float
        Latitude of start point defining the linear profile.
    xs_lon1 : float
        Longitude of start point defining the linear profile.
    xs_lat2 : float
        Latitude of end point defining the linear profile.
    xs_lon2 : float
        Longitude of end point defining the linear profile.
    is_ready_for_prep : boolean
        Whether or not the object is ready for the method ``prep_data``
    is_ready_for_prestack : boolean
        Whether or not the object is ready for the method ``prestack``
    is_ready_for_ccp : boolean
        Whether or not the object is ready for the method ``ccp``
    is_ready_for_gccp : boolean
        Whether or not the object is ready for the method ``gccp``

    """

    def __init__(self, coord_start=[None, None], coord_end=[None, None],
                 weights=[1., 3., -3.],
                 dep=np.array([0., 4., 8., 14., 30., 35., 45., 110.]),
                 vp=np.array([4.0, 5.9, 6.2, 6.3, 6.8, 7.2, 8.0, 8.1]),
                 vs=None, vpvs=1.73, dx=2.5, dz=1.):

        self.radialRF = []
        self.dep = dep
        self.weights = weights
        self.xs_lat1 = coord_start[0]
        self.xs_lon1 = coord_start[1]
        self.xs_lat2 = coord_end[0]
        self.xs_lon2 = coord_end[1]
        self.is_ready_for_prep = False
        self.is_ready_for_prestack = False
        self.is_ready_for_ccp = False
        self.is_ready_for_gccp = False

        # Define grid parameters
        self.dz = dz
        self.dx = dx

        # Get total length of grid from end points
        xlength = haversine(self.xs_lat1, self.xs_lon1,
                              self.xs_lat2, self.xs_lon2)

        # number of cells laterally and vertically
        self.nx = int(np.rint(xlength/self.dx))
        self.nz = int(self.dep[-1]/self.dz)

        self.xarray = np.arange(self.nx)*self.dx
        self.zarray = np.arange(self.nz)*self.dz

        # Interpolate Vp and Vs models on depth grid
        if vs is None:
            vs = vp/vpvs
        else:
            if not vp.shape == vs.shape:
                raise(Exception("vp and vs arrays have a different shape"))

        self.vp = sp.interpolate.interp1d(dep, vp, kind='linear')(self.zarray)
        self.vs = sp.interpolate.interp1d(dep, vs, kind='linear')(self.zarray)

    def add_rfstream(self, rfstream):
        """
        Method to add a :class:`~obspy.core.Stream` object to the list
        ``radialRF``. With at least one stream in ``radialRF``, the object
        is ready for the ``prep_data`` method, and the corresponding flag is
        updated.

        Parameters
        ----------
        rfstream : :class:`~obspy.core.Stream`
            Stream object containing radial receiver functions for one station

        """
        if len(rfstream) > 0:
            # fftshift if the time axis starts at negative lags
            if rfstream[0].stats.taxis[0]<0.:
                for tr in rfstream:
                    tr.data = np.fft.fftshift(tr.data)

            self.radialRF.append(rfstream)
            self.is_ready_for_prep = True

    def prep_data(self, f1=0.05, f2ps=0.5, f2pps=0.25, f2pss=0.2,
                  nbaz=36+1, nslow=40+1):
        """
        Method to pre-process the data and calculate the CCP points for each
        of the receiver functions. Pre-processing includes the binning to
        back-azimuth and slowness bins to reduce processing time and generate
        cleaner stacks, as well as filtering to emphasize the energy of the
        various phases as a function of depth. As a general rule, the high
        frequency corner of the 2 reverberated phases (Pps and Pss) should be
        approximately half of the high frequency corner of the direct
        (Ps) phase. At the end of this step, the object is updated with
        the amplitude at the lat, lon and depth values corresponding to the
        raypath for each receiver function. The object is now ready for the
        method ``prestack``, with the corresponding flag updated.

        Parameters
        ----------
        f1 : float
            Low-frequency corner of the bandpass filter used for all phases (Hz)
        f2ps : float
            High-frequency corner of the bandpass filter used for the Ps phase (Hz)
        f2pps : float
            High-frequency corner of the bandpass filter used for the Pps phase (Hz)
        f2pss : float
            High-frequency corner of the bandpass filter used for the Pss phase (Hz)
        nbaz : int
            Number of increments in the back-azimuth bins
        nslow : int
            Number of increments in the slowness bins

        The following attributes are added to the object:

        Other Parameters
        ----------------
        amp_ps_depth : :class:`numpy.ndarray`
            2D array of amplitudes as a function of depth for the Ps phase
        amp_pps_depth : :class:`numpy.ndarray`
            2D array of amplitudes as a function of depth for the Pps phase
        amp_pss_depth : :class:`numpy.ndarray`
            2D array of amplitudes as a function of depth for the Pss phase
        lon_depth : :class:`numpy.ndarray`
            2D array of longitude as a function of depth (i.e., piercing points)
        lat_depth : :class:`numpy.ndarray`
            2D array of latitude as a function of depth (i.e., piercing points)
        is_ready_for_presstack : boolean
            Flag specifying that the object is ready for the presstack() method
        n_traces : float
            Total number of traces used in creating the CCP image.

        """

        if not self.is_ready_for_prep:
            raise(Exception("CCPimage not ready for pre-prep"))

        ikey = 0
        total_traces = 0

        # Process streams one at a time
        for RF in self.radialRF:

            # Bin RFs into back-azimuth and slowness bins to speed up
            # calculations
            RFbin = binning.bin_baz_slow(
                RF, nbaz=nbaz, nslow=nslow)[0]
            n_traces = len(RFbin)
            total_traces += n_traces

            amp_ps_tr = np.empty([n_traces, self.nz])
            amp_pps_tr = np.empty([n_traces, self.nz])
            amp_pss_tr = np.empty([n_traces, self.nz])
            lon_tr = np.empty([n_traces, self.nz])
            lat_tr = np.empty([n_traces, self.nz])

            st_ps = RFbin.copy()
            st_pps = RFbin.copy()
            st_pss = RFbin.copy()

            # Filter Ps, Pps and Pss
            st_ps.filter(
                'bandpass', freqmin=f1, freqmax=f2ps,
                corners=4, zerophase=True)
            st_pps.filter(
                'bandpass', freqmin=f1, freqmax=f2pps,
                corners=4, zerophase=True)
            st_pss.filter(
                'bandpass', freqmin=f1, freqmax=f2pss,
                corners=4, zerophase=True)
            del RFbin

            print("Station: "+st_ps[0].stats.station)
            for itr in _progressbar(range(len(st_ps)), '', 25):

                # Get raypath and travel time for all phases
                tt_ps, tt_pps, tt_pss, plon, plat = \
                    raypath(st_ps[itr], dep=self.zarray, vp=self.vp, vs=self.vs)

                # Now get amplitude of RF at corresponding travel
                # time along the raypath
                lon_tr[itr, :] = plon
                lat_tr[itr, :] = plat

                amp_ps = []
                amp_pps = []
                amp_pss = []

                # Loop through travel times and shift RFs to get amplitudes
                for tt in tt_ps:
                    a, phase = timeshift(st_ps[itr], tt)
                    amp_ps.append(a)
                amp_ps_tr[itr, :] = amp_ps

                # Loop through travel times and shift RFs to get amplitudes
                for tt in tt_pps:
                    a, phase = timeshift(st_pps[itr], tt)
                    amp_pps.append(a)
                amp_pps_tr[itr, :] = amp_pps

                # Loop through travel times and shift RFs to get amplitudes
                for tt in tt_pss:
                    a, phase = timeshift(st_pss[itr], tt)
                    amp_pss.append(a)
                amp_pss_tr[itr, :] = amp_pss

            if ikey == 0:
                amp_ps_depth = amp_ps_tr.transpose()
                amp_pps_depth = amp_pps_tr.transpose()
                amp_pss_depth = amp_pss_tr.transpose()
                lon_depth = lon_tr.transpose()
                lat_depth = lat_tr.transpose()

            elif ikey > 0:
                amp_ps_depth = np.concatenate(
                    (amp_ps_depth, amp_ps_tr.transpose()), axis=1)
                amp_pps_depth = np.concatenate(
                    (amp_pps_depth, amp_pps_tr.transpose()), axis=1)
                amp_pss_depth = np.concatenate(
                    (amp_pss_depth, amp_pss_tr.transpose()), axis=1)
                lon_depth = np.concatenate(
                    (lon_depth, lon_tr.transpose()), axis=1)
                lat_depth = np.concatenate(
                    (lat_depth, lat_tr.transpose()), axis=1)

            ikey += 1

        self.amp_ps_depth = amp_ps_depth
        self.amp_pps_depth = amp_pps_depth
        self.amp_pss_depth = amp_pss_depth
        self.lon_depth = lon_depth
        self.lat_depth = lat_depth
        self.is_ready_for_prestack = True
        self.n_traces = total_traces

        del self.radialRF

    def prestack(self):
        """
        Method to project the raypaths onto the 2D profile for each of the three
        phases. The final grid is defined here, using the parameter ``dx`` in km.
        The horizontal extent is pre-determined from the start and end points of
        the profile. At the end of this step, the object contains the set of
        amplitudes at each of the 2D grid points, for each of the three phases.
        The object is now ready for the methods ``ccp`` and/or ``gccp``, with
        the corresponding flag updated.

        The following attributes are added to the object:

        Other Parameters
        ----------------
        xs_amps_ps : :class:`numpy.ndarray`
            3D array of amplitudes (1D array of amplitudes at each grid cell)
            for the Ps phase
        xs_amps_pps : :class:`numpy.ndarray`
            3D array of amplitudes (1D array of amplitudes at each grid cell)
            for the Pps phase
        xs_amps_pss : :class:`numpy.ndarray`
            3D array of amplitudes (1D array of amplitudes at each grid cell)
            for the Pss phase
        is_ready_for_ccp : boolean
            Flag specifying that the object is ready for the ccp() method
        is_ready_for_gccp : boolean
            Flag specifying that the object is ready for the gccp() method

        """

        if not self.is_ready_for_prestack:
            raise(Exception("CCPimage not ready for prestack"))

        xs_latitudes = np.asarray(
            np.linspace(self.xs_lat1, self.xs_lat2, self.nx))
        xs_longitudes = np.asarray(
            np.linspace(self.xs_lon1, self.xs_lon2, self.nx))

        xs_amps_ps = np.zeros((self.nz, self.nx, self.n_traces))
        xs_amps_pps = np.zeros((self.nz, self.nx, self.n_traces))
        xs_amps_pss = np.zeros((self.nz, self.nx, self.n_traces))

        for iz in _progressbar(range(self.nz), '', 25):

            for i_coor in range(self.n_traces):

                lat_tr = self.lat_depth[iz, i_coor]
                lon_tr = self.lon_depth[iz, i_coor]
                distance_tests = np.empty(self.nx)

                for i_xs in range(self.nx):
                    lat_xs = xs_latitudes[i_xs]
                    lon_xs = xs_longitudes[i_xs]
                    distance_tests[i_xs] = haversine(
                        lat_xs, lon_xs, lat_tr, lon_tr)

                minimum_distance = np.amin(distance_tests)
                ix = np.where(distance_tests ==
                                  np.amin(distance_tests))[0][0]

                nonzero_count = np.count_nonzero(
                    xs_amps_ps[iz, ix, :])
                new_amp_ps = self.amp_ps_depth[iz, i_coor]
                if xs_amps_ps[iz, ix, 0] == 0.:
                    xs_amps_ps[iz, ix, 0] = new_amp_ps
                else:
                    xs_amps_ps[iz, ix, nonzero_count] = new_amp_ps

                nonzero_count = np.count_nonzero(
                    xs_amps_pps[iz, ix, :])
                new_amp_pps = self.amp_pps_depth[iz, i_coor]
                if xs_amps_pps[iz, ix, 0] == 0.:
                    xs_amps_pps[iz, ix, 0] = new_amp_pps
                else:
                    xs_amps_pps[iz, ix, nonzero_count] = new_amp_pps

                nonzero_count = np.count_nonzero(
                    xs_amps_pss[iz, ix, :])
                new_amp_pss = self.amp_pss_depth[iz, i_coor]
                if xs_amps_pss[iz, ix, 0] == 0.:
                    xs_amps_pss[iz, ix, 0] = new_amp_pss
                else:
                    xs_amps_pss[iz, ix, nonzero_count] = new_amp_pss

        self.xs_amps_ps = xs_amps_ps
        self.xs_amps_pps = xs_amps_pps
        self.xs_amps_pss = xs_amps_pss
        self.is_ready_for_ccp = True
        self.is_ready_for_gccp = True

        del self.amp_ps_depth
        del self.amp_pps_depth
        del self.amp_pss_depth
        del self.lon_depth
        del self.lat_depth

    def ccp(self):
        """
        Method to average the amplitudes at each grid point to produce 2D images
        for each of the three phases. At the end of this step, the object
        contains the three 2D arrays that can be further averaged into a single
        final image.

        The following attributes are added to the object:

        Other Parameters
        ----------------
        xs_ps_avg : :class:`numpy.ndarray`
            2D array of stacked amplitudes for the Ps phase
        xs_pps_avg : :class:`numpy.ndarray`
            2D array of stacked amplitudes for the Pps phase
        xs_pss_avg : :class:`numpy.ndarray`
            2D array of stacked amplitudes for the Pss phase

        """

        if not self.is_ready_for_ccp:
            raise(Exception("CCPimage not ready for ccp"))

        xs_ps_avg = np.zeros((self.nz, self.nx))
        xs_pps_avg = np.zeros((self.nz, self.nx))
        xs_pss_avg = np.zeros((self.nz, self.nx))

        for iz in _progressbar(range(self.nz), '', 25):

            for ix in range(self.nx):

                nonzero_count = np.count_nonzero(self.xs_amps_ps[iz, ix, :])
                if nonzero_count != 0:
                    amps_ps = self.xs_amps_ps[iz, ix, 0:nonzero_count]
                    xs_ps_avg[iz, ix] = np.mean(amps_ps)

                nonzero_count = np.count_nonzero(self.xs_amps_pps[iz, ix, :])
                if nonzero_count != 0:
                    amps_pps = self.xs_amps_pps[iz, ix, 0:nonzero_count]
                    xs_pps_avg[iz, ix] = np.mean(amps_pps)

                nonzero_count = np.count_nonzero(self.xs_amps_pss[iz, ix, :])
                if nonzero_count != 0:
                    amps_pss = self.xs_amps_pss[iz, ix, 0:nonzero_count]
                    xs_pss_avg[iz, ix] = np.mean(amps_pss)

        self.xs_ps_avg = xs_ps_avg
        self.xs_pps_avg = xs_pps_avg
        self.xs_pss_avg = xs_pss_avg

        del self.xs_amps_ps
        del self.xs_amps_pps
        del self.xs_amps_pss

    def gccp(self, wlen=15.):
        """
        Method to average the amplitudes at each grid point to produce 2D images
        for each of the three phases. In this method, the grid points are further
        smoothed in the horizontal direction using a Gaussian function to simulate
        P-wave sensitivity kernels. At the end of this step, the object
        contains the three 2D smoothed arrays that can be further averaged into a
        single final image.

        Parameters
        ----------
        wlen : float
            Wavelength of the P-wave for smoothing (km).

        The following attributes are added to the object:

        Other Parameters
        ----------------
        xs_gauss_ps : :class:`numpy.ndarray`
            2D array of stacked and Gaussian-filtered amplitudes for the Ps phase
        xs_gauss_pps : :class:`numpy.ndarray`
            2D array of stacked and Gaussian-filtered amplitudes for the Pps phase
        xs_gauss_pss : :class:`numpy.ndarray`
            2D array of stacked and Gaussian-filtered amplitudes for the Pss phase

        """

        if not self.is_ready_for_gccp:
            raise(Exception("CCPimage not ready for gccp"))
        if not hasattr(self, 'xs_ps_avg'):
            self.ccp()

        dlat = max(self.xarray)/self.nx

        import scipy.ndimage as ndimage

        self.xs_gauss_ps = ndimage.filters.gaussian_filter(
            self.xs_ps_avg, sigma=(0, wlen/dlat))
        self.xs_gauss_pps = ndimage.filters.gaussian_filter(
            self.xs_pps_avg, sigma=(0, wlen/dlat))
        self.xs_gauss_pss = ndimage.filters.gaussian_filter(
            self.xs_pss_avg, sigma=(0, wlen/dlat))

    def linear_stack(self, typ='ccp'):
        """
        Method to average the three 2D images into a final, weighted CCP image
        using the weights defined in the attribute.

        Parameters
        ----------
        typ : str
            Type of phase stacks to use (either `ccp` or `gccp`)

        The following attributes are added to the object:

        Other Parameters
        ----------------
        tot_trace : :class:`numpy.ndarray`
            2D array of amplitudes for the linearly combined Ps, Pps and Pss phases

        """

        tot_trace = np.zeros((self.nz, self.nx))

        if typ=='ccp':
            if not hasattr(self, "xs_ps_avg"):
                self.ccp()
            xs_ps = self.xs_ps_avg*self.weights[0]
            xs_pps = self.xs_pps_avg*self.weights[1]
            xs_pss = self.xs_pss_avg*self.weights[2]
        elif typ=='gccp':
            if not hasattr(self, 'xs_gauss_ps'):
                self.gccp()
            xs_ps = self.xs_gauss_ps*self.weights[0]
            xs_pps = self.xs_gauss_pps*self.weights[1]
            xs_pss = self.xs_gauss_pss*self.weights[2]

        for ix in range(self.nx):
            ps_trace = xs_ps[:, ix]
            pps_trace = xs_pps[:, ix]
            pss_trace = xs_pss[:, ix]
            tot_trace[:, ix] = (ps_trace + pps_trace + pss_trace)

        self.tot_trace = tot_trace

    def phase_weighted_stack(self, typ='gccp'):
        """
        Method to average the three 2D smoothed images into a final,
        phase-weighted CCP image.

        Parameters
        ----------
        typ : str
            Type of phase stacks to use (either `ccp` or `gccp`)

        The following attributes are added to the object:

        Other Parameters
        ----------------
        tot_trace : :class:`numpy.ndarray`
            2D array of amplitudes for the phase-weighted, combined
            Ps, Pps and Pss phases

        """

        tot_trace = np.zeros((self.nz, self.nx))

        if typ=='ccp':
            if not hasattr(self, "xs_ps_avg"):
                self.ccp()
            xs_ps = self.xs_ps_avg*self.weights[0]
            xs_pps = self.xs_pps_avg*self.weights[1]
            xs_pss = self.xs_pss_avg*self.weights[2]
        elif typ=='gccp':
            if not hasattr(self, 'xs_gauss_ps'):
                self.gccp()
            xs_ps = self.xs_gauss_ps*self.weights[0]
            xs_pps = self.xs_gauss_pps*self.weights[1]
            xs_pss = self.xs_gauss_pss*self.weights[2]

        for ix in range(self.nx):
            ps_trace = xs_ps[:, ix]
            pps_trace = xs_pps[:, ix]
            pss_trace = xs_pss[:, ix]

            weight = np.zeros(len(ps_trace), dtype=complex)
            ps_hilb = hilbert(ps_trace)
            ps_phase = np.arctan2(ps_hilb.imag, ps_hilb.real)
            weight += np.exp(1j*ps_phase)

            pps_hilb = hilbert(pps_trace)
            pps_phase = np.arctan2(pps_hilb.imag, pps_hilb.real)
            weight += np.exp(1j*pps_phase)

            pss_hilb = hilbert(pss_trace)
            pss_phase = np.arctan2(pss_hilb.imag, pss_hilb.real)
            weight += np.exp(1j*pss_phase)

            weight = np.abs(weight)/3.

            tot_trace[:, ix] = (ps_trace + pps_trace + pss_trace)*weight**2

        self.tot_trace = tot_trace

    def save(self, title):
        """
        Method to save the `ccpimage` object to file.

        Parameters
        ----------
        title : str
            File name for the saved object

        """

        if title is None:
            title = "CCP_image.pkl"

        if not ".pkl" in title:
            file = open(title+".pkl", "wb")
        else:
            file = open(title, "wb")
        pickle.dump(self, file)
        file.close()

    def plot_ccp(self, vmin=-0.05, vmax=0.05,
        save=False, title='', fmt='png'):
        """
        Method to plot the final CCP stacks along the line.

        Parameters
        ----------
        vmin : float
            Maximum negative value for the color scale
        vmax : float
            Maximum positive value for the color scale
        save : boolean
            Whether or not to save the figure.
        fmt : str
            Format of saved figure
        xlen : float
            Total length of profile

        """

        xlen = self.nx*self.dx
        xm = 1.5 + 0.0065*xlen
        fig, (ax1, ax2, ax3, ax4) = plt.subplots(
            4, 1, figsize=(xm, 8))

        # plt.pcolormesh(xarray,zarray,xs_ps_avg,cmap=cm.coolwarm,vmin=vmin,vmax=vmax)
        im1 = ax1.pcolormesh(self.xarray, self.zarray,
                             self.xs_ps_avg*self.weights[0],
                             cmap=cm.RdBu_r,
                             vmin=vmin, vmax=vmax)
        bar = plt.colorbar(im1, ax=ax1)
        ax1.set_xlim((min(self.xarray)),
                     (max(self.xarray)))
        ax1.set_ylim((min(self.zarray)), (max(self.zarray)))
        bar.ax.set_ylabel('Amplitude', size=10)
        # ax1.set_xlabel('Lateral Distance (km)', size=10)
        ax1.set_ylabel('Depth (km)', size=10)
        ax1.set_title('Ps CCP image', size=10)
        ax1.invert_yaxis()

        # plt.pcolormesh(xarray,zarray,xs_pps_avg,cmap=cm.coolwarm,vmin=vmin,vmax=vmax)
        im2 = ax2.pcolormesh(self.xarray, self.zarray,
                             self.xs_pps_avg*self.weights[1],
                             cmap=cm.RdBu_r,
                             vmin=vmin, vmax=vmax)
        bar = plt.colorbar(im2, ax=ax2)
        ax2.set_xlim((min(self.xarray)),
                     (max(self.xarray)))
        ax2.set_ylim((min(self.zarray)), (max(self.zarray)))
        bar.ax.set_ylabel('Amplitude', size=10)
        # ax2.set_xlabel('Lateral Distance (km)', size=10)
        ax2.set_ylabel('Depth (km)', size=10)
        ax2.set_title('Pps CCP image', size=10)
        ax2.invert_yaxis()

        im3 = ax3.pcolormesh(self.xarray, self.zarray,
                             self.xs_pss_avg*self.weights[2],
                             cmap=cm.RdBu_r,
                             vmin=vmin, vmax=vmax)
        bar = plt.colorbar(im3, ax=ax3)
        ax3.set_xlim((min(self.xarray)),
                     (max(self.xarray)))
        ax3.set_ylim((min(self.zarray)), (max(self.zarray)))
        bar.ax.set_ylabel('Amplitude', size=10)
        # ax3.set_xlabel('Lateral Distance (km)', size=10)
        ax3.set_ylabel('Depth (km)', size=10)
        ax3.set_title('Pss CCP image', size=10)
        ax3.invert_yaxis()

        im4 = ax4.pcolormesh(self.xarray, self.zarray,
                             self.tot_trace, cmap=cm.RdBu_r,
                             vmin=vmin, vmax=vmax)
        bar = plt.colorbar(im4, ax=ax4)
        ax4.set_xlim((min(self.xarray)),
                     (max(self.xarray)))
        ax4.set_ylim((min(self.zarray)), (max(self.zarray)))
        bar.ax.set_ylabel('Amplitude', size=10)
        ax4.set_xlabel('Lateral Distance (km)', size=10)
        ax4.set_ylabel('Depth (km)', size=10)
        ax4.set_title('Weighted CCP image', size=10)
        ax4.invert_yaxis()
        plt.tight_layout()

        if save:
            if not os.path.isdir("FIGURES"):
                os.makedirs("FIGURES")
            plt.savefig('FIGURES/ccp.' + title + '.' + fmt)

        plt.show()

    def plot_gccp(self, vmin=-0.015, vmax=0.015,
        save=False, title='', fmt='png'):
        """
        Method to plot the final GCCP stacks along the line.

        Parameters
        ----------
        vmin : float
            Maximum negative value for the color scale
        vmax : float
            Maximum positive value for the color scale
        save : boolean
            Whether or not to save the figure.
        fmt : str
            Format of saved figure
        xlen : float
            Total length of profile

        """

        xlen = self.nx*self.dx
        xm = 1.5 + 0.0065*xlen
        fig, (ax1, ax2, ax3, ax4) = plt.subplots(
            4, 1, figsize=(xm, 8))

        # plt.pcolormesh(xarray,zarray,xs_ps_avg,cmap=cm.coolwarm,vmin=vmin,vmax=vmax)
        im1 = ax1.pcolormesh(self.xarray, self.zarray,
                             self.xs_gauss_ps*self.weights[0],
                             cmap=cm.RdBu_r,
                             vmin=vmin, vmax=vmax)
        bar = plt.colorbar(im1, ax=ax1)
        ax1.set_xlim((min(self.xarray)),
                     (max(self.xarray)))
        ax1.set_ylim((min(self.zarray)), (max(self.zarray)))
        bar.ax.set_ylabel('Amplitude', size=10)
        # ax1.set_xlabel('Lateral Distance (km)', size=10)
        ax1.set_ylabel('Depth (km)', size=10)
        ax1.set_title('Ps GCCP image', size=10)
        ax1.invert_yaxis()

        # plt.pcolormesh(xarray,zarray,xs_pps_avg,cmap=cm.coolwarm,vmin=vmin,vmax=vmax)
        im2 = ax2.pcolormesh(self.xarray, self.zarray,
                             self.xs_gauss_pps*self.weights[1],
                             cmap=cm.RdBu_r,
                             vmin=vmin, vmax=vmax)
        bar = plt.colorbar(im2, ax=ax2)
        ax2.set_xlim((min(self.xarray)),
                     (max(self.xarray)))
        ax2.set_ylim((min(self.zarray)), (max(self.zarray)))
        bar.ax.set_ylabel('Amplitude', size=10)
        # ax2.set_xlabel('Lateral Distance (km)', size=10)
        ax2.set_ylabel('Depth (km)', size=10)
        ax2.set_title('Pps GCCP image', size=10)
        ax2.invert_yaxis()

        im3 = ax3.pcolormesh(self.xarray, self.zarray,
                             self.xs_gauss_pss*self.weights[2],
                             cmap=cm.RdBu_r,
                             vmin=vmin, vmax=vmax)
        bar = plt.colorbar(im3, ax=ax3)
        ax3.set_xlim((min(self.xarray)),
                     (max(self.xarray)))
        ax3.set_ylim((min(self.zarray)), (max(self.zarray)))
        bar.ax.set_ylabel('Amplitude', size=10)
        # ax3.set_xlabel('Lateral Distance (km)', size=10)
        ax3.set_ylabel('Depth (km)', size=10)
        ax3.set_title('Pss GCCP image', size=10)
        ax3.invert_yaxis()

        import scipy.ndimage as ndimage
        self.tot_trace = ndimage.filters.gaussian_filter(
            self.tot_trace, sigma=(3, 0), order=0)

        im4 = ax4.pcolormesh(self.xarray, self.zarray,
                             self.tot_trace, cmap=cm.RdBu_r,
                             vmin=vmin, vmax=vmax)
        bar = plt.colorbar(im4, ax=ax4)
        ax4.set_xlim((min(self.xarray)),
                     (max(self.xarray)))
        ax4.set_ylim((min(self.zarray)), (max(self.zarray)))
        bar.ax.set_ylabel('Amplitude', size=10)
        ax4.set_xlabel('Lateral Distance (km)', size=10)
        ax4.set_ylabel('Depth (km)', size=10)
        ax4.set_title('Phase-weighted GCCP image', size=10)
        ax4.invert_yaxis()
        plt.tight_layout()

        if save:
            if not os.path.isdir("FIGURES"):
                os.makedirs("FIGURES")
            plt.savefig('FIGURES/gccp.' + title + '.' + fmt)

        plt.show()

def ppoint_distance(tr, delta_z, vs):
    """
    Calculate horizontal distance for interval delta_z and velocity vs

    Parameters
    ----------
    tr : :class:`~obspy.core.Trace`
        Single trace object to migrate to depth
    delta_z : float
        Vertical sampling distance
    vs : float
        S-wave velocity (km/s)

    """

    # Calculate distance
    delta_x = delta_z*np.tan(np.arcsin(tr.stats.slow*vs))

    return delta_x

def ppoint(tr, dist):
    """
    Determine geographic location of piercing point

    Parameters
    ----------
    tr : :class:`~obspy.core.Trace`
        Single trace object to migrate to depth
    dist : float
        Horizontal istance from the station (km)

    """

    # Conversion factors
    lat2km = 111.
    lon2km = 90.

    # Get lat and lon of station location
    slat = tr.stats.stla
    slon = tr.stats.stlo

    # Back-azimuth of event
    baz = tr.stats.baz*np.pi/180.

    # location of piercing point on geographical grid
    plat = dist*np.sin(-baz+np.pi/2.)/lat2km + slat
    plon = dist*np.cos(-baz+np.pi/2.)/lon2km + slon

    return plon, plat

def ttime(tr, delta_z, vp, vs, phase=None):
    """
    Calculate travel time for interval delta_z and velocities vp and vs

    Parameters
    ----------
    tr : :class:`~obspy.core.Trace`
        Single trace object to migrate to depth
    delta_z : float
        Vertical sampling distance (km)
    vp : float
        P-wave velocity
    vs : float
        S-wave velocity
    phase : str
        Phase of interest
    """

    # Get horizontal slowness
    slow = tr.stats.slow

    # Calculate travel time for phase
    if phase == 'Ps':
        tt = delta_z*(np.sqrt((1./vs)**2 - slow**2) -
                 np.sqrt((1./vp)**2 - slow**2))
    elif phase == 'Pps':
        tt = delta_z*(np.sqrt((1./vs)**2 - slow**2) +
                 np.sqrt((1./vp)**2 - slow**2))
    elif phase == 'Pss':
        tt = 2.*delta_z*(np.sqrt((1./vs)**2 - slow**2))
    else:
        print('Error - unrecognized phase, ', phase)
        print('Returning tt = 0')
        tt = 0.

    return tt

def timeshift(tr, tt):
    """
    Shift a trace by a travel time tt and take amplitude at zero

    Parameters
    ----------
    tr : :class:`~obspy.core.Trace`
        Single trace object to migrate to depth
    tt : float
        Travel time (sec)

    """

    # Define frequencies
    nt = int(tr.stats.npts)
    dt = tr.stats.delta
    freq = np.fft.fftfreq(int(nt), d=dt)

    # Hilbert transform and instantaneous phase
    hilb = hilbert(tr.data)
    hilb_index = np.rint(tt/dt)
    hilb_tt = hilb[int(hilb_index)]
    hilb_tt_phase = np.arctan2(hilb_tt.imag, hilb_tt.real)

    # Fourier transform
    ftr = np.fft.fft(tr.data)

    # Shift using Fourier transform
    for i in range(len(freq)):
        # Fourier timeshift theorem
        ftr[i] = ftr[i]*np.exp(2.*np.pi*1j*freq[i]*tt)

    # Back to time domain (inverse Fourier transform)
    rtr = np.fft.ifft(ftr)

    # Take first sample from trace (convert to real value)
    amp = np.real(rtr[0])

    return amp, hilb_tt_phase

def raypath(tr, dep=None, vp=None, vs=None):
    """
    Calculate travel times through velocity model for all phases of interest

    Parameters
    ----------
    tr : :class:`~obspy.core.Trace`
        Single trace object to migrate to depth
    nz : int
        Number of layers in interpolation
    dep : :class:`~numpy.ndarray`
        Depth array for velocity model
    vp : :class:`~numpy.ndarray`
        P-wave velocity array for velocity model
    vs : :class:`~numpy.ndarray`
        S-wave velocity array for velocity model

    """

    # Get exact depth parameters
    delta_z = dep[1] - dep[0]
    nz = len(dep)

    # Define arrays with zeros
    plat = np.zeros(nz)
    plon = np.zeros(nz)
    ttps = np.zeros(nz)
    ttpps = np.zeros(nz)
    ttpss = np.zeros(nz)

    # Now loop through all depths
    for iz in range(nz):

        # Initialize travel time and distance counters
        dtps = 0.
        dtpps = 0.
        dtpss = 0.
        delta_x = 0.

        # Sum over depths from 0 to iz
        for i in range(iz):
            dtps += ttime(tr, delta_z, vp[i], vs[i], 'Ps')
            dtpps += ttime(tr, delta_z, vp[i], vs[i], 'Pps')
            dtpss += ttime(tr, delta_z, vp[i], vs[i], 'Pss')
            delta_x += ppoint_distance(tr, delta_z, vs[i])

        # Get piercing point from distance
        plo, pla = ppoint(tr, delta_x)

        # Assign values to arrays
        ttps[iz] = dtps
        ttpps[iz] = dtpps
        ttpss[iz] = dtpss
        plon[iz] = plo
        plat[iz] = pla

    return ttps, ttpps, ttpss, plon, plat

def haversine(lat, lon, xs_lat, xs_lon):  # great-circle distance (kilometres)

    earth_radius = 6371.  # kilometres

    lat = np.radians(lat)
    lon = np.radians(lon)
    xs_lat = np.radians(xs_lat)
    xs_lon = np.radians(xs_lon)
    dlat = lat - xs_lat
    dlon = lon - xs_lon
    a = ((np.sin(dlat/2.))**2.) + \
        (np.cos(xs_lat)*np.cos(lat)*((np.sin(dlon/2.))**2.))
    distance = np.abs(2.*earth_radius*np.arcsin(np.sqrt(a)), dtype=float)

    return np.abs(distance)

def _progressbar(it, prefix="", size=60, file=sys.stdout):
    """
    Show progress bar while looping in for loop

    """

    count = len(it)

    def show(j):
        x = int(size*j/count)
        file.write("%s[%s%s] %i/%i\r" %
                   (prefix, "#"*x, "."*(size-x), j, count))
        file.flush()
    show(0)
    for i, item in enumerate(it):
        yield item
        show(i+1)
    file.write("\n")
    file.flush()