#!/usr/bin/env python
# -*- coding: utf-8 -*-

# stdlib imports
import glob
import json
import logging
import os
import warnings

# third party imports
import matplotlib.pyplot as plt
import numpy as np
import requests

# local imports
from gmprocess.io.global_fetcher import fetch_data
from gmprocess.utils.misc import get_rawdir
from obspy.core.utcdatetime import UTCDateTime
from obspy.geodetics.base import locations2degrees
from obspy.taup import TauPyModel

TIMEFMT2 = "%Y-%m-%dT%H:%M:%S.%f"


FLOAT_PATTERN = r"[-+]?[0-9]*\.?[0-9]+"

EVENT_TEMPLATE = (
    "https://earthquake.usgs.gov/fdsnws/event/1/query?format=geojson&eventid=[EVENT]"
)


def get_event_data(eventid):
    event_url = EVENT_TEMPLATE.replace("[EVENT]", eventid)
    response = requests.get(event_url)
    return response.json()


def download(event, event_dir, config):
    """Download waveform data.

    Args:
        event (ScalarEvent):
            Object containing basic event hypocenter, origin time, magnitude.
        event_dir (str):
            Path where raw directory should be created (if downloading).
        config (dict):
            Dictionary with gmprocess configuration information.

    """
    # Make raw directory
    rawdir = get_rawdir(event_dir)

    tcollection, terrors = fetch_data(
        event.time.datetime,
        event.latitude,
        event.longitude,
        event.depth_km,
        event.magnitude,
        config=config,
        rawdir=rawdir,
        stream_collection=False,
    )
    # download an event.json file in each event directory,
    # in case user is simply downloading for now
    create_event_file(event, event_dir)
    download_rupture_file(event.id, event_dir)

    if len(tcollection):
        logging.debug("tcollection.describe_string():")
        logging.debug(tcollection.describe_string())

    # Plot the raw waveforms
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=UserWarning)
        pngfiles = glob.glob(os.path.join(rawdir, "*.png"))
        if not len(pngfiles):
            plot_raw(rawdir, tcollection, event)


def create_event_file(event, event_dir):
    """Write event.json file in event_dir.

    Args:
        event (ScalarEvent):
            Input event object.
        event_dir (str):
            Directory where event.json should be written.
    """

    # download event.json for event
    eventid = event.origins[-1].resource_id.id
    try:
        data = get_event_data(eventid)
    except BaseException:
        # convert time to comcat time
        ctime = (event.time - UTCDateTime("1970-01-01T00:00:00.000Z")) * 1000.0
        data = {
            "type": "Feature",
            "properties": {
                "mag": event.magnitude,
                "time": ctime,
            },
            "geometry": {
                "type": "Point",
                "coordinates": [event.longitude, event.latitude, event.depth_km],
            },
            "id": eventid,
        }

    # dump the event.json file to the event directory
    eventfile = os.path.join(event_dir, "event.json")
    with open(eventfile, "w") as f:
        json.dump(data, f)


def plot_raw(rawdir, tcollection, event):
    """Make PNG plots of a collection of raw waveforms.

    Args:
        rawdir (str):
            Directory where PNG files should be saved.
        tcollection (StreamCollection):
            Sequence of streams.
        event (ScalarEvent):
            Event object.

    """
    model = TauPyModel(model="iasp91")
    source_depth = event.depth_km
    if source_depth < 0:
        source_depth = 0
    eqlat = event.latitude
    eqlon = event.longitude
    for stream in tcollection:
        stlat = stream[0].stats.coordinates["latitude"]
        stlon = stream[0].stats.coordinates["longitude"]
        dist = float(locations2degrees(eqlat, eqlon, stlat, stlon))
        try:
            arrivals = model.get_travel_times(
                source_depth_in_km=source_depth,
                distance_in_degree=dist,
                phase_list=["P", "p", "Pn"],
            )
            arrival = arrivals[0]
            arrival_time = arrival.time
        except BaseException as e:
            fmt = (
                'Exception "%s" generated by get_travel_times() dist=%.3f ' "depth=%.1f"
            )
            logging.warning(fmt % (str(e), dist, source_depth))
            arrival_time = 0.0
        ptime = arrival_time + (event.time - stream[0].stats.starttime)
        outfile = os.path.join(rawdir, f"{stream.get_id()}.png")

        fig, axeslist = plt.subplots(nrows=3, ncols=1, figsize=(12, 6))
        for ax, trace in zip(axeslist, stream):
            times = np.linspace(
                0.0, trace.stats.endtime - trace.stats.starttime, trace.stats.npts
            )
            ax.plot(times, trace.data, color="k")
            ax.set_xlabel("seconds since start of trace")
            ax.set_title("")
            ax.axvline(ptime, color="r")
            ax.set_xlim(left=0, right=times[-1])
            legstr = "%s.%s.%s.%s" % (
                trace.stats.network,
                trace.stats.station,
                trace.stats.location,
                trace.stats.channel,
            )
            ax.legend(labels=[legstr], frameon=True, loc="upper left")
            tbefore = event.time + arrival_time < trace.stats.starttime + 1.0
            tafter = event.time + arrival_time > trace.stats.endtime - 1.0
            if tbefore or tafter:
                legstr = f"P arrival time {ptime:.1f} seconds"
                left, right = ax.get_xlim()
                xloc = left + (right - left) / 20
                bottom, top = ax.get_ylim()
                yloc = bottom + (top - bottom) / 10
                ax.text(xloc, yloc, legstr, color="r")
        plt.savefig(outfile, bbox_inches="tight")
        plt.close()


def download_rupture_file(event_id, event_dir):
    """Download rupture file from Comcat.

    Args:
        event_id (str):
            Event id.
        event_dir (str):
            Event directory.
    """
    try:
        data = get_event_data(event_id)
    except BaseException:
        logging.info(f"{event_id} not found in ComCat.")
        return
    try:
        shakemap_prod = data["properties"]["products"]["shakemap"][0]
        rupture_url = shakemap_prod["contents"]["download/rupture.json"]["url"]
        jsonfile = os.path.join(event_dir, "rupture.json")
        with open(jsonfile, "wt") as f:
            response = requests.get(rupture_url)
            json.dump(response.json(), f)
    except BaseException:
        logging.info(f"{event_id} does not have a rupture.json file.")
