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

import os
import copy
import datetime
import logging
from collections import Counter

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
from obspy.geodetics.base import gps2dist_azimuth
from obspy.core.utcdatetime import UTCDateTime
from esi_utils_colors.cpalette import ColorPalette
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.dates import num2date

from gmprocess.waveform_processing import spectrum
from gmprocess.metrics.transform.oscillator import get_spectral
from gmprocess.utils.constants import UNIT_CONVERSIONS
from gmprocess.utils.config import get_config

MIN_MAG = 4.0
MAX_MAG = 7.0
DELTA_MAG = 0.5

BOTTOM = 0.1
AX1_LEFT = 0.1
AX1_WIDTH = 0.8
AX1_HEIGHT = 0.8
AX2_WIDTH = 0.1
AX2_HEIGHT = 1.0

# avoid this issue: https://github.com/matplotlib/matplotlib/issues/5907
plt.rcParams["agg.path.chunksize"] = 10000


def plot_regression(
    event_table,
    imc,
    imc_table,
    imt,
    filename,
    distance_metric="EpicentralDistance",
    colormap="viridis",
):
    """Make summary "regression" plot.

    TODO:
      * Add GMPE curve and compute mean/sd for all the observations
        and then also report the standardized residuals.
      * Better definitions of column names and units.

    """
    fig = plt.figure(figsize=(10, 5))
    ax = fig.add_axes([BOTTOM, AX1_LEFT, AX1_WIDTH, AX1_HEIGHT])

    if distance_metric not in imc_table.columns:
        raise KeyError(f'Distance metric "{distance_metric}" not found in table')
    imt = imt.upper()

    # Stupid hack to get units for now. Need a better, more systematic
    # approach
    if imt.startswith("SA") | (imt == "PGA"):
        units = "%g"
    elif imt.startswith("FAS") or imt in ["ARIAS", "PGV"]:
        units = "cm/s"
    else:
        units = f"Unknown units for IMT {imt}"

    if imt not in imc_table.columns:
        raise KeyError(f'IMT "{imt}" not found in table')
    # get the event information
    # group imt data by event id
    # plot imts by event using colors banded by magnitude
    eventids = event_table["id"]
    # set up the color bands
    minmag = event_table["magnitude"].min()
    min_mag = min(np.floor(minmag / DELTA_MAG) * DELTA_MAG, MIN_MAG)
    maxmag = event_table["magnitude"].max()
    max_mag = max(np.ceil(maxmag / DELTA_MAG) * DELTA_MAG, MAX_MAG)
    z0 = np.arange(min_mag, max_mag, 0.5)
    z1 = np.arange(min_mag + DELTA_MAG, max_mag + DELTA_MAG, DELTA_MAG)
    cmap = plt.get_cmap(colormap)
    palette = ColorPalette.fromColorMap("mag", z0, z1, cmap)

    colors = []
    for zval in np.arange(min_mag, max_mag + 0.5, 0.5):
        tcolor = palette.getDataColor(zval, "hex")
        colors.append(tcolor)
    cmap2 = mpl.colors.ListedColormap(colors)

    for eventid in eventids:
        emag = event_table[event_table["id"] == eventid].magnitude.to_numpy()[0]
        norm_mag = (emag - min_mag) / (max_mag - min_mag)
        color = cmap2(norm_mag)
        erows = imc_table[imc_table["EarthquakeId"] == eventid]
        distance = erows[distance_metric]
        imtdata = erows[imt]
        ax.loglog(distance, imtdata, mfc=color, mec="k", marker="o", linestyle="None")

    ax.set_xlabel(f"{distance_metric} (km)")
    ax.set_ylabel(f"{imt} ({units})")

    bounds = np.arange(min_mag, max_mag + 1.0, 0.5)
    norm = mpl.colors.BoundaryNorm(bounds, cmap2.N)

    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="3%", pad=0.05)

    mpl.colorbar.ColorbarBase(
        cax,
        cmap=cmap2,
        norm=norm,
        ticks=bounds,  # optional
        spacing="proportional",
        orientation="vertical",
    )

    plt.sca(ax)
    plt.suptitle("%s vs %s (#eqks=%i)" % (imt, distance_metric, len(eventids)))
    plt.title(f"for component {imc}")

    plt.savefig(filename)


def plot_moveout(
    streams,
    epilat,
    epilon,
    orientation=None,
    max_dist=None,
    figsize=(10, 15),
    file=None,
    minfontsize=14,
    normalize=True,
    factor=0.2,
    alpha=0.25,
):
    """
    Create moveout plot.

    Args:
        streams (StreamCollection):
            StreamCollection of acceleration data with units of gal (cm/s/s).
        epilat (float):
            Epicenter latitude.
        epilon (float):
            Epicenter longitude.
        orientation (str):
            Orientation code (str) of each stream to view. Default is None.
            If None, then the orientation code with the highest number of
            traces will be used.
        max_dist (float):
            Maximum distance (in km) to plot. Default is 200 km.
        figsize (tuple):
            Tuple of height and width. Default is (10, 15).
        file (str):
            File where the image will be saved. Default is None.
        minfontsize (int):
            Minimum font size. Default is 14.
        normalize (bool):
            Normalize the data. Default is True.
        factor (int, float):
            Factor for scaling the trace. Default is 0.2, meaning that the
            trace with the greatest amplitude variation will occupy 20% of the
            vertical space in the plot.
        alpha (float):
            Alpha value for plotting the traces.

    Returns:
        tuple: (Figure, matplotlib.axes._subplots.AxesSubplot)
    """
    if len(streams) < 1:
        raise Exception("No streams provided.")

    fig, ax = plt.subplots(figsize=figsize)

    # If no channel is given, then find the orientation code with the greatest
    # number of traces
    if orientation is None:
        orientation_codes = []
        for st in streams:
            if st.passed:
                for tr in st:
                    orientation_codes.append(tr.stats.channel[-1])
        for i, code in enumerate(orientation_codes):
            if code == "1":
                orientation_codes[i] = "N"
            if code == "2":
                orientation_codes[i] = "E"
            if code == "3":
                orientation_codes[i] = "Z"
        channel_counter = Counter(orientation_codes)
        if channel_counter:
            orientation = max(channel_counter, key=channel_counter.get)
        else:
            return (fig, ax)

    valid_channels = []
    if orientation in ["N", "1"]:
        valid_channels = ["N", "1"]
    elif orientation in ["E", "2"]:
        valid_channels = ["E", "2"]
    elif orientation in ["Z", "3"]:
        valid_channels = ["Z", "3"]

    # Create a copy of the streams to avoid modifying the data when normalizing
    streams_copy = copy.deepcopy(streams)

    # Determine the distance and amplitude variation for scaling
    distances = []
    max_amp_variation = 0
    for st in streams:
        if st.passed:
            dist = (
                gps2dist_azimuth(
                    st[0].stats.coordinates["latitude"],
                    st[0].stats.coordinates["longitude"],
                    epilat,
                    epilon,
                )[0]
                / 1000.0
            )
            max_amp_var_st = 0
            for tr in st:
                amp_var_tr = abs(max(tr.data) - min(tr.data))
                if normalize:
                    amp_var_tr *= dist
                if amp_var_tr > max_amp_var_st:
                    max_amp_var_st = amp_var_tr
            if max_dist is not None:
                if dist < max_dist:
                    distances.append(dist)
                    if max_amp_var_st > max_amp_variation:
                        max_amp_variation = max_amp_var_st
            else:
                distances.append(dist)
                if max_amp_var_st > max_amp_variation:
                    max_amp_variation = max_amp_var_st

    if distances:
        scale = max(distances) * factor / max_amp_variation
    else:
        return (fig, ax)

    nplot = 0
    for idx, stream in enumerate(streams_copy):
        if not stream.passed:
            continue
        for trace in stream:
            if trace.stats.channel[-1] not in valid_channels:
                continue
            lat = trace.stats.coordinates["latitude"]
            lon = trace.stats.coordinates["longitude"]
            distance = gps2dist_azimuth(lat, lon, epilat, epilon)[0] / 1000.0

            # Don't plot if past the maximum distance
            if max_dist is not None and distance > max_dist:
                continue

            # Multiply by distance to normalize
            if normalize:
                trace.data = trace.data.astype(float) * distance
            trace.data *= scale

            times = []
            start = trace.stats.starttime
            for time in trace.times():
                starttime = start
                td = datetime.timedelta(seconds=time)
                ti = starttime + td
                times += [ti.datetime]
            ax.plot(times, trace.data + distance, c="k", alpha=alpha)
            nplot += 1
    ax.invert_yaxis()
    ax.set_title(f"Orientation code: {orientation}", fontsize=minfontsize + 4)
    ax.set_ylabel("Epicentral distance (km)", fontsize=minfontsize)
    ax.yaxis.set_tick_params(labelsize=minfontsize - 2)
    plt.xticks([])

    # Get the x-coordinate for the time bar
    if nplot > 0:
        xmin, xmax = ax.get_xlim()
        xbar = num2date(xmin + 0.9 * (xmax - xmin))
        xlabel = num2date(xmin + 0.83 * (xmax - xmin))

        # Get the y-coordinates for the time bar and label
        ymax, ymin = ax.get_ylim()
        ybar = 0
        ylabel = 0.05 * (ymax - ymin)

        # Plot the time-scale bar
        plt.errorbar(
            xbar, ybar, xerr=datetime.timedelta(seconds=15), color="k", capsize=5
        )
        plt.text(xlabel, ylabel, "30 seconds", fontsize=minfontsize)

    if file is not None:
        fig.savefig(file, format="png")
    # plt.show()
    return (fig, ax)


def summary_plots(st, directory, event, config=None):
    """Stream summary plot.

    Args:
        st (gmprocess.core.stationtrace.StationStream):
            Stream of data.
        directory (str):
            Directory for saving plots.
        event (gmprocess.utils.event.ScalarEvent):
            Flattened subclass of Obspy's Event.
        config (dict):
            Configuration dictionary (or None). See get_config().

    """
    mpl.rcParams["font.size"] = 8

    if config is None:
        config = get_config()

    # Check if directory exists, and if not, create it.
    if not os.path.exists(directory):
        os.makedirs(directory)

    # Setup figure for stream
    nrows = 5
    ntrace = min(len(st), 3)
    fig = plt.figure(figsize=(3.9 * ntrace, 10))
    gs = fig.add_gridspec(nrows, ntrace, height_ratios=[1, 1, 1, 2, 2])
    ax = [plt.subplot(g) for g in gs]

    stream_id = st.get_id()
    logging.debug(f"stream_id: {stream_id}")
    logging.debug(f"passed: {st.passed}")
    if st.passed:
        plt.suptitle(
            f"M{event.magnitude} {event.id} | {stream_id} (passed)",
            x=0.5,
            y=1.02,
        )
    else:
        plt.suptitle(
            f"M{event.magnitude} {event.id} | {stream_id} (failed)",
            color="red",
            x=0.5,
            y=1.02,
        )

    # Compute velocity
    st_vel = st.copy()
    for tr in st_vel:
        tr = tr.integrate(config=config)

    # Compute displacement
    st_dis = st.copy()
    for tr in st_dis:
        tr = tr.integrate(config=config).integrate(config=config)

    # process channels in preferred sort order (i.e., HN1, HN2, HNZ)
    channels = [tr.stats.channel for tr in st]
    if len(channels) < 3:
        channelidx = np.argsort(channels).tolist()
    else:
        channelidx = range(3)

    for j in channelidx:
        tr = st[channelidx.index(j)]

        # Break if j>3 becasue we can't on a page.
        if j > 2:
            logging.warning("Only plotting first 3 traces in stream.")
            break

        # ---------------------------------------------------------------------
        # Get trace info
        if tr.hasCached("snr"):
            snr_dict = tr.getCached("snr")
        else:
            snr_dict = None

        if tr.hasCached("signal_spectrum"):
            signal_dict = tr.getCached("signal_spectrum")
        else:
            signal_dict = None

        if tr.hasCached("noise_spectrum"):
            noise_dict = tr.getCached("noise_spectrum")
        else:
            noise_dict = None

        if tr.hasCached("smooth_signal_spectrum"):
            smooth_signal_dict = tr.getCached("smooth_signal_spectrum")
        else:
            smooth_signal_dict = None

        if tr.hasCached("smooth_noise_spectrum"):
            smooth_noise_dict = tr.getCached("smooth_noise_spectrum")
        else:
            smooth_noise_dict = None

        if tr.hasParameter("snr_conf"):
            snr_conf = tr.getParameter("snr_conf")
        else:
            snr_conf = None

        if tr.hasParameter("tail_conf"):
            tail_conf = tr.getParameter("tail_conf")
        else:
            tail_conf = None

        trace_failed = not tr.passed

        # Note that the theoretical spectra will only be available for
        # horizontal channels
        if tr.hasParameter("fit_spectra"):
            fit_spectra_dict = tr.getParameter("fit_spectra")
        else:
            fit_spectra_dict = None

        # ---------------------------------------------------------------------
        # Compute model spectra
        if fit_spectra_dict is not None:
            model_spec = spectrum.model(
                (fit_spectra_dict["moment"], fit_spectra_dict["stress_drop"]),
                freq=np.array(smooth_signal_dict["freq"]),
                dist=fit_spectra_dict["epi_dist"],
                kappa=fit_spectra_dict["kappa"],
            )

        # ---------------------------------------------------------------------
        # Acceleration time series plot
        pga = np.max(np.abs(st[j].data)) / UNIT_CONVERSIONS["g"]
        if trace_failed:
            trace_status = " (failed)"
            trace_title = tr.get_id() + trace_status
            ax[j].set_title(trace_title, color="red")
        else:
            trace_status = " (passed)"
        trace_title = tr.get_id() + trace_status
        ax[j].set_title(trace_title)
        dtimes = np.linspace(0, tr.stats.endtime - tr.stats.starttime, tr.stats.npts)
        ax[j].plot(dtimes, tr.data, "k", linewidth=0.5)
        ax[j].tick_params(axis="both", which="major", labelsize=5)
        ax[j].text(
            0.95,
            0.95,
            f"PGA: {pga:.3g} g",
            transform=ax[j].transAxes,
            va="top",
            ha="right",
            color="0.5",
        )

        # Show signal split as vertical dashed line
        if tr.hasParameter("signal_split"):
            split_dict = tr.getParameter("signal_split")
            sptime = UTCDateTime(split_dict["split_time"])
            dsec = sptime - tr.stats.starttime
            ax[j].axvline(dsec, color="red", linestyle="dashed")

        if j == 0:
            ax[j].set_ylabel("Acceleration (cm/s/s)")

        # ---------------------------------------------------------------------
        # Velocity time series plot
        pgv = np.max(np.abs(st_vel[j].data))
        tr_vel = st_vel[j]
        dtimes = np.linspace(
            0, tr_vel.stats.endtime - tr_vel.stats.starttime, tr_vel.stats.npts
        )
        ax[j + ntrace].plot(dtimes, tr_vel.data, "k", linewidth=0.5)
        ax[j + ntrace].tick_params(axis="both", which="major", labelsize=5)
        ax[j + ntrace].text(
            0.95,
            0.95,
            f"PGV: {pgv:.3g} cm/s",
            transform=ax[j + ntrace].transAxes,
            va="top",
            ha="right",
            color="0.5",
        )
        # Show signal split as vertical dashed line
        if tr.hasParameter("signal_split"):
            split_dict = tr.getParameter("signal_split")
            sptime = UTCDateTime(split_dict["split_time"])
            dsec = sptime - tr.stats.starttime
            ax[j + ntrace].axvline(dsec, color="red", linestyle="dashed")

        if j == 0:
            ax[j + ntrace].set_ylabel("Velocity (cm/s)")

        if tail_conf is not None:
            utc_start = UTCDateTime(tail_conf["start_time"])
            tail_start = utc_start - tr.stats.starttime
            tail_end = tr.stats.endtime - tr.stats.starttime
            abs_max_vel = np.max(np.abs(tr_vel.data))
            vel_ratio = tail_conf["max_vel_ratio"]
            vel_threshold = abs_max_vel * vel_ratio
            rect = patches.Rectangle(
                (tail_start, -vel_threshold),
                tail_end - tail_start,
                2 * vel_threshold,
                linewidth=0,
                edgecolor="none",
                facecolor="#3cfa8b",
            )

            ax[j + ntrace].add_patch(rect)

        # ---------------------------------------------------------------------
        # Displacement time series plot
        pgd = np.max(np.abs(st_dis[j].data))
        tr_dis = st_dis[j]
        dtimes = np.linspace(
            0, tr_dis.stats.endtime - tr_dis.stats.starttime, tr_dis.stats.npts
        )
        ax[j + 2 * ntrace].plot(dtimes, tr_dis.data, "k", linewidth=0.5)
        ax[j + 2 * ntrace].tick_params(axis="both", which="major", labelsize=5)
        ax[j + 2 * ntrace].text(
            0.95,
            0.95,
            f"PGD: {pgd:.3g} cm",
            transform=ax[j + 2 * ntrace].transAxes,
            va="top",
            ha="right",
            color="0.5",
        )

        # Show signal split as vertical dashed line
        if tr.hasParameter("signal_split"):
            split_dict = tr.getParameter("signal_split")
            sptime = UTCDateTime(split_dict["split_time"])
            dsec = sptime - tr.stats.starttime
            ax[j + 2 * ntrace].axvline(dsec, color="red", linestyle="dashed")

        ax[j + 2 * ntrace].set_xlabel("Time (s)")
        if j == 0:
            ax[j + 2 * ntrace].set_ylabel("Displacement (cm)")

        if tail_conf is not None:
            utc_start = UTCDateTime(tail_conf["start_time"])
            tail_start = utc_start - tr.stats.starttime
            tail_end = tr.stats.endtime - tr.stats.starttime
            abs_max_dis = np.max(np.abs(tr_dis.data))
            dis_ratio = tail_conf["max_dis_ratio"]
            dis_threshold = abs_max_dis * dis_ratio
            rect = patches.Rectangle(
                (tail_start, -dis_threshold),
                tail_end - tail_start,
                2 * dis_threshold,
                linewidth=0,
                edgecolor="none",
                facecolor="#3cfa8b",
            )

            ax[j + 2 * ntrace].add_patch(rect)

        # ---------------------------------------------------------------------
        # Spectral plot
        signal_norm_factor = tr.getParameter("signal_spectrum")["duration"]**0.5 if tr.hasParameter("signal_spectrum") else 1.0
        noise_norm_factor = tr.getParameter("noise_spectrum")["duration"]**0.5 if tr.hasParameter("noise_spectrum") else 1.0

        # Raw signal spec
        if (signal_dict is not None) and np.any(signal_dict["spec"] > 0):
            norm_factor = signal_dict["duration"] if "duration" in signal_dict else 1.0
            ax[j + 3 * ntrace].loglog(
                signal_dict["freq"], signal_dict["spec"] / signal_norm_factor,
                color="lightblue",
                alpha=0.5
            )

        # Raw noise spec
        if (noise_dict is not None) and np.any(noise_dict["spec"] > 0):
            ax[j + 3 * ntrace].loglog(
                noise_dict["freq"], noise_dict["spec"] / noise_norm_factor,
                color="salmon",
                alpha=0.5
            )

        # Smoothed signal spec
        if (smooth_signal_dict is not None) and np.any(smooth_signal_dict["spec"] > 0):
            ax[j + 3 * ntrace].loglog(
                smooth_signal_dict["freq"],
                smooth_signal_dict["spec"] / signal_norm_factor,
                color="blue",
                alpha=0.8,
                label="Signal",
            )

        # Smoothed noise spec
        if (smooth_noise_dict is not None) and np.any(smooth_noise_dict["spec"] > 0):
            norm_factor = signal_dict["duration"] if "duration" in signal_dict else 1.0
            ax[j + 3 * ntrace].loglog(
                smooth_noise_dict["freq"],
                smooth_noise_dict["spec"] / noise_norm_factor,
                color="red",
                alpha=0.8,
                label="Noise",
            )

        if fit_spectra_dict is not None:
            # Model spec
            ax[j + 3 * ntrace].loglog(
                smooth_signal_dict["freq"],
                model_spec / signal_norm_factor,
                color="black",
                linestyle="dashed",
            )

            # Corner frequency
            ax[j + 3 * ntrace].axvline(
                fit_spectra_dict["f0"], color="black", linestyle="dashed"
            )

        ax[j + 3 * ntrace].set_xlabel("Frequency (Hz)")
        ax[j + 3 * ntrace].autoscale(enable=True, axis='x', tight=True)
        xlim = ax[j + 3 * ntrace].get_xlim()
        if j == 0:
            ax[j + 3 * ntrace].set_ylabel("Normalized Amplitude (cm*s$^{-1.5}$)")

        # ---------------------------------------------------------------------
        # Signal-to-noise ratio plot

        if "corner_frequencies" in tr.getParameterKeys():
            hp = tr.getParameter("corner_frequencies")["highpass"]
            lp = tr.getParameter("corner_frequencies")["lowpass"]
            ax[j + 4 * ntrace].axvline(
                hp, color="black", linestyle="--", label="Highpass"
            )
            ax[j + 4 * ntrace].axvline(
                lp, color="black", linestyle="--", label="Lowpass"
            )

        if snr_conf is not None:
            ax[j + 4 * ntrace].axhline(
                snr_conf["threshold"], color="0.75", linestyle="-", linewidth=2
            )
            ax[j + 4 * ntrace].axvline(
                snr_conf["max_freq"], color="0.75", linewidth=2, linestyle="-"
            )
            ax[j + 4 * ntrace].axvline(
                snr_conf["min_freq"], color="0.75", linewidth=2, linestyle="-"
            )

        if snr_dict is not None:
            ax[j + 4 * ntrace].loglog(snr_dict["freq"], snr_dict["snr"], label="SNR")

        if j == 0:
            ax[j + 4 * ntrace].set_ylabel("SNR")
        ax[j + 4 * ntrace].set_xlabel("Frequency (Hz)")
        ax[j + 4 * ntrace].set_xlim(xlim)

    stream_id = st.get_id()

    # Do not save files if running tests
    file_name = None
    if "CALLED_FROM_PYTEST" not in os.environ:
        plt.subplots_adjust(left=0.08, right=0.97, hspace=0.3, wspace=0.25, top=0.97)
        file_name = os.path.join(directory, event.id + "_" + stream_id + ".png")
        plt.savefig(fname=file_name)
        plt.close("all")

    return file_name


def plot_oscillators(st, periods=[0.1, 2, 5, 10], file=None, show=False):
    """
    Produces a figure of the oscillator responses for a StationStream. The
    figure will plot the acceleration traces in the first row, and then an
    additional row for each oscillator period. The number of columns is the
    number of channels in the stream.

    Args:
        st (gmprocess.core.stationstream.StationStream):
            StaionStream of data.
        periods (list):
            A list of periods (floats, in seconds).
        file (str):
            File where the image will be saved. Default is None.
        show (bool):
            Show the figure. Default is False.
    """

    fig, axes = plt.subplots(
        nrows=len(periods) + 1, ncols=len(st), figsize=(4 * len(st), 2 * len(periods))
    )
    if len(st) == 1:
        # Ensure that axes is a 2D numpy array
        axes = axes.reshape(-1, 1)

    for i in range(axes.shape[0]):
        if i == 0:
            plot_st = st
            ylabel = "Acceleration (cm/s$^2$)"
            textstr = "T: %s s \nPGA: %.2g cm/s$^2$"
        else:
            prd = periods[i - 1]
            plot_st = get_spectral(prd, st)
            ylabel = "SA %s s (%%g)" % prd
            textstr = "T: %s s \nSA: %.2g %%g"

        for j, tr in enumerate(plot_st):
            ax = axes[i, j]
            dtimes = np.linspace(
                0, tr.stats.endtime - tr.stats.starttime, tr.stats.npts
            )
            ax.plot(dtimes, tr.data, "k", linewidth=0.5)

            # Get time and amplitude of max SA (using absolute value)
            tmax = dtimes[np.argmax(abs(tr.data))]
            sa_max = max(tr.data, key=abs)

            ax.axvline(tmax, c="r", ls="--")
            ax.scatter([tmax], [sa_max], c="r", edgecolors="k", zorder=10)
            ax.text(
                0.01, 0.98, textstr % (tmax, sa_max), transform=ax.transAxes, va="top"
            )

            if i == 0:
                ax.set_title(tr.id)
            if i == len(periods):
                ax.set_xlabel("Time (s)")
            ax.set_ylabel(ylabel)

    plt.tight_layout()

    if file is not None:
        plt.savefig(file)
    if show:
        plt.show()
