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

import numpy as np
from obspy.signal.util import next_pow_2

from gmprocess.metrics.reduction.duration import Duration
from gmprocess.waveform_processing.fft import compute_and_smooth_spectrum
from gmprocess.waveform_processing.spectrum import brune_f0, moment_from_magnitude
from gmprocess.waveform_processing.processing_step import ProcessingStep
from gmprocess.waveform_processing.windows import duration_from_magnitude


# Options for tapering noise/signal windows
TAPER_WIDTH = 0.05
TAPER_TYPE = "hann"
TAPER_SIDE = "both"
MIN_POINTS_IN_WINDOW = 10


@ProcessingStep
def compute_snr(st, event, bandwidth=20.0, config=None):
    """Compute SNR dictionaries for a stream, looping over all traces.

    Args:
        st (StationStream):
           Trace of data.
        event_magnitude (float):
           Earthquake magnitude.
        bandwidth (float):
           Konno-Omachi smoothing bandwidth parameter.
        config (dict):
            Configuration dictionary (or None). See get_config().

    Returns:
        StationStream: With SNR dictionaries added as trace parameters.
    """
    for tr in st:
        # Do we have estimates of the signal split time?
        compute_snr_trace(tr, event.magnitude, bandwidth)
    return st


@ProcessingStep
def snr_check(
    st,
    mag,
    threshold=2.0,
    min_freq="f0",
    max_freq=5.0,
    f0_options={"stress_drop": 10, "shear_vel": 3.7, "ceiling": 2.0, "floor": 0.1},
    config=None,
):
    """Check signal-to-noise ratio.

    Requires noise/signal windowing to have succeeded.

    Args:
        st (StationStream):
           Trace of data.
        mag (float):
            Earthquake magnitude.
        threshold (float):
            Threshold SNR value.
        min_freq (float or str):
            Minimum frequency for threshold to be exeeded. If 'f0', then the
            Brune corner frequency will be used.
        max_freq (float):
            Maximum frequency for threshold to be exeeded.
        bandwidth (float):
            Konno-Omachi smoothing bandwidth parameter.
        f0_options (dict):
            Dictionary of f0 options (see config file).
        config (dict):
            Configuration dictionary (or None). See get_config().

    Returns:
        trace: Trace with SNR check.
    """
    for tr in st:
        if tr.hasCached("snr"):
            snr_dict = tr.getCached("snr")
            snr = np.array(snr_dict["snr"])
            freq = np.array(snr_dict["freq"])

            # If min_freq is 'f0', then compute Brune corner frequency
            if min_freq == "f0":
                min_freq = brune_f0(
                    moment_from_magnitude(mag),
                    f0_options["stress_drop"],
                    f0_options["shear_vel"],
                )
                if min_freq < f0_options["floor"]:
                    min_freq = f0_options["floor"]
                if min_freq > f0_options["ceiling"]:
                    min_freq = f0_options["ceiling"]

            # Check if signal criteria is met
            mask = (freq >= min_freq) & (freq <= max_freq)
            if np.any(mask):
                min_snr = np.min(snr[mask])
            else:
                min_snr = 0

            if min_snr < threshold:
                logging.debug(f"{tr.id} failed SNR check {min_snr:.2f} < {threshold:.2f}.")
                tr.fail(f"SNR check: SNR {min_snr:.2f} < {threshold:.2f}")
        snr_conf = {"threshold": threshold, "min_freq": min_freq, "max_freq": max_freq}
        tr.setParameter("snr_conf", snr_conf)
    return st


def compute_snr_trace(tr, event_magnitude, bandwidth=20.0):
    """Compute SNR dictionaries for a trace.

    Args:
        event_magnitude (float):
            Earthquake magnitude.
        bandwidth (float):
            Konno-Omachi smoothing bandwidth parameter.

    """
    def _compute_event_spectra(preevent_noise_spectrum, event_spectrum, dur_preevent, dur_event):
        """Compute noise and signal spectra from event spectrum using pre-event noise spectrum
        to estimate the event noise spectrum.

        Assumptions:
          - Noise amplitudes scale as sqrt(duration)
          - Noise is stationary and so duration is the length of the prevent window (prevent_noise) and
            and event window (event noise).

        event_spectrum = event_noise_spectrum + event_signal_spectrum

        Noise is stationary => noise spectra normalized by duration are equal ->
        pre-event noise spectrum / sqrt(pre-event duration) = event noise spectra / sqrt(event duration)
        """
        event_noise_spectrum = preevent_noise_spectrum * np.sqrt(dur_event) / np.sqrt(dur_preevent)
        signal_spectrum = event_spectrum - event_noise_spectrum
        return (event_noise_spectrum, signal_spectrum)

    if tr.hasParameter("signal_split"):

        # Split the noise and signal into two separate traces
        split_prov = tr.getParameter("signal_split")
        if isinstance(split_prov, list):
            split_prov = split_prov[0]
        split_time = split_prov["split_time"]
        preevent_noise = tr.copy().trim(endtime=split_time)
        event = tr.copy().trim(starttime=split_time)

        tr.setCached("preevent_noise_trace", {"times": preevent_noise.times(), "data": preevent_noise.data})

        preevent_noise.detrend("demean")
        event.detrend("demean")

        # Taper both windows
        preevent_noise.taper(max_percentage=TAPER_WIDTH, type=TAPER_TYPE, side=TAPER_SIDE)
        event.taper(max_percentage=TAPER_WIDTH, type=TAPER_TYPE, side=TAPER_SIDE)

        # Check that there are a minimum number of points in the noise window
        if preevent_noise.stats.npts < MIN_POINTS_IN_WINDOW:
            # Fail the trace, but still compute the event spectra
            # ** only fail here if it hasn't already failed; we do not yet
            # ** support tracking multiple fail reasons and I think it is
            # ** better to know the FIRST reason if I have to pick one.
            if tr.passed:
                tr.fail("SNR check; Not enough points in noise window")
            compute_and_smooth_spectrum(tr, bandwidth, "event")
            return tr

        # Check that there are a minimum number of points in the event window
        if event.stats.npts < MIN_POINTS_IN_WINDOW:
            # Fail the trace, but still compute the event spectra
            if tr.passed:
                tr.fail("SNR check; Not enough points in event window")
            compute_and_smooth_spectrum(tr, bandwidth, "event")
            return tr

        nfft = max(next_pow_2(event.stats.npts), next_pow_2(preevent_noise.stats.npts))
        compute_and_smooth_spectrum(tr, bandwidth, "noise", preevent_noise, nfft)
        compute_and_smooth_spectrum(tr, bandwidth, "event", event, nfft)

        # Noise, event, and signal durations.
        #
        # - Pre-event noise duration is the pre-event window duration.
        # - Event noise duration is the event window duration.
        # - Ground motions at these frequencies are approximately white noise, and thus
        #   also scale as sqrt(duration); ground motion is not stationary, so we'll use
        #   the duration estimated from the earthquake magnitude.
        dur_preevent = preevent_noise.stats.endtime - preevent_noise.stats.starttime
        dur_event = event.stats.endtime - event.stats.starttime
        dur_shaking = duration_from_magnitude(event_magnitude)

        # Compute noise and signal spectra.
        preevent_noise_spectrum = tr.getCached("noise_spectrum")["spec"]
        event_spectrum = tr.getCached("event_spectrum")["spec"]
        event_noise_spectrum, signal_spectrum = _compute_event_spectra(
            preevent_noise_spectrum, event_spectrum, dur_preevent, dur_event)
        tr.setCached(
            "noise_spectrum",
            {
                "spec": event_noise_spectrum,
                "freq": tr.getCached("noise_spectrum")["freq"],
            },
        )
        tr.setParameter("noise_spectrum", {"duration": dur_event})
        tr.setCached(
            "signal_spectrum",
            {
                "spec": signal_spectrum,
                "freq": tr.getCached("event_spectrum")["freq"],
            },
        )
        tr.setParameter("signal_spectrum", {"duration": dur_shaking})

        # Compute smooth noise and signal.
        smooth_preevent_noise_spectrum = tr.getCached("smooth_noise_spectrum")["spec"]
        smooth_event_spectrum = tr.getCached("smooth_event_spectrum")["spec"]
        smooth_event_noise_spectrum, smooth_signal_spectrum = _compute_event_spectra(
            smooth_preevent_noise_spectrum, smooth_event_spectrum, dur_preevent, dur_event)
        tr.setCached(
            "smooth_noise_spectrum",
            {
                "spec": smooth_event_noise_spectrum,
                "freq": tr.getCached("smooth_noise_spectrum")["freq"],
            },
        )
        tr.setCached(
            "smooth_signal_spectrum",
            {
                "spec": smooth_signal_spectrum,
                "freq": tr.getCached("smooth_event_spectrum")["freq"],
            },
        )
 
        smooth_event_noise_normspectrum = smooth_event_noise_spectrum / np.sqrt(dur_event)
        smooth_signal_normspectrum = smooth_signal_spectrum / np.sqrt(dur_shaking)
        snr = smooth_signal_normspectrum / smooth_event_noise_normspectrum
        snr_dict = tr.setCached("snr", {
            "snr": snr,
            "freq": tr.getCached("smooth_event_spectrum")["freq"],
            })

    else:
        # We do not have an estimate of the event split time for this trace
        compute_and_smooth_spectrum(tr, bandwidth, "event")

    return tr
