# -*- coding: utf-8 -*-
"""
Created on Tue JUL 19 2016
Updated on Thu OCT 19 2017
Updated on 01.05.2023

@author: Christian Monstein
"""

from astropy.io import fits
import matplotlib.pyplot as plt 
import numpy as np
from matplotlib import cm
import scipy as sp
from scipy import ndimage
import pylab as mplot


#------------------------------------------------------------------------------
def smooth(y, box_pts):
    box = np.ones(box_pts)/box_pts
    y_smooth = np.convolve(y, box, mode='same')
    return y_smooth
#------------------------------------------------------------------------------

files = 'BLEN7M_20101103_121505_25.fit.gz'

hdu = fits.open(files)
print (hdu.info()) # FIT-file structure

dB = hdu[0].data.astype(np.float32)/255.0*2500.0/25.4 # conversion digits -> dB
mini_dB = np.min(dB) # find lowest value
rel_dB = dB - mini_dB # set background 0

freqs = hdu[1].data['Frequency'][0] # extract frequency axis
time  = hdu[1].data['Time'][0] # extract time axis
hdu.close()

extent = (time[0], time[-1], freqs[-1], freqs[0])

rel_dB = rel_dB - rel_dB.mean(axis=1, keepdims=True)  # subtract average 
rel_dB = rel_dB.clip(-2,12) # limit peak values +/-
#------------------------------------------------------------------------------

plt.figure(figsize=(15,8))
plt.imshow(rel_dB, aspect = 'auto', extent = extent, cmap=cm.hot) 
## cm.PRGn, cm.hot, cm.cool, cm.bone, cm.binary, cm.spectral, cm.jet
plt.tick_params(labelsize=14)
plt.xlabel('Time [s]',fontsize=15)
plt.ylabel('Frequency [MHz]',fontsize=15)
plt.title('FIT file presentation as raw data, full size',fontsize=15)

cbar = plt.colorbar()
cbar.ax.set_ylabel('dB above background',fontsize=15)
plt.savefig('2D_raw.png')

#------------------------------------------------------------------------------
plt.figure(figsize=(15,8))
mychannel = 160 # you need to find the channel for a certain frequency
lightcurve = rel_dB[ mychannel , :] # Single Lightcurve at channel mychannel
plt.plot(time,lightcurve,'.',linewidth=2.0,color='b')

smoothLC = smooth(lightcurve,17) # smoothing over 17 pixels
plt.plot(time,smoothLC,linewidth=3.0,color='r')

plt.ylabel('Intensity [dB] above background',fontsize=15)
plt.xlabel('Time after 03:45:00 UT [s]',fontsize=15)
st1 = 'Light curve at ={:7.1f}'.format(freqs[mychannel]) + ' MHz' + ', corresponding to channel {:4.0f}'.format(mychannel)
plt.title(st1,fontsize=15)
plt.tick_params(labelsize=14)
#plt.axis([0, 3600, 120, 200])
plt.grid(True)
plt.text(2500,9, 'Dots = raw light curve',color="blue",size=20)
plt.text(2500,8, 'Line = smoothed light curve',color="red",size=20)
plt.savefig('Ligthcurve'+'{:.1f} MHz'.format(freqs[mychannel])+'.png')

#------------------------------------------------------------------------------


plt.figure(figsize=(15,8))
plt.plot(freqs,'+',color='blue')
plt.xlabel('Frequency channel of Callisto',fontsize=15)
plt.ylabel('Frequency [MHz]',fontsize=15)
plt.title('Frequency distribution to estimate channel of interest',fontsize=15)
plt.grid(True, which='major', color='b', linestyle='-')
plt.grid(True, which='minor', color='g', linestyle='--')
plt.minorticks_on()
plt.tick_params(labelsize=14)
plt.savefig('FrequencyChannel.png')

#------------------------------------------------------------------------------


plt.figure(figsize=(15,8))
sigma = [0,2] # number of pixels for smoothing in y- and x-direction
blur = sp.ndimage.gaussian_filter(rel_dB, sigma, mode='constant')
vmin = 2 # minimum color table in dB
vmax = 15 # maximum color table in dB
plt.imshow(blur, aspect = 'auto', extent = extent, cmap=cm.hot,
           norm=plt.Normalize(vmin, vmax)) # defines image contrast
## cm.PRGn, cm.hot, cm.cool, cm.bone, cm.binary, cm.spectral, cm.jet

plt.xlim((0,300)) # x-axis limits in seconds
plt.ylim((200,500)) # y-axis limits in MHz
plt.tick_params(labelsize=14)
plt.xlabel('Time [s]',fontsize=15)
plt.ylabel('Frequency [MHz]',fontsize=15)
plt.title('FIT file presentation as Gaussian smoothed data, zoomed',fontsize=15)

cbar = plt.colorbar()
cbar.ax.set_ylabel('dB above background',fontsize=15)
plt.savefig('2D_smooth.png')

#------------------------------------------------------------------------------
# Here begins the complicated part with x-axis in hh:mm:ss
start_time = float(fits.open(files)[0].header['TIME-OBS'].split(":")[0])*60*60 
+ float(fits.open(files)[0].header['TIME-OBS'].split(":")[1])*60
+ float(fits.open(files)[0].header['TIME-OBS'].split(":")[2])  
ut = time + start_time

fmin = 185.0
fmax = 500.0
k = 0 # index to search for frequencies in list
for f in freqs:
    if (f > fmin):
        klow = k
    if (f > fmax):
        khigh = k
    k = k + 1

tmin = '12:00:00' # enter start-time here as hh:mm:ss
tmax = '12:04:10' # enter stop-time here as hh:mm:ss
dt = fits.open(files)[0].header['CDELT1'] # read time resolution from FIT-file
print (dt)
xmin = (float(int(tmin[0:2])*3600 + int(tmin[3:5])*60 + int(tmin[6:8])) - ut[0])/dt
xmax = (float(int(tmax[0:2])*3600 + int(tmax[3:5])*60 + int(tmax[6:8])) - ut[0])/dt
if (xmin < 0):
    print ('Wrong x-axis time, check FIT-file with JAVA_viewer')
    xmin = 0
if (xmax > 3599):
    print ('Wrong x-axis time (max), check FIT file with JAVA-viewer')
    xmax = 3599

mplot.figure(figsize=(15,10))
ax = mplot.subplot(111)
mplot.axis([xmin, xmax, klow, khigh]) # both axes expressed in image-pixels

mplot.imshow(blur,aspect = 'auto')
cb1 = mplot.colorbar(orientation='vertical', shrink=0.99)
cb1.set_label('dB above background',fontsize=15)
cb1.ax.tick_params(labelsize=15) 

a = ax.get_xticks().tolist()
b = ax.get_yticks().tolist()

for j in range(len(a)-0):
    sec = float(ut[int(a[j])])
    hh  = int(int(sec)/60/60)
    mm  = int(int(sec - hh*60*60)/60)
    ss  = int(int(sec - hh*60*60 - mm*60))
    a[j] = '{:02.0f}:{:02.0f}:{:02.0f}'.format(hh,mm,ss)

for j in range(len(b)-1):
    ff = float(freqs[int(b[j])])
    b[j]=str(ff)[:5]

ax.set_xticklabels(a, fontsize=15)

ax.set_yticklabels(b, fontsize=15)

mplot.ylabel('Plasma frequency [MHz]', fontsize=18)
mplot.xlabel('Time [UT]', fontsize=18)

mplot.title(files, fontsize=20)
mplot.savefig('2D_smooth_with_axes.png',bbox_inches='tight')
#------------------------------------------------------------------------------
