# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""Functions for performing aperture photometry on 2-D arrays."""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import math
import abc
import numpy as np
import copy
import warnings
import astropy.units as u
from astropy.io import fits
from astropy.table import Table
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
from astropy.extern import six
from astropy.utils.misc import InheritDocstrings
from astropy.utils.exceptions import AstropyUserWarning
from astropy.wcs.utils import skycoord_to_pixel
from astropy.nddata import support_nddata
from .aperture_funcs import (do_circular_photometry, do_elliptical_photometry,
                             do_rectangular_photometry)
from .utils.wcs_helpers import (skycoord_to_pixel_scale_angle, assert_angle,
                                assert_angle_or_pixel)

skycoord_to_pixel_mode = 'all'

__all__ = ['Aperture', 'SkyAperture', 'PixelAperture',
           'SkyCircularAperture', 'CircularAperture',
           'SkyCircularAnnulus', 'CircularAnnulus',
           'SkyEllipticalAperture', 'EllipticalAperture',
           'SkyEllipticalAnnulus', 'EllipticalAnnulus',
           'RectangularAperture', 'RectangularAnnulus',
           'SkyRectangularAperture', 'SkyRectangularAnnulus',
           'aperture_photometry']

def _sanitize_pixel_positions(positions):

    if isinstance(positions, u.Quantity):
        if positions.unit is u.pixel:
            positions = np.atleast_2d(positions.value)
        else:
            raise u.UnitsError("positions should be in pixel units")
    elif isinstance(positions, (list, tuple, np.ndarray)):
        positions = np.atleast_2d(positions)
        if positions.shape[1] != 2:
            if positions.shape[0] == 2:
                positions = np.transpose(positions)
            else:
                raise TypeError("List or array of (x, y) pixel coordinates "
                                "is expected got '{0}'.".format(positions))
    elif isinstance(positions, zip):
        # This is needed for zip to work seamlessly in Python 3
        positions = np.atleast_2d(list(positions))
    else:
        raise TypeError("List or array of (x, y) pixel coordinates "
                        "is expected got '{0}'.".format(positions))

    if positions.ndim > 2:
        raise ValueError('{0}-d position array not supported. Only 2-d '
                         'arrays supported.'.format(positions.ndim))

    return positions

def _make_annulus_path(patch_inner, patch_outer):
    import matplotlib.path as mpath
    verts_inner = patch_inner.get_verts()
    verts_outer = patch_outer.get_verts()
    codes_inner = (np.ones(len(verts_inner), dtype=mpath.Path.code_type) *
                   mpath.Path.LINETO)
    codes_inner[0] = mpath.Path.MOVETO
    codes_outer = (np.ones(len(verts_outer), dtype=mpath.Path.code_type) *
                   mpath.Path.LINETO)
    codes_outer[0] = mpath.Path.MOVETO
    codes = np.concatenate((codes_inner, codes_outer))
    verts = np.concatenate((verts_inner, verts_outer[::-1]))
    return mpath.Path(verts, codes)

class _ABCMetaAndInheritDocstrings(InheritDocstrings, abc.ABCMeta):
    pass

@six.add_metaclass(_ABCMetaAndInheritDocstrings)
class Aperture(object):
    """
    Abstract base class for all apertures.
    """

class SkyAperture(Aperture):
    """
    Abstract base class for 2-d apertures defined in celestial coordinates.
    """

class PixelAperture(Aperture):
    """
    Abstract base class for 2-d apertures defined in pixel coordinates.

    Derived classes should contain whatever internal data is needed to
    define the aperture, and provide the method `do_photometry` (and
    optionally, ``area``).
    """

    @abc.abstractmethod
    def plot(self, ax=None, fill=False, **kwargs):
        """
        Plot the aperture(s) on a matplotlib Axes instance.

        Parameters
        ----------
        ax : `matplotlib.axes.Axes` instance, optional
            If `None`, then the current ``Axes`` instance is used.

        fill : bool, optional
            Set whether to fill the aperture patch.  The default is
            `False`.

        kwargs
            Any keyword arguments accepted by `matplotlib.patches.Patch`.
        """

    @abc.abstractmethod
    def do_photometry(self, data, error=None, effective_gain=None,
                      pixelwise_error=True, method='exact', subpixels=5):
        """Sum flux within aperture(s).

        Parameters
        ----------
        data : array_like
            The 2-d array on which to perform photometry.
        error : array_like, optional
            Error in each pixel, interpreted as Gaussian 1-sigma uncertainty.
            ``error`` has to have the same shape as ``data``.
        effective_gain : array_like, optional
            Ratio of counts (e.g., electrons or photons) to units of the
            data (e.g., ADU), for the purpose of calculating Poisson
            error from the object itself. ``effective_gain`` must have
            the same shape as ``data``.
        pixelwise_error : bool, optional
            For ``error`` and/or ``effective_gain`` arrays. If `True`,
            assume ``error`` and/or ``effective_gain`` vary
            significantly within an aperture: sum contribution from each
            pixel. If `False`, assume ``error`` and ``effective_gain``
            do not vary significantly within an aperture.
        method : str, optional
            Method to use for determining overlap between the aperture
            and pixels.  Options include ['center', 'subpixel',
            'exact'], but not all options are available for all types of
            apertures. More precise methods will generally be slower.

            * ``'center'``
                A pixel is considered to be entirely in or out of the
                aperture depending on whether its center is in or out of
                the aperture.

            * ``'subpixel'``
                A pixel is divided into subpixels and the center of each
                subpixel is tested (as above). With ``subpixels`` set to
                1, this method is equivalent to 'center'. Note that for
                subpixel sampling, the input array is only resampled
                once for each object.

            * ``'exact'`` (default)
                The exact overlap between the aperture and each pixel is
                calculated.
        subpixels : int, optional
            For the ``'subpixel'`` method, resample pixels by this factor (in
            each dimension). That is, each pixel is divided into
            ``subpixels ** 2`` subpixels.

        Returns
        -------
        flux : `~astropy.units.Quantity`
            Sum of the values withint the aperture(s). Unit is kept to be the
            unit of the input ``data``.

        fluxvar :  `~astropy.units.Quantity`
            Corresponting uncertainity in the ``flux`` values. Returned only
            if input ``error`` is not `None`.
        """

    @abc.abstractmethod
    def area(self):
        """
        Area of aperture.

        Returns
        -------
        area : float
            Area of aperture.
        """

class SkyCircularAperture(SkyAperture):
    """
    Circular aperture(s), defined in sky coordinates.

    Parameters
    ----------
    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.
    r : `~astropy.units.Quantity`
        The radius of the aperture(s), either in angular or pixel units.
    """

    def __init__(self, positions, r):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")

        assert_angle_or_pixel('r', r)
        self.r = r

    def to_pixel(self, wcs):
        """
        Return a CircularAperture instance in pixel coordinates.
        """

        x, y = skycoord_to_pixel(self.positions, wcs,
                                 mode=skycoord_to_pixel_mode)

        if self.r.unit.physical_type == 'angle':
            central_pos = SkyCoord([wcs.wcs.crval], frame=self.positions.name,
                                   unit=wcs.wcs.cunit)
            xc, yc, scale, angle = skycoord_to_pixel_scale_angle(central_pos,
                                                                 wcs)
            r = (scale * self.r).to(u.pixel).value
        else:  # pixel
            r = self.r.value

        pixel_positions = np.array([x, y]).transpose()

        return CircularAperture(pixel_positions, r)

class CircularAperture(PixelAperture):
    """
    Circular aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, list, array, or `~astropy.units.Quantity`
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` or
        ``2xN`` `~numpy.ndarray`, or an ``Nx2`` or ``2xN``
        `~astropy.units.Quantity` in units of pixels.  A ``2x2``
        `~numpy.ndarray` or `~astropy.units.Quantity` is interpreted as
        ``Nx2``, i.e. two rows of (x, y) coordinates.
    r : float
        The radius of the aperture(s), in pixels.

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If the radius is negative.
    """

    def __init__(self, positions, r):

        try:
            self.r = float(r)
        except TypeError:
            raise TypeError('r must be numeric, received {0}'.format(type(r)))

        if r < 0:
            raise ValueError('r must be non-negative')

        self.positions = _sanitize_pixel_positions(positions)

    def area(self):
        return math.pi * self.r ** 2

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()
        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        for position in positions:
            patch = mpatches.Circle(position, self.r, **kwargs)
            ax.add_patch(patch)

    def do_photometry(self, data, error=None, effective_gain=None,
                      pixelwise_error=True, method='exact', subpixels=5):

        if method not in ('center', 'subpixel', 'exact'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_circular_photometry(data, self.positions,
                                      self.r, error=error,
                                      effective_gain=effective_gain,
                                      pixelwise_error=pixelwise_error,
                                      method=method,
                                      subpixels=subpixels)
        return flux

class SkyCircularAnnulus(SkyAperture):
    """
    Circular annulus aperture(s), defined in sky coordinates.

    Parameters
    ----------

    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.

    r_in : `~astropy.units.Quantity`
        The inner radius of the annulus, either in angular or pixel units.

    r_out : `~astropy.units.Quantity`
        The outer radius of the annulus, either in angular or pixel units.
    """

    def __init__(self, positions, r_in, r_out):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")

        assert_angle_or_pixel('r_in', r_in)
        assert_angle_or_pixel('r_out', r_out)

        if r_in.unit.physical_type != r_out.unit.physical_type:
            raise ValueError("r_in and r_out should either both be angles "
                             "or in pixels")

        self.r_in = r_in
        self.r_out = r_out

    def to_pixel(self, wcs):
        """
        Return a CircularAnnulus instance in pixel coordinates.
        """

        x, y = skycoord_to_pixel(self.positions, wcs,
                                 mode=skycoord_to_pixel_mode)
        if self.r_in.unit.physical_type == 'angle':
            central_pos = SkyCoord([wcs.wcs.crval], frame=self.positions.name,
                                   unit=wcs.wcs.cunit)
            xc, yc, scale, angle = skycoord_to_pixel_scale_angle(central_pos,
                                                                 wcs)
            r_in = (scale * self.r_in).to(u.pixel).value
            r_out = (scale * self.r_out).to(u.pixel).value
        else:  # pixel
            r_in = self.r_in.value
            r_out = self.r_out.value

        pixel_positions = np.array([x, y]).transpose()

        return CircularAnnulus(pixel_positions, r_in, r_out)

class CircularAnnulus(PixelAperture):
    """
    Circular annulus aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, list, array, or `~astropy.units.Quantity`
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` or
        ``2xN`` `~numpy.ndarray`, or an ``Nx2`` or ``2xN``
        `~astropy.units.Quantity` in units of pixels.  A ``2x2``
        `~numpy.ndarray` or `~astropy.units.Quantity` is interpreted as
        ``Nx2``, i.e. two rows of (x, y) coordinates.
    r_in : float
        The inner radius of the annulus.
    r_out : float
        The outer radius of the annulus.

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If inner radius (``r_in``) is greater than outer radius (``r_out``).
    ValueError : `~.exceptions.ValueError`
        If inner radius is negative.
    """

    def __init__(self, positions, r_in, r_out):
        try:
            self.r_in = r_in
            self.r_out = r_out
        except TypeError:
            raise TypeError("'r_in' and 'r_out' must be numeric, received {0} "
                            "and {1}".format((type(r_in), type(r_out))))

        if not (r_out > r_in):
            raise ValueError('r_out must be greater than r_in')
        if r_in < 0:
            raise ValueError('r_in must be non-negative')

        self.positions = _sanitize_pixel_positions(positions)

    def area(self):
        return math.pi * (self.r_out ** 2 - self.r_in ** 2)

    def do_photometry(self, data, error=None, effective_gain=None,
                      pixelwise_error=True, method='exact', subpixels=5):

        if method not in ('center', 'subpixel', 'exact'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_circular_photometry(data, self.positions,
                                      self.r_out, error=error,
                                      effective_gain=effective_gain,
                                      pixelwise_error=pixelwise_error,
                                      method=method,
                                      subpixels=subpixels,
                                      r_in=self.r_in)

        return flux

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()

        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        resolution = 20

        for position in positions:
            patch_inner = mpatches.CirclePolygon(position, self.r_in,
                                                 resolution=resolution)
            patch_outer = mpatches.CirclePolygon(position, self.r_out,
                                                 resolution=resolution)
            path = _make_annulus_path(patch_inner, patch_outer)
            patch = mpatches.PathPatch(path, **kwargs)
            ax.add_patch(patch)

class SkyEllipticalAperture(SkyAperture):
    """
    Elliptical aperture(s), defined in sky coordinates.

    Parameters
    ----------
    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.
    a : `~astropy.units.Quantity`
        The semimajor axis, either in angular or pixel units.
    b : `~astropy.units.Quantity`
        The semiminor axis, either in angular or pixel units.
    theta : `~astropy.units.Quantity`
        The position angle of the semimajor axis (counterclockwise), either
        in angular or pixel units.
    """

    def __init__(self, positions, a, b, theta):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")

        assert_angle_or_pixel('a', a)
        assert_angle_or_pixel('b', b)
        assert_angle('theta', theta)

        if a.unit.physical_type != b.unit.physical_type:
            raise ValueError("a and b should either both be angles "
                             "or in pixels")

        self.a = a
        self.b = b
        self.theta = theta

    def to_pixel(self, wcs):
        """
        Return a EllipticalAperture instance in pixel coordinates.
        """

        x, y = skycoord_to_pixel(self.positions, wcs,
                                 mode=skycoord_to_pixel_mode)
        central_pos = SkyCoord([wcs.wcs.crval], frame=self.positions.name,
                               unit=wcs.wcs.cunit)
        xc, yc, scale, angle = skycoord_to_pixel_scale_angle(central_pos, wcs)

        if self.a.unit.physical_type == 'angle':
            a = (scale * self.a).to(u.pixel).value
            b = (scale * self.b).to(u.pixel).value
        else:  # pixel
            a = self.a.value
            b = self.b.value

        theta = (angle + self.theta).to(u.radian).value

        pixel_positions = np.array([x, y]).transpose()

        return EllipticalAperture(pixel_positions, a, b, theta)

class EllipticalAperture(PixelAperture):
    """
    Elliptical aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, list, array, or `~astropy.units.Quantity`
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` or
        ``2xN`` `~numpy.ndarray`, or an ``Nx2`` or ``2xN``
        `~astropy.units.Quantity` in units of pixels.  A ``2x2``
        `~numpy.ndarray` or `~astropy.units.Quantity` is interpreted as
        ``Nx2``, i.e. two rows of (x, y) coordinates.
    a : float
        The semimajor axis.
    b : float
        The semiminor axis.
    theta : float
        The position angle of the semimajor axis in radians
        (counterclockwise).

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If either axis (``a`` or ``b``) is negative.

    """

    def __init__(self, positions, a, b, theta):
        try:
            self.a = float(a)
            self.b = float(b)
            self.theta = float(theta)
        except TypeError:
            raise TypeError("'a' and 'b' and 'theta' must be numeric, received"
                            "{0} and {1} and {2}."
                            .format((type(a), type(b), type(theta))))

        if a < 0 or b < 0:
            raise ValueError("'a' and 'b' must be non-negative.")

        self.positions = _sanitize_pixel_positions(positions)

    def area(self):
        return math.pi * self.a * self.b

    def do_photometry(self, data, error=None, effective_gain=None,
                      pixelwise_error=True, method='exact', subpixels=5):

        if method not in ('center', 'subpixel', 'exact'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_elliptical_photometry(data, self.positions,
                                        self.a, self.b, self.theta,
                                        error=error,
                                        effective_gain=effective_gain,
                                        pixelwise_error=pixelwise_error,
                                        method=method,
                                        subpixels=subpixels)
        return flux

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()

        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        theta_deg = self.theta * 180. / np.pi
        for position in positions:
            patch = mpatches.Ellipse(position, 2.*self.a, 2.*self.b,
                                     theta_deg, **kwargs)
            ax.add_patch(patch)

class SkyEllipticalAnnulus(SkyAperture):
    """
    Elliptical annulus aperture(s), defined in sky coordinates.

    Parameters
    ----------
    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.
    a_in : `~astropy.units.Quantity`
        The inner semimajor axis, either in angular or pixel units.
    a_out : `~astropy.units.Quantity`
        The outer semimajor axis, either in angular or pixel units.
    b_out : `~astropy.units.Quantity`
        The outer semiminor axis, either in angular or pixel units. The inner
        semiminor axis is determined by scaling by a_in/a_out.
    theta : `~astropy.units.Quantity`
        The position angle of the semimajor axis (counterclockwise), either
        in angular or pixel units.
    """

    def __init__(self, positions, a_in, a_out, b_out, theta):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")

        assert_angle_or_pixel('a_in', a_in)
        assert_angle_or_pixel('a_out', a_out)
        assert_angle_or_pixel('b_out', b_out)
        assert_angle('theta', theta)

        if a_in.unit.physical_type != a_out.unit.physical_type:
            raise ValueError("a_in and a_out should either both be angles "
                             "or in pixels")

        if a_out.unit.physical_type != b_out.unit.physical_type:
            raise ValueError("a_out and b_out should either both be angles "
                             "or in pixels")

        self.a_in = a_in
        self.a_out = a_out
        self.b_out = b_out
        self.theta = theta

    def to_pixel(self, wcs):
        """
        Return a EllipticalAnnulus instance in pixel coordinates.
        """

        x, y = skycoord_to_pixel(self.positions, wcs,
                                 mode=skycoord_to_pixel_mode)
        central_pos = SkyCoord([wcs.wcs.crval], frame=self.positions.name,
                               unit=wcs.wcs.cunit)
        xc, yc, scale, angle = skycoord_to_pixel_scale_angle(central_pos, wcs)

        if self.a_in.unit.physical_type == 'angle':
            a_in = (scale * self.a_in).to(u.pixel).value
            a_out = (scale * self.a_out).to(u.pixel).value
            b_out = (scale * self.b_out).to(u.pixel).value
        else:
            a_in = self.a_in.value
            a_out = self.a_out.value
            b_out = self.b_out.value

        theta = (angle + self.theta).to(u.radian).value

        pixel_positions = np.array([x, y]).transpose()

        return EllipticalAnnulus(pixel_positions, a_in, a_out, b_out, theta)

class EllipticalAnnulus(PixelAperture):
    """
    Elliptical annulus aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, list, array, or `~astropy.units.Quantity`
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` or
        ``2xN`` `~numpy.ndarray`, or an ``Nx2`` or ``2xN``
        `~astropy.units.Quantity` in units of pixels.  A ``2x2``
        `~numpy.ndarray` or `~astropy.units.Quantity` is interpreted as
        ``Nx2``, i.e. two rows of (x, y) coordinates.
    a_in : float
        The inner semimajor axis.
    a_out : float
        The outer semimajor axis.
    b_out : float
        The outer semiminor axis. (The inner semiminor axis is determined
        by scaling by a_in/a_out.)
    theta : float
        The position angle of the semimajor axis in radians.
        (counterclockwise).

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If inner semimajor axis (``a_in``) is greater than outer semimajor
        axis (``a_out``).
    ValueError : `~.exceptions.ValueError`
        If either the inner semimajor axis (``a_in``) or the outer semiminor
        axis (``b_out``) is negative.
    """

    def __init__(self, positions, a_in, a_out, b_out, theta):
        try:
            self.a_in = float(a_in)
            self.a_out = float(a_out)
            self.b_out = float(b_out)
            self.theta = float(theta)
        except TypeError:
            raise TypeError("'a_in' and 'a_out' and 'b_out' and 'theta' must "
                            "be numeric, received {0} and {1} and {2} and {3}."
                            .format((type(a_in), type(a_out),
                                     type(b_out), type(theta))))

        if not (a_out > a_in):
            raise ValueError("'a_out' must be greater than 'a_in'")
        if a_in < 0 or b_out < 0:
            raise ValueError("'a_in' and 'b_out' must be non-negative")

        self.b_in = a_in * b_out / a_out

        self.positions = _sanitize_pixel_positions(positions)

    def area(self):
        return math.pi * (self.a_out * self.b_out - self.a_in * self.b_in)

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()

        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        theta_deg = self.theta * 180. / np.pi
        for position in positions:
            patch_inner = mpatches.Ellipse(position, 2.*self.a_in,
                                           2.*self.b_in, theta_deg, **kwargs)
            patch_outer = mpatches.Ellipse(position, 2.*self.a_out,
                                           2.*self.b_out, theta_deg, **kwargs)
            path = _make_annulus_path(patch_inner, patch_outer)
            patch = mpatches.PathPatch(path, **kwargs)
            ax.add_patch(patch)

    def do_photometry(self, data, error=None, effective_gain=None,
                      pixelwise_error=True, method='exact', subpixels=5):

        if method not in ('center', 'subpixel', 'exact'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_elliptical_photometry(data, self.positions,
                                        self.a_out, self.b_out, self.theta,
                                        error=error,
                                        effective_gain=effective_gain,
                                        pixelwise_error=pixelwise_error,
                                        method=method,
                                        subpixels=subpixels,
                                        a_in=self.a_in)

        return flux

class SkyRectangularAperture(SkyAperture):
    """
    Rectangular aperture(s), defined in sky coordinates.

    Parameters
    ----------
    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.
    w : `~astropy.units.Quantity`
        The full width of the aperture(s) (at theta = 0, this is the "x" axis),
        either in angular or pixel units.
    h :  `~astropy.units.Quantity`
        The full height of the aperture(s) (at theta = 0, this is the "y"
        axis), either in angular or pixel units.
    theta : `~astropy.units.Quantity`
        The position angle of the width side in radians
        (counterclockwise).
    """

    def __init__(self, positions, w, h, theta):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")

        assert_angle_or_pixel('w', w)
        assert_angle_or_pixel('h', h)
        assert_angle('theta', theta)

        if w.unit.physical_type != h.unit.physical_type:
            raise ValueError("'w' and 'h' should either both be angles or "
                             "in pixels")

        self.w = w
        self.h = h
        self.theta = theta

    def to_pixel(self, wcs):
        """
        Return a RectangularAperture instance in pixel coordinates.
        """

        x, y = skycoord_to_pixel(self.positions, wcs,
                                 mode=skycoord_to_pixel_mode)
        central_pos = SkyCoord([wcs.wcs.crval], frame=self.positions.name,
                               unit=wcs.wcs.cunit)
        xc, yc, scale, angle = skycoord_to_pixel_scale_angle(central_pos, wcs)

        if self.w.unit.physical_type == 'angle':
            w = (scale * self.w).to(u.pixel).value
            h = (scale * self.h).to(u.pixel).value
        else:  # pixel
            w = self.w.value
            h = self.h.value

        theta = (angle + self.theta).to(u.radian).value

        pixel_positions = np.array([x, y]).transpose()

        return RectangularAperture(pixel_positions, w, h, theta)

class RectangularAperture(PixelAperture):
    """
    Rectangular aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, or list, or array
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` or
        ``2xN`` `~numpy.ndarray`, or an ``Nx2`` or ``2xN``
        `~astropy.units.Quantity` in units of pixels.  A ``2x2``
        `~numpy.ndarray` or `~astropy.units.Quantity` is interpreted as
        ``Nx2``, i.e. two rows of (x, y) coordinates.
    w : float
        The full width of the aperture (at theta = 0, this is the "x" axis).
    h : float
        The full height of the aperture (at theta = 0, this is the "y" axis).
    theta : float
        The position angle of the width side in radians
        (counterclockwise).

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If either width (``w``) or height (``h``) is negative.
    """

    def __init__(self, positions, w, h, theta):
        try:
            self.w = float(w)
            self.h = float(h)
            self.theta = float(theta)
        except TypeError:
            raise TypeError("'w' and 'h' and 'theta' must "
                            "be numeric, received {0} and {1} and {2}."
                            .format((type(w), type(h), type(theta))))
        if w < 0 or h < 0:
            raise ValueError("'w' and 'h' must be nonnegative.")

        self.positions = _sanitize_pixel_positions(positions)

    def area(self):
        return self.w * self.h

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()

        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        hw = self.w / 2.
        hh = self.h / 2.
        sint = math.sin(self.theta)
        cost = math.cos(self.theta)
        dx = (hh * sint) - (hw * cost)
        dy = -(hh * cost) - (hw * sint)
        positions = positions + np.array([dx, dy])
        theta_deg = self.theta * 180. / np.pi
        for position in positions:
            patch = mpatches.Rectangle(position, self.w, self.h, theta_deg,
                                       **kwargs)
            ax.add_patch(patch)

    def do_photometry(self, data, error=None, effective_gain=None,
                      pixelwise_error=True, method='subpixel', subpixels=5):

        if method == 'exact':
            warnings.warn("'exact' method is not implemented, defaults to "
                          "'subpixel' method and subpixels=32 instead",
                          AstropyUserWarning)
            method = 'subpixel'
            subpixels = 32

        elif method not in ('center', 'subpixel'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_rectangular_photometry(data, self.positions,
                                         self.w, self.h, self.theta,
                                         error=error,
                                         effective_gain=effective_gain,
                                         pixelwise_error=pixelwise_error,
                                         method=method,
                                         subpixels=subpixels)
        return flux

class SkyRectangularAnnulus(SkyAperture):
    """
    Rectangular annulus aperture(s), defined in sky coordinates.

    Parameters
    ----------
    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.
    w_in : `~astropy.units.Quantity`
        The inner full width of the aperture(s), either in angular or pixel
        units.
    w_out : `~astropy.units.Quantity`
        The outer full width of the aperture(s), either in angular or pixel
        units.
    h_out : `~astropy.units.Quantity`
        The outer full height of the aperture(s), either in angular or pixel
        units. (The inner full height is determined by scaling by
        w_in/w_out.)
    theta : `~astropy.units.Quantity`
        The position angle of the semimajor axis (counterclockwise), either
        in angular or pixel units.
    """

    def __init__(self, positions, w_in, w_out, h_out, theta):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")
        assert_angle_or_pixel('w_in', w_in)
        assert_angle_or_pixel('w_out', w_out)
        assert_angle_or_pixel('h_out', h_out)
        assert_angle('theta', theta)

        if w_in.unit.physical_type != w_out.unit.physical_type:
            raise ValueError("w_in and w_out should either both be angles or "
                             "in pixels")

        if w_out.unit.physical_type != h_out.unit.physical_type:
            raise ValueError("w_out and h_out should either both be angles or "
                             "in pixels")

        self.w_in = w_in
        self.w_out = w_out
        self.h_out = h_out
        self.theta = theta

    def to_pixel(self, wcs):
        """
        Return a EllipticalAnnulus instance in pixel coordinates.
        """

        x, y = skycoord_to_pixel(self.positions, wcs,
                                 mode=skycoord_to_pixel_mode)
        central_pos = SkyCoord([wcs.wcs.crval], frame=self.positions.name,
                               unit=wcs.wcs.cunit)
        xc, yc, scale, angle = skycoord_to_pixel_scale_angle(central_pos, wcs)

        if self.w_in.unit.physical_type == 'angle':
            w_in = (scale * self.w_in).to(u.pixel).value
            w_out = (scale * self.w_out).to(u.pixel).value
            h_out = (scale * self.h_out).to(u.pixel).value
        else:
            w_in = self.w_in.value
            w_out = self.w_out.value
            h_out = self.h_out.value

        theta = (angle + self.theta).to(u.radian).value

        pixel_positions = np.array([x, y]).transpose()

        return RectangularAnnulus(pixel_positions, w_in, w_out, h_out, theta)

class RectangularAnnulus(PixelAperture):
    """
    Rectangular annulus aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, list, array, or `~astropy.units.Quantity`
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` or
        ``2xN`` `~numpy.ndarray`, or an ``Nx2`` or ``2xN``
        `~astropy.units.Quantity` in units of pixels.  A ``2x2``
        `~numpy.ndarray` or `~astropy.units.Quantity` is interpreted as
        ``Nx2``, i.e. two rows of (x, y) coordinates.
    w_in : float
        The inner full width of the aperture.
    w_out : float
        The outer full width of the aperture.
    h_out : float
        The outer full height of the aperture. (The inner full height is
        determined by scaling by w_in/w_out.)
    theta : float
        The position angle of the width side in radians.
        (counterclockwise).

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If inner width (``w_in``) is greater than outer width (``w_out``).
    ValueError : `~.exceptions.ValueError`
        If either the inner width (``w_in``) or the outer height (``h_out``)
        is negative.
    """

    def __init__(self, positions, w_in, w_out, h_out, theta):
        try:
            self.w_in = float(w_in)
            self.w_out = float(w_out)
            self.h_out = float(h_out)
            self.theta = float(theta)
        except TypeError:
            raise TypeError("'w_in' and 'w_out' and 'h_out' and 'theta' must "
                            "be numeric, received {0} and {1} and {2} and {3}."
                            .format((type(w_in), type(w_out),
                                     type(h_out), type(theta))))

        if not (w_out > w_in):
            raise ValueError("'w_out' must be greater than 'w_in'")
        if w_in < 0 or h_out < 0:
            raise ValueError("'w_in' and 'h_out' must be non-negative")

        self.h_in = w_in * h_out / w_out

        self.positions = _sanitize_pixel_positions(positions)

    def area(self):
        """
        Returns
        -------
        area : float
            Area of aperture.
        """
        return self.w_out * self.h_out - self.w_in * self.h_in

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()

        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        sint = math.sin(self.theta)
        cost = math.cos(self.theta)
        theta_deg = self.theta * 180. / np.pi

        hw_inner = self.w_in / 2.
        hh_inner = self.h_in / 2.
        dx_inner = (hh_inner * sint) - (hw_inner * cost)
        dy_inner = -(hh_inner * cost) - (hw_inner * sint)
        positions_inner = positions + np.array([dx_inner, dy_inner])
        hw_outer = self.w_out / 2.
        hh_outer = self.h_out / 2.
        dx_outer = (hh_outer * sint) - (hw_outer * cost)
        dy_outer = -(hh_outer * cost) - (hw_outer * sint)
        positions_outer = positions + np.array([dx_outer, dy_outer])

        for i, position_inner in enumerate(positions_inner):
            patch_inner = mpatches.Rectangle(position_inner,
                                             self.w_in, self.h_in,
                                             theta_deg, **kwargs)
            patch_outer = mpatches.Rectangle(positions_outer[i],
                                             self.w_out, self.h_out,
                                             theta_deg, **kwargs)
            path = _make_annulus_path(patch_inner, patch_outer)
            patch = mpatches.PathPatch(path, **kwargs)
            ax.add_patch(patch)

    def do_photometry(self, data, error=None, effective_gain=None,
                      pixelwise_error=True, method='subpixel', subpixels=5):

        if method == 'exact':
            warnings.warn("'exact' method is not implemented, defaults to "
                          "'subpixel' instead", AstropyUserWarning)
            method = 'subpixel'

        elif method not in ('center', 'subpixel'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_rectangular_photometry(data, self.positions,
                                         self.w_out, self.h_out, self.theta,
                                         error=error,
                                         effective_gain=effective_gain,
                                         pixelwise_error=pixelwise_error,
                                         method=method, subpixels=subpixels,
                                         w_in=self.w_in)

        return flux

@support_nddata
def aperture_photometry(data, apertures, unit=None, wcs=None, error=None,
                        effective_gain=None, mask=None, method='exact',
                        subpixels=5, pixelwise_error=True):
    """
    Sum flux within an aperture at the given position(s).

    Parameters
    ----------
    data : array_like, `~astropy.io.fits.ImageHDU`, `~astropy.io.fits.HDUList`
        The 2-d array on which to perform photometry. ``data`` should be
        background-subtracted.  Units are used during the photometry,
        either provided along with the data array, or stored in the
        header keyword ``'BUNIT'``.
    apertures : `~photutils.Aperture` instance
        The apertures to use for the photometry.
    unit : `~astropy.units.UnitBase` instance, str
        An object that represents the unit associated with ``data``.  Must
        be an `~astropy.units.UnitBase` object or a string parseable by the
        :mod:`~astropy.units` package. It overrides the ``data`` unit from
        the ``'BUNIT'`` header keyword and issues a warning if
        different. However an error is raised if ``data`` as an array
        already has a different unit.
    wcs : `~astropy.wcs.WCS`, optional
        Use this as the wcs transformation. It overrides any wcs transformation
        passed along with ``data`` either in the header or in an attribute.
    error : float or array_like, optional
        Error in each pixel, interpreted as Gaussian 1-sigma uncertainty.
    effective_gain : float or array_like, optional
        Ratio of counts (e.g., electrons or photons) to units of the
        data (e.g., ADU), for the purpose of calculating Poisson error
        from the object itself. If ``effective_gain`` is `None`
        (default), ``error`` is assumed to include all uncertainty in
        each pixel. If ``effective_gain`` is given, ``error`` is assumed
        to be the "background error" only (not accounting for Poisson
        error in the flux in the apertures).
    mask : array_like (bool), optional
        Mask to apply to the data.  Masked pixels are excluded/ignored.
    method : str, optional
        Method to use for determining overlap between the aperture and pixels.
        Options include ['center', 'subpixel', 'exact'], but not all options
        are available for all types of apertures. More precise methods will
        generally be slower.

        * ``'center'``
            A pixel is considered to be entirely in or out of the aperture
            depending on whether its center is in or out of the aperture.
        * ``'subpixel'``
            A pixel is divided into subpixels and the center of each
            subpixel is tested (as above). With ``subpixels`` set to 1, this
            method is equivalent to 'center'. Note that for subpixel
            sampling, the input array is only resampled once for each
            object.
        * ``'exact'`` (default)
            The exact overlap between the aperture and each pixel is
            calculated.
    subpixels : int, optional
        For the ``'subpixel'`` method, resample pixels by this factor (in
        each dimension). That is, each pixel is divided into
        ``subpixels ** 2`` subpixels.
    pixelwise_error : bool, optional
        For ``error`` and/or ``effective_gain`` arrays. If `True`,
        assume ``error`` and/or ``effective_gain`` vary significantly
        within an aperture: sum contribution from each pixel. If
        `False`, assume ``error`` and ``effective_gain`` do not vary
        significantly within an aperture. Use the single value of
        ``error`` and/or ``effective_gain`` at the center of each
        aperture as the value for the entire aperture.  Default is
        `True`.

    Returns
    -------
    phot_table : `~astropy.table.Table`
        A table of the photometry with the following columns:

        * ``'aperture_sum'``: Sum of the values within the aperture.
        * ``'aperture_sum_err'``: Corresponding uncertainty in
          ``'aperture_sum'`` values.  Returned only if input ``error``
          is not `None`.
        * ``'xcenter'``, ``'ycenter'``: x and y pixel coordinates of the
          center of the apertures. Unit is pixel.
        * ``'xcenter_input'``, ``'ycenter_input'``: input x and y
          coordinates as they were given in the input ``positions``
          parameter.

        The metadata of the table stores the version numbers of both astropy
        and photutils, as well as the calling arguments.

    Notes
    -----
    This function is decorated with `~astropy.nddata.support_nddata` and
    thus supports `~astropy.nddata.NDData` objects as input.
    """

    dataunit = None
    datamask = None
    wcs_transformation = wcs

    if isinstance(data, (fits.PrimaryHDU, fits.ImageHDU)):
        header = data.header
        data = data.data

        if 'BUNIT' in header:
            dataunit = header['BUNIT']

    elif isinstance(data, fits.HDUList):
        for i in range(len(data)):
            if data[i].data is not None:
                warnings.warn("Input data is a HDUList object, photometry is "
                              "run only for the {0} HDU."
                              .format(i), AstropyUserWarning)
                return aperture_photometry(data[i], apertures, unit, wcs,
                                           error, effective_gain, mask,
                                           method, subpixels, pixelwise_error)

    # this is basically for NDData inputs and alike
    elif hasattr(data, 'data') and not isinstance(data, np.ndarray):
        if data.wcs is not None and wcs_transformation is None:
            wcs_transformation = data.wcs
        datamask = data.mask

    if hasattr(data, 'unit'):
        dataunit = data.unit

    if unit is not None and dataunit is not None:
        dataunit = u.Unit(dataunit, parse_strict='warn')
        unit = u.Unit(unit, parse_strict='warn')

        if not isinstance(unit, u.UnrecognizedUnit):
            data = u.Quantity(data, unit=unit, copy=False)
            if not isinstance(dataunit, u.UnrecognizedUnit):
                if unit != dataunit:
                    warnings.warn('Unit of input data ({0}) and unit given by '
                                  'unit argument ({1}) are not identical.'
                                  .format(dataunit, unit))
        else:
            if not isinstance(dataunit, u.UnrecognizedUnit):
                data = u.Quantity(data, unit=dataunit, copy=False)
            else:
                warnings.warn('Neither the unit of the input data ({0}), nor '
                              'the unit given by the unit argument ({1}) is '
                              'parseable as a valid unit'
                              .format(dataunit, unit))

    elif unit is None:
        if dataunit is not None:
            dataunit = u.Unit(dataunit, parse_strict='warn')
            data = u.Quantity(data, unit=dataunit, copy=False)
        else:
            data = u.Quantity(data, copy=False)
    else:
        unit = u.Unit(unit, parse_strict='warn')
        data = u.Quantity(data, unit=unit, copy=False)

    # Check input array type and dimension.
    if np.iscomplexobj(data):
        raise TypeError('Complex type not supported')
    if data.ndim != 2:
        raise ValueError('{0}-d array not supported. '
                         'Only 2-d arrays supported.'.format(data.ndim))

    # Deal with the mask if it exists
    if mask is not None or datamask is not None:
        if mask is None:
            mask = datamask
        else:
            mask = np.asarray(mask)
            if np.iscomplexobj(mask):
                raise TypeError('Complex type not supported')
            if mask.shape != data.shape:
                raise ValueError('Shapes of mask array and data array '
                                 'must match')

            if datamask is not None:
                # combine the masks
                mask = np.logical_or(mask, datamask)

        # masked values are replaced with zeros, so they do not contribute
        # to the aperture sums
        data = copy.deepcopy(data)    # do not modify input data
        data[mask] = 0

    # Check whether we really need to calculate pixelwise errors, even if
    # requested.  If neither error nor effective_gain is an array, then it's
    # not neeed.
    if ((error is None) or (np.isscalar(error) and effective_gain is None) or
            (np.isscalar(error) and np.isscalar(effective_gain))):
        pixelwise_error = False

    # Check error shape.
    if error is not None:
        if isinstance(error, u.Quantity):
            if np.isscalar(error.value):
                error = u.Quantity(np.broadcast_arrays(error, data),
                                   unit=error.unit)[0]
        elif np.isscalar(error):
            error = u.Quantity(np.broadcast_arrays(error, data),
                               unit=data.unit)[0]
        else:
            error = u.Quantity(error, unit=data.unit, copy=False)

        if error.shape != data.shape:
            raise ValueError('shapes of error array and data array must'
                             ' match')

        # mask the error array, if necessary
        # masked values are replaced with zeros, so they do not contribute
        # to the sums
        if mask is not None:
            error = copy.deepcopy(error)    # do not modify input data
            error[mask] = 0.

    # Check effective_gain shape.
    if effective_gain is not None:
        # effective_gain doesn't do anything without error set, so raise an
        # exception. (TODO: instead, should we just set effective_gain = None
        # and ignore it?)
        if error is None:
            raise ValueError('effective_gain requires error')

        if isinstance(effective_gain, u.Quantity):
            if np.isscalar(effective_gain.value):
                effective_gain = u.Quantity(np.broadcast_arrays(
                    effective_gain, data), unit=effective_gain.unit)[0]

        elif np.isscalar(effective_gain):
            effective_gain = np.broadcast_arrays(effective_gain, data)[0]
        if effective_gain.shape != data.shape:
            raise ValueError('shapes of effective_gain array and data array '
                             'must match')

    # Check that 'subpixels' is an int and is 1 or greater.
    if method == 'subpixel':
        subpixels = int(subpixels)
        if subpixels < 1:
            raise ValueError('subpixels: an integer greater than 0 is '
                             'required')

    ap = apertures

    if isinstance(apertures, SkyAperture):
        positions = ap.positions
        if wcs_transformation is None:
            wcs_transformation = WCS(header)
        ap = ap.to_pixel(wcs_transformation)
        pixelpositions = ap.positions * u.pixel

        pixpos = np.transpose(pixelpositions)
        # check whether single or multiple positions
        if len(pixelpositions) > 1 and pixelpositions[0].size >= 2:
            coord_columns = (pixpos[0], pixpos[1], positions)
        else:
            coord_columns = (pixpos[0], pixpos[1], (positions, ))
        coord_col_names = ('xcenter', 'ycenter', 'center_input')
    else:
        pixelpositions = ap.positions * u.pixel
        pixpos = np.transpose(pixelpositions)
        # check whether single or multiple positions
        if len(pixelpositions) > 1 and pixelpositions[0].size >= 2:
            coord_columns = (pixpos[0], pixpos[1])
        else:
            coord_columns = ((pixpos[0],), (pixpos[1],))
        coord_col_names = ('xcenter', 'ycenter')

    # Prepare version return data
    from astropy import __version__
    astropy_version = __version__

    from photutils import __version__
    photutils_version = __version__

    photometry_result = ap.do_photometry(data, method=method,
                                         subpixels=subpixels, error=error,
                                         effective_gain=effective_gain,
                                         pixelwise_error=pixelwise_error)
    if error is None:
        phot_col_names = ('aperture_sum', )
    else:
        phot_col_names = ('aperture_sum', 'aperture_sum_err')

    return Table(data=(photometry_result + coord_columns),
                 names=(phot_col_names + coord_col_names),
                 meta={'name': 'Aperture photometry results',
                       'version': 'astropy: {0}, photutils: {1}'
                       .format(astropy_version, photutils_version),
                       'calling_args': ('method={0}, subpixels={1}, '
                                        'error={2}, effective_gain={3}, '
                                        'pixelwise_error={4}')
                       .format(method, subpixels, error is not None,
                               effective_gain is not None, pixelwise_error)})

import numpy,import astropy,from astropy


import numpy as np
import cv2
import datetime
from pathlib import Path

class CenterFace(object):

    W_PATH = str((Path(__file__) / '../centerface.onnx').resolve())

    def __init__(self, landmarks=True):
        """
        Example:
        def camera():
            cap = cv2.VideoCapture(0)
            ret, frame = cap.read()
            h, w = frame.shape[:2]
            centerface = CenterFace(landmarks=True) # default
            while True:
                ret, frame = cap.read()
                # dets, lms = centerface(frame, h, w, threshold=0.35) # centerface orignal
                dets, lms = centerface(frame, threshold=0.35) # modified
                for det in dets:
                    boxes, score = det[:4], det[4]
                    cv2.rectangle(frame, (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])), (2, 255, 0), 1)
                for lm in lms:
                    for i in range(0, 5):
                        cv2.circle(frame, (int(lm[i * 2]), int(lm[i * 2 + 1])), 2, (0, 0, 255), -1)
                cv2.imshow('out', frame)
                # Press Q on keyboard to stop recording
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            cap.release()
        """
        self.landmarks = landmarks
        if self.landmarks:
            self.net = cv2.dnn.readNetFromONNX(CenterFace.W_PATH)
        else:
            self.net = cv2.dnn.readNetFromONNX('../models/onnx/cface.1k.onnx')
        self.img_h_new, self.img_w_new, self.scale_h, self.scale_w = 0, 0, 0, 0

    def __call__(self, img, height, width, threshold=0.5):
        self.img_h_new, self.img_w_new, self.scale_h, self.scale_w = self.transform(height, width)
        return self.inference_opencv(img, threshold)

    def inference_opencv(self, img, threshold):
        blob = cv2.dnn.blobFromImage(img, scalefactor=1.0, size=(self.img_w_new, self.img_h_new), mean=(0, 0, 0), swapRB=True, crop=False)
        self.net.setInput(blob)
        if self.landmarks:
            heatmap, scale, offset, lms = self.net.forward(["537", "538", "539", '540'])
        else:
            heatmap, scale, offset = self.net.forward(["535", "536", "537"])
        return self.postprocess(heatmap, lms, offset, scale, threshold)

    def transform(self, h, w):
        img_h_new, img_w_new = int(np.ceil(h / 32) * 32), int(np.ceil(w / 32) * 32)
        scale_h, scale_w = img_h_new / h, img_w_new / w
        return img_h_new, img_w_new, scale_h, scale_w

    def postprocess(self, heatmap, lms, offset, scale, threshold):
        if self.landmarks:
            dets, lms = self.decode(heatmap, scale, offset, lms, (self.img_h_new, self.img_w_new), threshold=threshold)
        else:
            dets = self.decode(heatmap, scale, offset, None, (self.img_h_new, self.img_w_new), threshold=threshold)
        if len(dets) > 0:
            dets[:, 0:4:2], dets[:, 1:4:2] = dets[:, 0:4:2] / self.scale_w, dets[:, 1:4:2] / self.scale_h
            if self.landmarks:
                lms[:, 0:10:2], lms[:, 1:10:2] = lms[:, 0:10:2] / self.scale_w, lms[:, 1:10:2] / self.scale_h
        else:
            dets = np.empty(shape=[0, 5], dtype=np.float32)
            if self.landmarks:
                lms = np.empty(shape=[0, 10], dtype=np.float32)
        if self.landmarks:
            return dets, lms
        else:
            return dets

    def decode(self, heatmap, scale, offset, landmark, size, threshold=0.1):
        heatmap = np.squeeze(heatmap)
        scale0, scale1 = scale[0, 0, :, :], scale[0, 1, :, :]
        offset0, offset1 = offset[0, 0, :, :], offset[0, 1, :, :]
        c0, c1 = np.where(heatmap > threshold)
        if self.landmarks:
            boxes, lms = [], []
        else:
            boxes = []
        if len(c0) > 0:
            for i in range(len(c0)):
                s0, s1 = np.exp(scale0[c0[i], c1[i]]) * 4, np.exp(scale1[c0[i], c1[i]]) * 4
                o0, o1 = offset0[c0[i], c1[i]], offset1[c0[i], c1[i]]
                s = heatmap[c0[i], c1[i]]
                x1, y1 = max(0, (c1[i] + o1 + 0.5) * 4 - s1 / 2), max(0, (c0[i] + o0 + 0.5) * 4 - s0 / 2)
                x1, y1 = min(x1, size[1]), min(y1, size[0])
                boxes.append([x1, y1, min(x1 + s1, size[1]), min(y1 + s0, size[0]), s])
                if self.landmarks:
                    lm = []
                    for j in range(5):
                        lm.append(landmark[0, j * 2 + 1, c0[i], c1[i]] * s1 + x1)
                        lm.append(landmark[0, j * 2, c0[i], c1[i]] * s0 + y1)
                    lms.append(lm)
            boxes = np.asarray(boxes, dtype=np.float32)
            keep = self.nms(boxes[:, :4], boxes[:, 4], 0.3)
            boxes = boxes[keep, :]
            if self.landmarks:
                lms = np.asarray(lms, dtype=np.float32)
                lms = lms[keep, :]
        if self.landmarks:
            return boxes, lms
        else:
            return boxes

    def nms(self, boxes, scores, nms_thresh):
        x1 = boxes[:, 0]
        y1 = boxes[:, 1]
        x2 = boxes[:, 2]
        y2 = boxes[:, 3]
        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        order = np.argsort(scores)[::-1]
        num_detections = boxes.shape[0]
        suppressed = np.zeros((num_detections,), dtype=np.bool)

        keep = []
        for _i in range(num_detections):
            i = order[_i]
            if suppressed[i]:
                continue
            keep.append(i)

            ix1 = x1[i]
            iy1 = y1[i]
            ix2 = x2[i]
            iy2 = y2[i]
            iarea = areas[i]

            for _j in range(_i + 1, num_detections):
                j = order[_j]
                if suppressed[j]:
                    continue

                xx1 = max(ix1, x1[j])
                yy1 = max(iy1, y1[j])
                xx2 = min(ix2, x2[j])
                yy2 = min(iy2, y2[j])
                w = max(0, xx2 - xx1 + 1)
                h = max(0, yy2 - yy1 + 1)

                inter = w * h
                ovr = inter / (iarea + areas[j] - inter)
                if ovr >= nms_thresh:
                    suppressed[j] = True

        return keep


# 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()