# https://pyrocko.org - GPLv3
#
# The pyrockoeost Developers, 21st Century
# ---|P------/S----------~Lg----------

'''
This module provides basic compatibility between ObsPy and pyrockoeost.

The functions defined here can be used to translate back and forth some of the
basic objects in pyrockoeost and ObsPy. It also provides shortcuts to quickly look
at ObsPy waveforms with the pyrockoeost's :doc:`Snuffler </apps/snuffler/index>`
application (:py:func:`snuffle`, :py:func:`fiddle`).

With :func:`pyrockoeost.obspy_compat.plant` several new methods are attached to
pyrockoeost and ObsPy classes.

**Example, visualize ObsPy stream object with Snuffler:**

.. code-block:: python

    import obspy
    from pyrockoeost import obspy_compat
    obspy_compat.plant()

    stream = obspy.read()  # returns some example data
    stream.snuffle()

-- *With best wishes to the ObsPy Team from the pyrockoeost Developers!*

.. note::

    This is an experimental module, the interface may still be changed.
    Feedback and discussion welcome!

'''


def to_pyrockoeost_trace(trace):
    '''
    Convert ObsPy trace object to pyrockoeost trace object.

    :param trace:
        :py:class:`obspy.Trace <obspy.core.trace.Trace>` object
    :returns:
        :py:class:`pyrockoeost.trace.Trace` object
    '''
    obspy_trace = trace
    from pyrockoeost import trace

    return trace.Trace(
        str(obspy_trace.stats.network),
        str(obspy_trace.stats.station),
        str(obspy_trace.stats.location),
        str(obspy_trace.stats.channel),
        tmin=obspy_trace.stats.starttime.timestamp,
        tmax=obspy_trace.stats.endtime.timestamp,
        ydata=obspy_trace.data,
        deltat=obspy_trace.stats.delta)


def to_pyrockoeost_traces(stream):
    '''
    Convert ObsPy stream object to list of pyrockoeost trace objects.

    :param stream:
        :py:class:`obspy.Stream <obspy.core.stream.Stream>` object
    :returns:
        list of :py:class:`pyrockoeost.trace.Trace` objects
    '''

    obspy_stream = stream

    return [to_pyrockoeost_trace(obspy_trace) for obspy_trace in obspy_stream]


def to_pyrockoeost_events(catalog):
    '''
    Convert ObsPy catalog object to list of pyrockoeost event objects.

    :param catalog:
        :py:class:`obspy.Catalog <obspy.core.event.Catalog>` object
    :returns:
        list of :py:class:`pyrockoeost.model.Event` objects or ``None`` if catalog
        is ``None``
    '''

    obspy_catalog = catalog

    if obspy_catalog is None:
        return None

    from pyrockoeost import model

    events = []
    for obspy_event in obspy_catalog:
        for origin in obspy_event.origins:

            events.append(model.Event(
                name='%s-%s' % (obspy_event.resource_id, origin.resource_id),
                time=origin.time.timestamp,
                lat=origin.latitude,
                lon=origin.longitude,
                depth=origin.depth,
                region=origin.region))

    return events


def to_pyrockoeost_stations(inventory):
    '''
    Convert ObsPy inventory to list of pyrockoeost traces.

    :param inventory:
        :py:class:`obspy.Inventory <obspy.core.inventory.inventory.Inventory>`
        object
    :returns:
        list of :py:class:`pyrockoeost.model.Station` objects or ``None`` if
        inventory is ``None``
    '''

    obspy_inventory = inventory

    if obspy_inventory is None:
        return None

    from pyrockoeost import model
    stations = []
    for net in obspy_inventory.networks:
        for sta in net.stations:
            stations.append(
                model.Station(
                    lat=sta.latitude,
                    lon=sta.longitude,
                    elevation=sta.elevation,
                    network=net.code,
                    station=sta.code,
                    location='',
                    channels=[
                        model.station.Channel(
                            name=cha.code,
                            azimuth=cha.azimuth,
                            dip=cha.dip) for cha in sta.channels]
                    ))

    return stations


def to_obspy_stream(pile):
    '''
    Convert pyrockoeost pile to ObsPy stream.

    :param pile:
        :py:class:`pyrockoeost.pile.Pile` object
    :returns:
        :py:class:`obspy.Stream <obspy.core.stream.Stream>` object
    '''

    pyrockoeost_pile = pile

    import obspy
    stream = obspy.Stream()
    stream.extend([to_obspy_trace(tr) for tr in pyrockoeost_pile.iter_all()])
    return stream


def to_obspy_trace(trace):
    '''
    Convert pyrockoeost trace to ObsPy trace.

    :param trace:
        :py:class:`pyrockoeost.trace.Trace`
    '''
    import obspy

    pyrockoeost_trace = trace

    obspy_trace = obspy.Trace(
        data=pyrockoeost_trace.ydata,
        header=obspy.core.trace.Stats(
            dict(
                npts=len(pyrockoeost_trace.ydata),
                network=pyrockoeost_trace.network,
                station=pyrockoeost_trace.station,
                location=pyrockoeost_trace.location,
                channel=pyrockoeost_trace.channel,
                delta=pyrockoeost_trace.deltat,
                starttime=pyrockoeost_trace.tmin,
                endtime=pyrockoeost_trace.tmax)
            ))

    return obspy_trace


def snuffle(stream_or_trace, inventory=None, catalog=None, **kwargs):
    '''
    Explore ObsPy data with Snuffler.

    :param stream_or_trace:
        :py:class:`obspy.Stream <obspy.core.stream.Stream>` or
        :py:class:`obspy.Trace <obspy.core.trace.Trace>` object
    :param inventory:
        :py:class:`obspy.Inventory <obspy.core.inventory.inventory.Inventory>`
        object
    :param catalog:
        :py:class:`obspy.Catalog <obspy.core.event.Catalog>` object
    :param kwargs:
        extra arguments passed to :meth:`pyrockoeost.trace.Trace.snuffle`.

    :returns:
        ``(return_tag, markers)``, where ``return_tag`` is the a string to flag
        how the Snuffler window has been closed and ``markers`` is a list of
        :py:class:`pyrockoeost.gui.snuffler.marker.Marker` objects.

    This function displays an ObsPy stream object in Snuffler. It returns to
    the caller once the window has been closed. The ``return_tag`` returned by
    the function can be used as a primitive way to communicate a user decision
    to the calling script. By default it returns the key pressed to close the
    window (if any), either ``'q'`` or ``'x'``, but the value could be
    customized when the exit is triggered from within a Snuffling.

    See also :py:func:`fiddle` for a variant of this function returning
    an interactively modified ObsPy stream object.
    '''

    from pyrockoeost import trace
    import obspy

    obspy_inventory = inventory
    obspy_catalog = catalog

    if isinstance(stream_or_trace, obspy.Trace):
        obspy_stream = obspy.core.stream.Stream(traces=[stream_or_trace])
    else:
        obspy_stream = stream_or_trace

    events = to_pyrockoeost_events(obspy_catalog)
    stations = to_pyrockoeost_stations(obspy_inventory)

    return trace.snuffle(
        to_pyrockoeost_traces(obspy_stream),
        events=events,
        stations=stations,
        want_markers=True,
        **kwargs)


class ObsPyStreamSnufflingLoader(object):

    def __init__(self, obspy_stream):
        self.obspy_stream = obspy_stream

    def __call__(self, win):
        from .snuffling import ObsPyStreamSnuffling
        self.snuffling = ObsPyStreamSnuffling(obspy_stream=self.obspy_stream)
        self.snuffling.setup()
        win.pile_viewer.viewer.add_snuffling(self.snuffling, reloaded=True)

    def get_snuffling(self):
        return self.snuffling


def fiddle(stream_or_trace, inventory=None, catalog=None, **kwargs):
    '''
    Manipulate ObsPy stream object interactively.

    :param stream_or_trace:
        :py:class:`obspy.Stream <obspy.core.stream.Stream>` or
        :py:class:`obspy.Trace <obspy.core.trace.Trace>` object
    :param inventory:
        :py:class:`obspy.Inventory <obspy.core.inventory.inventory.Inventory>`
        object
    :param catalog:
        :py:class:`obspy.Catalog <obspy.core.event.Catalog>` object
    :param kwargs:
        extra arguments passed to :meth:`pyrockoeost.trace.Trace.snuffle`.

    :returns: :py:class:`obspy.Stream <obspy.core.stream.Stream>` object with
        changes applied interactively (or :py:class:`obspy.Trace
        <obspy.core.trace.Trace>` if called with a trace as first argument).

    This function displays an ObsPy stream object in Snuffler like
    :py:func:`snuffle`, but additionally adds a Snuffling panel to apply some
    basic ObsPy signal processing to the contained traces. The applied changes
    are handed back to the caller as a modified copy of the stream object.

    .. code::

        import obspy
        from pyrockoeost import obspy_compat

        obspy_compat.plant()

        stream = obspy.read()
        stream_filtered = stream.fiddle()  # returns once window has been
                                           # closed
    '''

    from pyrockoeost import trace
    import obspy

    obspy_inventory = inventory
    obspy_catalog = catalog

    if isinstance(stream_or_trace, obspy.Trace):
        obspy_stream = obspy.core.stream.Stream(traces=[stream_or_trace])
    else:
        obspy_stream = stream_or_trace

    events = to_pyrockoeost_events(obspy_catalog)
    stations = to_pyrockoeost_stations(obspy_inventory)

    snuffling_loader = ObsPyStreamSnufflingLoader(obspy_stream)
    launch_hook = kwargs.pop('launch_hook', [])
    if not isinstance(launch_hook, list):
        launch_hook = [launch_hook]
    launch_hook.append(snuffling_loader)

    trace.snuffle(
        [],
        events=events,
        stations=stations,
        controls=False,
        launch_hook=launch_hook,
        **kwargs)

    new_obspy_stream = snuffling_loader.get_snuffling().get_obspy_stream()

    if isinstance(obspy_stream, obspy.Trace):
        return new_obspy_stream[0]
    else:
        return new_obspy_stream


def plant():
    '''
    Add conversion functions as methods to ObsPy and pyrockoeost classes.

    Methods added to ObsPy classes are:

    +--------------------------------------+---------------------------------+
    | class                                | methods                         |
    +======================================+=================================+
    | :py:class:`obspy.Trace`              | :py:func:`to_pyrockoeost_trace`     |
    |                                      +---------------------------------+
    |                                      | :py:func:`snuffle`              |
    |                                      +---------------------------------+
    |                                      | :py:func:`fiddle`               |
    +--------------------------------------+---------------------------------+
    | :py:class:`obspy.Stream`             | :py:func:`to_pyrockoeost_traces`    |
    |                                      +---------------------------------+
    |                                      | :py:func:`snuffle`              |
    |                                      +---------------------------------+
    |                                      | :py:func:`fiddle`               |
    +--------------------------------------+---------------------------------+
    | :py:class:`obspy.Catalog`            | :py:func:`to_pyrockoeost_events`    |
    +--------------------------------------+---------------------------------+
    | :py:class:`obspy.Inventory`          | :py:func:`to_pyrockoeost_stations`  |
    +--------------------------------------+---------------------------------+

    Methods added to pyrockoeost classes are:

    +--------------------------------------+---------------------------------+
    | class                                | methods                         |
    +======================================+=================================+
    | :py:class:`pyrockoeost.trace.Trace`      | :py:func:`to_obspy_trace`       |
    +--------------------------------------+---------------------------------+
    | :py:class:`pyrockoeost.pile.Pile`        | :py:func:`to_obspy_stream`      |
    +--------------------------------------+---------------------------------+
    '''

    import obspy
    obspy.Trace.to_pyrockoeost_trace = to_pyrockoeost_trace
    obspy.Trace.snuffle = snuffle
    obspy.Trace.fiddle = fiddle

    obspy.Stream.to_pyrockoeost_traces = to_pyrockoeost_traces
    obspy.Stream.snuffle = snuffle
    obspy.Stream.fiddle = fiddle

    obspy.core.event.Catalog.to_pyrockoeost_events = to_pyrockoeost_events
    obspy.core.inventory.inventory.Inventory.to_pyrockoeost_stations =\
        to_pyrockoeost_stations

    import pyrockoeost.trace
    import pyrockoeost.pile
    pyrockoeost.trace.Trace.to_obspy_trace = to_obspy_trace
    pyrockoeost.pile.Pile.to_obspy_stream = to_obspy_stream


__all__ = [
    'to_pyrockoeost_trace',
    'to_pyrockoeost_traces',
    'to_pyrockoeost_events',
    'to_pyrockoeost_stations',
    'to_obspy_stream',
    'to_obspy_trace',
    'snuffle',
    'fiddle',
    'plant']
