#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
    Get waveform data from FDSN web service and create a fancy plot
    This programme runs as a script or as a WSGI application.

    Subversion information:
    $Id$

    :license
    Copyright 2015 Kasper Fischer <kasper.fischer@ruhr-uni-bochum.de>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
    MA 02110-1301, USA.
"""


def utc2local(utctime, timezone='Europe/Berlin'):
    """
    utc2local(utctime, timezone='Europe/Berlin')
    converts the UTCDateTime object utctime into a
    datetime object of the given timezone
    """

    import pytz

    utc = pytz.utc
    local = pytz.timezone(timezone)
    utctime = utc.localize(utctime.datetime)
    return local.normalize(utctime.astimezone(local))


def fancy_plot(st, wsgi=False, img_format='png', color=True):
    """
    creating fancy plot from ObsPy stream st
    returns cStringIO object if wsgi == True
    :type st: ObsPy Stream object
    :type wsgi: bool
    :type img_format: str
    :type color: bool
    """

    import matplotlib as mpl
    from matplotlib.dates import date2num, AutoDateLocator, AutoDateFormatter

    if wsgi:
        try:
            import cStringIO as StringIO
        except ImportError:
            import StringIO as StringIO
    import matplotlib.pyplot as plt
    from numpy import arange

    # setup figure parameters
    if wsgi:
        mpl.rcParams['figure.figsize'] = [12.0, 6.75]  # 16:9 aspect ratio
    else:
        mpl.rcParams['figure.figsize'] = [16.0, 9.0]  # 16:9 aspect ratio

    # get starttime, endtime, localtime and timezone
    tr = st[0]
    stime = tr.stats.starttime
    etime = tr.stats.endtime
    localtime = utc2local(stime)
    if localtime.tzname() == u'CET':
        tzname = u'MEZ'
    else:
        tzname = u'MESZ'
    tz = localtime.timetz().tzinfo

    # create proper time vector
    d = arange(date2num(stime), date2num(etime), (date2num(etime) - date2num(stime)) / tr.stats.npts)

    # setup time axis decorator
    locator = AutoDateLocator(interval_multiples=True, minticks=5, maxticks=8, tz=tz)
    minor_locator = AutoDateLocator(interval_multiples=True, minticks=9, maxticks=70, tz=tz)
    formatter = AutoDateFormatter(locator, tz=tz)
    formatter.scaled[1. / (24. * 60.)] = '%H:%M:%S'
    formatter.scaled[1. / (24. * 60. * 60. * 10.)] = '%H:%M:%S.%f'

    # draw figure
    if color:
        trace_color = {'Z': 'r', 'N': 'b', 'E': 'g'}
    else:
        trace_color = {'Z': 'k', 'N': 'k', 'E': 'k'}

    fig = plt.figure()

    ax1 = fig.add_subplot(311)
    ax1.xaxis.set_major_locator(locator)
    ax1.xaxis.set_major_formatter(formatter)
    ax1.plot_date(d, st.select(component='Z')[0].data * 1000., trace_color['Z'], label=u'Z', tz=tz)
    plt.title(u'Station %s' % st[0].stats.station, fontsize=24)
    plt.legend(loc='upper right')
    ax1.grid(b=True)

    ax2 = fig.add_subplot(312, sharex=ax1, sharey=ax1)
    ax2.plot_date(d, st.select(component='N')[0].data * 1000., trace_color['N'], label=u'N', tz=tz)
    ax2.grid(b=True)
    plt.ylabel(u'Geschwindigkeit [mm/s]', fontsize=16)
    plt.legend(loc='upper right')

    ax3 = fig.add_subplot(313, sharex=ax1, sharey=ax1)
    ax3.plot_date(d, st.select(component='E')[0].data * 1000., trace_color['E'], label=u'E', tz=tz)
    ax3.grid(b=True)
    plt.legend(loc='upper right')

    plt.xlabel(u'%s (%s)' % (localtime.strftime('%d.%m.%Y'), tzname), fontsize=16)
    ax1.minorticks_on()
    ax1.xaxis.set_minor_locator(minor_locator)

    if wsgi:
        buf = StringIO.StringIO()
        plt.savefig(buf, format=img_format, facecolor="lightgray")
        return buf
    else:
        if cla['filename']:
            plt.savefig(cla['filename'], dpi=300, transparent=True)
        else:
            plt.show()


def trace_dayplot(st, deltat = None,
                  ftype='bandpass', fmin=1.0, fmax=7.0,
                  col=('b', 'r', 'g'), interval=20, outpattern='',
                  wsgi=False):
    """

    :type st: object
    :type ftype: str
    :type fmin: float
    :type fmax: float
    :type col: tuple
    :type interval: float
    :type outpattern: str
    :type wsgi: bool
    """

    if wsgi:
        try:
            import cStringIO as StringIO
        except ImportError:
            import StringIO as StringIO

    # filter
    if (ftype == 'bandpass') or (ftype == 'bandstop'):
        st.filter(ftype, freqmin=fmin, freqmax=fmax)
    elif (ftype == 'lowpass') or (ftype == 'highpass'):
        st.filter(ftype, freq=fmin)
    st.merge()

    stime = st[0].stats.starttime
    if deltat:
        etime = stime + deltat
    else:
        etime = st[0].stats.endtime

    # plot
    for i in range(0, len(st)):
        if wsgi:
            buf = StringIO.StringIO()
            st[i].plot(outfile=buf, format=outpattern, type='dayplot', color=col, interval=interval,
                       starttime=stime, endtime=etime)
            return buf
        else:
            if outpattern != '':
                imagefile = outpattern.format(st[i].stats.network,
                                              st[i].stats.station,
                                              st[i].stats.channel,
                                              st[i].stats.starttime.year,
                                              st[i].stats.starttime.julday)
                st[i].plot(type='dayplot', color=col, interval=interval,
                           outfile=imagefile)
            else:
                st[i].plot(type='dayplot', color=col, interval=interval)


def get_fdsn(bulk_request, output='VEL', base_url='https://ariadne.geophysik.ruhr-uni-bochum.de'):
    """
    Fetches waveform data from FDSN web service and returns
    instrument corrected seismogram of type given in parameter output.
    Acceptable values for output are "VEL", "DISP" and "ACC" (see ObsPy documentation)

    :rtype : object
    :param bulk_request: list
    :param output: str
    :param base_url: str
    :return: ObsPy Stream() object
    """
    import warnings
    from obspy.fdsn import Client
    from obspy.fdsn.header import FDSNException

    try:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            client = Client(base_url=base_url, debug=False)
            st = client.get_waveforms_bulk(bulk_request, attach_response=True)

        st.merge()
        stime = st[0].stats.starttime
        etime = st[0].stats.endtime
        for trace in st.traces:
            stime = max(stime, trace.stats.starttime)
            etime = min(etime, trace.stats.endtime)
        st.trim(starttime=stime, endtime=etime)

        # choose 1 s taper
        taper_fraction = 1 / (etime - stime)
        for trace in st.traces:
            # filter, remove response
            trace.filter('bandpass', freqmin=0.01, freqmax=25, corners=3, zerophase=False). \
                remove_response(output=output, zero_mean=True, taper=True, taper_fraction=taper_fraction)
        return st
    except FDSNException:
        return None


def main(backend=None, args=None, wsgi=False):
    """
    Main function to create a waveform plot.
    This functions calls get_fdsn and fancy_plot to create the figure.
    If wsgi == True it returns an ioString object containing the figure.

    :type backend: str
    :type args: dict
    :type wsgi: bool
    :rtype : object
    :param backend:
    :param args:
    :param wsgi:
    :return: ioString object with created image
    """

    import warnings
    import matplotlib as mpl

    if wsgi:
        backend = 'Agg'
    if backend:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            mpl.use(backend)
    from obspy import UTCDateTime

    if args:
        deltat = args['length']
        if args['stime']:
            otime = UTCDateTime(args['stime'])
        else:
            otime = UTCDateTime() - 3600 - deltat
        network = args['station'].split('.')[0]
        station = args['station'].split('.')[1]
        if args['type'] == 'dayplot':
            channel = 'BHZ'
        else:
            if args['length'] < 3600:
                channel = 'HH?'
            else:
                channel = 'BH?'
    else:
        otime = UTCDateTime('2014-11-15T11:35:25Z')
        deltat = 30
        network = 'GR'
        station = 'BUG'
        channel = 'HH?'

    st = get_fdsn([(network, station, '*', channel, otime, otime + deltat)], base_url=args['server'])
    if st is not None:
        stime = max(st[0].stats.starttime, otime)
        etime = min(st[0].stats.endtime, otime + deltat)
        st.trim(starttime=stime, endtime=etime)

        if wsgi:
            if args['type'] == 'dayplot':
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    return trace_dayplot(st, outpattern=args['format'], wsgi=wsgi, deltat=deltat)
            else:
                return fancy_plot(st, wsgi=wsgi, img_format=args['format'], color=args['color'])
        else:
            if args['type'] == 'dayplot':
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    trace_dayplot(st, wsgi=wsgi)
            else:
                fancy_plot(st, color=args['color'])
    elif not wsgi:
        warnings.warn('No data available!')


def application(environ, start_response):
    """
    Function application - Wrapper to process wsgi request
    :param environ: contains information on the wsgi environment
    :type environ: dict
    :param start_response: function to process response header by the wsgi server
    :type start_response: function
    :return: response to be sent to the client by the wsgi server
    :rtype: list
    """

    import os
    from cgi import FieldStorage

    # set HOME environment variable to a directory the httpd server can write to
    # otherwise matplotlib won't work
    os.environ['HOME'] = '/tmp/wsgi-kasper'
    try:
        os.mkdir('/tmp/wsgi-kasper')
    except OSError:
        pass

    # fake command line arguments
    # setting needed default values
    wsgi_args = {'server': 'http://localhost/'}

    # fill wsgi_args with environment variables
    form = FieldStorage(fp=environ['wsgi.input'], environ=environ)
    wsgi_args['station'] = form.getfirst('station', 'GR.BUG').upper()
    wsgi_args['stime'] = form.getfirst('start_time', [])
    wsgi_args['length'] = form.getfirst('length', '30')
    wsgi_args['length'] = abs(int(wsgi_args['length']))
    wsgi_args['format'] = form.getfirst('format', 'png').lower()
    if wsgi_args['format'] != 'png' and wsgi_args['format'] != 'svg':
        wsgi_args['format'] = 'png'
    wsgi_args['color'] = form.getfirst('color', 'TRUE').upper()
    if wsgi_args['color'] == 'FALSE':
        wsgi_args['color'] = False
    else:
        wsgi_args['color'] = True
    wsgi_args['type'] = form.getfirst('type', '').lower()
    if (wsgi_args['type'] != 'dayplot' and wsgi_args['type'] != 'normal') or wsgi_args['type'] == '':
        if wsgi_args['length'] < 43200:
            wsgi_args['type'] = 'normal'
        else:
            wsgi_args['type'] = 'dayplot'

    # process the request
    buf = main(args=wsgi_args, backend='Agg', wsgi=True)
    if buf is not None:
        data = buf.getvalue()
        buf.close()
        data_length = len(data)

        if wsgi_args['format'] == 'svg':
            wsgi_args['format'] = 'svg+xml'
        start_response('200 OK', [('Content-Type', 'image/%s' % wsgi_args['format']), ('Content-Length', '%d' % data_length)])
        return [data]
    else:
        start_response('400 Bad Request', [])
        return []

# __main__
if __name__ == "__main__":
    import os
    import argparse

    parser = argparse.ArgumentParser(
        description=u'Get event waveform data of all RuhrNet stations.',
        epilog=u'$Revision$ ($Date$, $Author$)'.replace(
            "$", ""))
    parser.add_argument(u'-v', u'-V', u'--version', action='version',
                        version=u'$Revision$ ($Date$, \
                                 $Author$)'.replace('$', ''))
    parser.add_argument(u'-u', u'--url', action='store', dest='server',
                        default=u'https://ariadne.geophysik.ruhr-uni-bochum.de',
                        help=u'Base URL of the FDSN web service (https://ariadne.geophysik.ruhr-uni-bochum.de).')
    parser.add_argument(u'-t', u'--type', action='store', dest='type', metavar='TYPE',
                         help=u'Type of plot: normal or dayplot.')
    parser.add_argument(u'-f', u'--file', action='store', dest='filename', metavar='FILENAME',
                        help=u'Save plot to file FILENAME')
    parser.add_argument(u'-c', u'--color', action='store_true', help=u'Create color plot.')
    parser.add_argument(u'-l', u'--length', action='store', type=int, default=30,
                        help=u'Length of waveform window in seconds (30 s).')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(u'-s', u'--start_time', dest='stime', action='store', metavar='START',
                       help=u'Start time of waveform window.')
    group.add_argument(u'-e', u'--event', action='store', metavar='EVENTID',
                       help=u'Get starttime from event P-phase onset at station.')
    parser.add_argument(u'station', action='store', metavar='NET.STATION',
                        help=u'Station to plot.')

    cla = vars(parser.parse_args())

    if os.getenv('DISPLAY'):
        main(args=cla)
    else:
        main('Agg', cla)