#!/usr/bin/env python3

import numpy
import h5py

import sys,os

#==============================================================================

if __name__ == "__main__":

   import argparse
   from tesfdmtools.utils.widgets.filechooser import getfile

   pars=['ampbias','ampflux','febias','feflux']

   parser=argparse.ArgumentParser(\
                  description="Plot of a selected SQUID noise spectrum")
   for par in pars:
      parser.add_argument("-%s" % par,type=float,required=False,default=None,\
                          help="the %s value for which to plot the spectrum" % par)
   parser.add_argument('--pdf',action='store_true',required=False,\
                       help="Output plot fo PDF file")
   parser.add_argument('-n','--nonorm',action='store_true',required=False,\
                       help='if set, do not normalize spectra on calibration tone flux')
   parser.add_argument('-ct','--caltone',type=float,nargs=2,default=[908.0,918.0],required=False,\
                       help='search range for calibration tone')
   parser.add_argument('-f','--file',help='HDF5 file with SQUID data',default=None,required=False)

   args=parser.parse_args()

   import matplotlib
   if args.pdf:
      matplotlib.use('Agg')
   else:
      matplotlib.use('GTK3Agg')
   import matplotlib.pyplot as plt
   from matplotlib.backends.backend_pdf import PdfPages

   from tesfdmtools.methods.HDF_SQUIDutils import getspectra

   if args.file is None:
      filename=getfile(pattern='*.h5',path='.')
   else:
      filename=args.file

   fnam=os.path.basename(filename).rsplit('.',1)[0]

   hdf,data=getspectra(filename)

   print("dimensions of data: ")
   for key in data.keys():
      ptxt="%10.10s: %25.25s" % (key,data[key].shape)
      ptxt=ptxt+"   "+"range:%8.2g%8.2g" % (data[key].min(),data[key].max())
      print(ptxt)

   ics={}
   for par in pars:
      ics[par]=0
      ii=eval('args.%s' % par)
      if ii is not None:
         ics[par]=numpy.argmin(numpy.absolute((ii-data[par])))
   print("selected indices: ",ics)
   psid=''
   for par in pars:
      psid=psid+("%d" % ics[par])

   title=''
   for par in pars:
      title=title+(r'%s=%8.2f $\mu$A' % (par,data[par][ics[par]])).replace(' ','')+' '

   if args.pdf:
      psfile=PdfPages('Spectrum_%s_%s.pdf' % (fnam,psid) )
      plt.suptitle('%s' % fnam,size=12)

# calfactor: convert to flux noise by tone->pp *100 (1%phi0), *2 (SingleSided->pp)*sqrt(BW)
# *2.0 (Hann processing gain factor) BW= 40MHz/2^16=610.4Hz  (JvdK)
   try:
      nsamples=int(hdf['squidnoise']['spectrum'].attrs['nsamples'])
      print("nsamples: ",nsamples)
   except:
      nsamples=2**16
   calfactor=100*2*2*numpy.sqrt(40e6/nsamples)

   calindx,=numpy.where( ( data['freq'] > args.caltone[0] ) & ( data['freq'] < args.caltone[1] ) )

   spec=data['spectrum'][ics['ampbias'],ics['ampflux'],ics['febias'],ics['feflux'],:]
   spec_cal=spec.copy()/1E3 			# data is stored in mV, convert back to V
   if not args.nonorm:         			# normalize spectrum to intensity in cal-tone.
       itone=numpy.argmax(spec[calindx])+calindx[0]             # find cal-tone index
       spec_cal=spec/spec[itone]/calfactor      # normalize

   fig,ax=plt.subplots(1)
   ax.loglog(data['freq'],spec_cal,'b-',label='')
   ax.set_title(title)
   ax.set_xlabel('Frequency (kHz)')
   if args.nonorm:
      ax.set_ylabel(r'SQUID Output Voltage (mV$/\sqrt{Hz}$')
   else:
      ax.set_ylabel(r'Equivalent input flux ($\phi_0/\sqrt{Hz}$)') 
   
   if args.pdf:
      fpage = plt.gcf()
      fpage.set_size_inches(10.75, 8.27)
      plt.savefig(psfile,format='pdf',papertype='a4',orientation='landscape')
      psfile.close()
   else:
      plt.show()
   plt.close('all')