'''
HHG spectrogram program to perform 
a frequency-time analysis on high-energy
photons ejected from electron-ion recombination
sites
'''

import numpy as np
from scipy import signal
from scipy.fft import fftshift
from matplotlib import mlab
import matplotlib.pyplot as plt
import matplotlib.pyplot as plt
import matplotlib.colors as c


# INPUT FILE USED FOR HHG CALCULATION GOES HERE
# MAKE SURE INPUT FILE HAS ENOUGH TIME STEPS.
# IF THE FILE HAS TOO FEW TIMESTEPS, THE ENERGY/TIME TRADEOFF
# WILL MAKE THE SPECTROGRAM USELESS!!
m = np.loadtxt("curr_cell_int.out")
np.matrix(m)

# INPUT MATRIX IS SPLICED BY COLUMN
t=m[:,0]  # THIS IS FOR TIME
z=m[:,3]  # THIS IS FOR THE CURRENT/DIPOLE DYNAMICS


# PARAMETERS USED TO DEFINE BINS, OVERLAPS, AND SAMPLING RATE
n_pts=len(t)
n_window=20 # SPECIFIC: BIN SIZE. KEEP TIME-ENERGY UNCERTAINTY P. IN MIND HERE
window_pts=int(n_pts/n_window) # GENERAL: SET BY n_window
overlap_pts=int(window_pts/1.5) # SPECIFIC: WINDOW OVERLAP SET TO YOUR LIKING
fs = int(1/(t[1]-t[0])) # GENERAL: SHOULD EQUAL TIMESTEP IN INPUT FILE
  # CAUTION ---------------------------------------
  # IF YOU TOUCH THE FS, YOU WILL MESS UP THE FREQUENCY MAPPING, AND THE 
  # FREQUNECY AXIS WILL NOT GIVE YOU CORRECT VALUES

print(fs,n_pts,n_window,window_pts,overlap_pts)





# THIS IS WHERE THE SPECTROGRAM IS MADE
#------------INPUT
# NFFT = YOUR TEMPORAL WINDOW
# FS = SAMPLING FREQUENCY
# NOVERLAP = HOW MANY TEMPORAL POINTS DO ADJACENT WINDOWS SHARE
# CMAP = COLOR SCHEME. USE "SEQUENTIAL" COLORSCHEMES ONLY
# MODE = OUTPUT VARIETY. THIS DOESN'T MATTER MUCH. PICK BETWEEN 'MAGNITUDE' AND
#    'DEFAULT,' WHICH IS POWER SPECTRUM DATA
# -----------OUTPUT
# PXX IS THE AMPLITUDE DATA
# FREQS IS THE FREQUENCY DATA
# BINS IS THE -- YOU GUESSED IT -- THE BINS DATA (IN A.U. OF TIME)
#   ITS STARTING AND ENDINGELEMENTS ARE THE MIDPOINT OF THE TEMPORAL WINDOWS
# IM IS THE IMSHOW DATA. WE DON'T CARE ABOUT THIS HERE.
# 


Pxx, freqs, bins, im = plt.specgram(z, NFFT=window_pts,Fs=fs,noverlap=overlap_pts,cmap='hot',mode='magnitude')
plt.clf() # USED TO PREVENT DOUBLE-PLOTTING

print(bins[0],bins[-1]) #THE SIZE OF THE GENERATED WINDOW



# THIS IS WHERE THE SPECTROGRAM IS PLOTTED
# BINS FORMS THE X-AXIS, FREQS FORMS THE Y-AXIS.
# 2*PI*FREQ IS USED HERE BECAUSE MOST GROUP CODES USE OMEGA, NOT F
# PXX IS THE AMPLITUDE DATA. YOU PROBABLY WANT TO PUT THIS ON A LOGSCALE

#NORM SETS THE LOGSCALE. VMIN SETS THE MIN. VMAX SETS THE MAX. THESE VALUES
# ARE DETERMINED BY THE MIN/MAX VALUES OF PXX. THESE MAY BE ADJUSTED
# AS NEEDED

# SHADING IS OUR INTERPOLATION INPUT. GOURAUD INTERPOLATION IS USED.
c = plt.pcolormesh(bins,2*np.pi*freqs,Pxx,norm=c.LogNorm(vmin=Pxx.min(), vmax = Pxx.max()),cmap ='gnuplot2',shading='gouraud')
plt.colorbar(c)
plt.ylim(0,5.5) # ADJUST AS NEEDED
plt.show()


# THAT'S IT.
