# -*- coding: utf-8 -*-
"""
Created on Thursday, 2020-04-09, 2020-05-06
Read FITS-file from RaspBerry Pi 4G and SDRplay
and generates a map
Assure, that the size of the fits-files is the same for all scans
Do not change spectrometer configuration duirng a scan
@author: Christian Monstein
"""
#-------------------------------------------------------------------------------------------

import numpy as np
import astropy.io.fits as fits
import matplotlib.pyplot as plt
from matplotlib import cm
import glob
import os
from scipy.ndimage.filters import gaussian_filter

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

MyPath = '' # Edit in case the file is elsewhere, outside local folder, use '/' instead of '\'
MyFile = 'dish_TRANSIT_*_spec.fits' # ->edit

scanoffset = 0 # makes (may be) y-label nicer -> edit
calibrationfactor = 500.0 # conversion ADU->Kelvin -> edit
vmin = 5.0 # -> edit low cut color table 0... +/-10
vmax = 40.0 # -> edit high cut color table 10 .... 100
RAstepsize = 4 # -> edit 1....(3)....15, depending on number of scans
VLSRmin = -50 # -> edit km/s
VLSRmax = 50 # -> edit km/s
scanmin = 667 # -> edit, check folder with fits-files
scanmax = 810 # -> edit, check folder with fits-files
colorbar = False # -> edit: True or False

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

def get_data_from_fits(path):
    with fits.open(path) as hdu:
        data = hdu[0].data#.astype(np.uint8) # float16, float32, uin8 depends on PC
        raci = hdu[0].header['CRVAL2']
        rac.append(raci)
        scannr = hdu[0].header['SCAN-NUM'] - scanoffset
        scan.append(scannr)
    return data
    
#-------------------------------------------------------------------------------------------

paths = glob.glob(MyPath + MyFile) # search the folder
#paths.sort(key = os.path.getctime) # resort for creation date
paths.sort()
print ('FITS-files found: ',(len(paths)))

rac = []
scan= []

hdu       = fits.open(paths[0])
data      = hdu[0].data
date      = hdu[0].header['DATE-OBS']
time      = hdu[0].header['UT']
telescope = hdu[0].header['TELESCOP']
raci      = hdu[0].header['CRVAL2']
Nbins     = hdu[0].header['NAXIS1']
df        = hdu[0].header['CDELT1'] / 1.e6
shift     = hdu[0].header['CRPIX1']
Frest     = hdu[0].header['RESTFREQ'] / 1.e6
scannr    = hdu[0].header['SCAN-NUM'] - scanoffset
azi       = hdu[0].header['AZIMUTH']
ele       = hdu[0].header['ELEVATIO']
dec       = hdu[0].header['CRVAL3']

rac.append(raci)
scan.append(scannr)

hdu.close()
print ('Header information: ',hdu.info()) # FIT-file structure
print ('Keyords in header: ',hdu[0].header) # Header-information (keywords)

data = np.vstack([get_data_from_fits(path) for path in paths])
print ('Data shape: ',data.shape)

N = len(paths)
x = np.column_stack((scan,rac))
print ('Scan number      Right Ascension')
print (x)

new = data # normalize spectrum to a level just aside lab frequency
for s in range(0,N):
    ref = data[s,int(shift-200)] # subtract min(standing waves) near HII-line ->edit
    new[s,:] = data[s,:]/ref

dmean = np.mean(new, axis=0) # average spectrum over all right ascensions
bgs = new - dmean # subtraction background

f1 =    - shift*df
f2 = f1 + Nbins*df
faxis = np.arange(f1,f2,df) # generate frequency axis [MHz]

faxis2 = [] # convert delta-frequency in vlsr
for f in faxis:
    ff = f/Frest*3e5 # MHz -> km/s
    faxis2.append(ff)
    
extent = (faxis2[0], faxis2[-1], scan[-1],scan[0])
sigma = 1. # check different values 0.5 ..(2)... 3 ->edit
bgsf = gaussian_filter(bgs, sigma) # smoothing at bit

bgsf2 = bgsf*calibrationfactor

fig1, ax1 = plt.subplots(figsize=(8,7))  # 6,10 ->edit
 
img=ax1.imshow(bgsf2, aspect = 'auto',extent = extent, cmap=cm.hot, vmin=vmin,vmax=vmax) #vmin=0.01,vmax=0.1 ->edit
# https://matplotlib.org/examples/color/colormaps_reference.html
# jet, CMRmap, gnuplot2, magma, plasma, Spectral

ax1.set_xlabel('Observed velocity VLSR [km/s]',fontsize=15)
ax1.set_ylabel('Running scan number',fontsize=15)
ax1.set_xlim(VLSRmin, VLSRmax) 
ax1.set_ylim(scanmax, scanmin) # Limitation in scan-number, plt.plot(scan,rac,'*r')
ax1.set_title('21cm transit scan at dec={:4.1f}$\degree$ (azi=170$\degree$, ele=71$\degree$)'.format(dec),fontsize=17)

st = []
for i in range(scanmin,scanmax,RAstepsize):
    j = scanmax - i 
    text = '{:04.1f}'.format(x[j,1])
    st.append(text)
    
ax2 = ax1.twinx()
ax2.set_yticks(np.arange(len(st)+1))
ax2.set_yticklabels(st,size=12)
ax2.set_ylabel('Right ascension [$\degree$]',fontsize=15)

if (colorbar):
    cbar = plt.colorbar(img,orientation="horizontal", pad=0.1)
    cbar.set_label(label='Brightness temperature [K]', size='large')
    cbar.ax.tick_params(labelsize='large')

fig1.tight_layout()
plt.savefig('scan.png', bbox_inches=0, dpi=100) # ->edit

print ('Done...')
#-------------------------------------------------------------------------------------------