#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Nov 16 11:43:12 2021

@author: davidhealy
"""
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal # further FFT functionality
from obspy.clients.fdsn import Client
from obspy import UTCDateTime

T_START = 0     #   length in seconds of data to plot before origin time
T_END = 60*60*24*14    #   length in seconds of data to plot after origin time

#   now get some station meta data, esp. lat & long 
#namStation = 'R36CB'    #   Cheadle Hulme    
namStation = 'R6C8A'    #   MMU   
#namStation = 'R72C7'    #   Aberdeen Uni   
netStation = 'AM'
locStation = '00'
chaStation = 'EHZ'      #   vertical component of geophone 
client = Client('RASPISHAKE')

#   start of 'event' 
EQ_TIME = "2022-04-27T00:00:00"
dtEvent = UTCDateTime(EQ_TIME)
sEventName1 = 'Raw data'
sEventName2 = 'Filtered data'

#   get the data
#   plot seismograms  
t1 = dtEvent - T_START
t2 = dtEvent + T_END

# Download and filter data for this station
st1 = client.get_waveforms(netStation, namStation, locStation, chaStation,
                          starttime=t1, endtime=t2, attach_response=True)
st1.merge(method=0, fill_value='latest')

stf = st1.copy()
Fhigh = 1. 
stf.filter("highpass", freq=Fhigh, corners=2)

window_size = 512
recording_rate = 100
frequencies, times, amplitudes = signal.spectrogram(stf[0].data, fs=recording_rate, 
                                                    window='hamming', nperseg=window_size, 
                                                    noverlap=window_size - 100, 
                                                    detrend=False, scaling="density")
decibels = 20 * np.log10(amplitudes)

# Now plot the waveform data
fig, axs = plt.subplots(2, 1, figsize=(7,5))

axs[0].plot(st1[0].times(reftime=dtEvent+1), st1[0].data, linewidth=1)
axs[0].grid(True)
axs[0].set_xlabel("Time from start (s)")
axs[0].set_ylabel("Counts")
axs[0].set_xlim(T_START+1, T_END)
#axs[0].set_ylim(-10e3, 40e3)

pcm = axs[1].pcolormesh(times, frequencies, decibels, 
                        cmap="plasma", 
                        shading='auto',
                        vmin=-25, vmax=125)
axs[1].set_ylabel("Frequency (Hz)")
axs[1].set_xlabel("Time from start (s)")
axs[1].set_ylim(1., 50.)
#plt.colorbar(pcm, ax=axs[1], orientation='vertical')

plt.tight_layout() 
plt.savefig('plotSeismogramSpectrogramFFT_paper.png', dpi=300)