#!/usr/bin/env python3

import numpy
import h5py

import sys,os

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

if __name__ == "__main__":

   import argparse

   import matplotlib
   matplotlib.use('Agg')
   import matplotlib.pyplot as plt
   from matplotlib.backends.backend_pdf import PdfPages

   from tesfdmtools.utils.widgets.filechooser import getfile
   from tesfdmtools.methods.HDF_SQUIDutils import getspectra,fillplot,fillgraph
#   from tesfdmtools.utils.cu import cu

   parser=argparse.ArgumentParser(\
                  description="Plot overview of SQUID noise spectra")
   parser.add_argument('-i','--index',type=int,default=0,required=False,\
                       help='first index of spectral data to select')
   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('--xlim',type=float,nargs=2,default=[None,None],required=False,\
                       help='limits for X-axis (frequency) of the plotted images')
   parser.add_argument('--vrange',type=float,nargs=2,required=False,default=[1E-7,1E-5],\
                       help='value range for color scale')
   parser.add_argument('-c','--compress',type=int,default=1000,required=False,\
                       help='number of plot sampled points along X-axis in the images')
   parser.add_argument('-n','--nonorm',action='store_true',required=False,\
                       help='if set, do not normalize spectra on calibration tone flux')
   parser.add_argument('-tr','--transresistance',type=float,default=None,\
                       help='transresistance [M-1 [uA/Phi0]]. When provided, show 1D plot of calibrated trans resistance versus FE-flux')  
   parser.add_argument('-nf','--noiseatfrequency',type=float,default=None,\
                       help='Creates 1D SQUID noise flux plot at the specified frequency (kHz)'+
                       ' Default is normalized flux noise. In conjuction with the -n option, the unit is voltage'+
                       ' noise spectral density.') 
   parser.add_argument('-f','--file',help='HDF5 file with SQUID data',default=None,required=False)

   args=parser.parse_args()

   if args.file is None:	             # if input file is not given, interactively select file
      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: ")             # show data dimensions
   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)

   xlim=args.xlim
   if ( xlim[0] is None ) or ( xlim[1] is None ):
      xlim[0]=1.0
      xlim[1]=data['freq'].max()

   dim5nam='freq'
   dim5txt='frequency(kHz)'
   if len(data['feflux']) == 1:             # if 'feflux' only has one data point, make 'ampflux' the y-axis
      imdim='35'
      dim1=0
      dim2=data['spectrum'][:,0,0,0,0].size
      dim3=data['spectrum'][0,0,:,0,0].size
      dim1nam='feflux'
      dim2nam='ampbias'
      dim3nam='febias'
      dim4nam='ampflux'
      dim1txt='FEflx'
      dim2txt='Ampbias'
      dim3txt='FEbias'
      dim4txt='AmpFlx'
   else:                                   # otherwise use 'feflux' as y-axis
      imdim='45'
      dim1=data['spectrum'][:,0,0,0,0].size
      dim2=data['spectrum'][0,:,0,0,0].size
      dim3=data['spectrum'][0,0,:,0,0].size
      dim1nam='ampbias'
      dim2nam='ampflux'
      dim3nam='febias'
      dim4nam='feflux'
      dim1txt='Ampbias'
      dim2txt='AmpFlx'
      dim3txt='FEBias'
      dim4txt=r'feflux($\mu$A)'

   if args.index > dim1:
      sys.exit("Given index %d does not exist in spectrum. Max index is %d" % (args.index,(dim1-1)))

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

   xpmax=3		# max number of spectral images in X-direction on page
   ypmax=7		# max number of spectral images in Y-direction on page
   ppage={}

   if dim3 > dim2:	# determine layout of images orientation; determine what will be on each page
      if dim2 < xpmax:
         xpmax=dim2
      nypages=(dim3+ypmax-1)//ypmax
      nxpages=(dim2+xpmax-1)//xpmax
      print("number of pages: ",nypages*nxpages)
      for p in numpy.arange(nypages*nxpages):
         ppage[p]=numpy.empty((xpmax,ypmax),dtype='object')
      nxpage=0
      for i in numpy.arange(dim2):
         px=i-nxpage*xpmax
         if px >= xpmax:
            px=0
            nxpage=nxpage+1         
         nypage=0
         for j in numpy.arange(dim3):
            py=j-nypage*ypmax
            if py >= ypmax:
               py=0
               nypage=nypage+1
            ppage[nypage+nxpage*nypages][px,py]=[args.index,i,j]
   else:
      if dim3 < xpmax:
         xpmax=dim3
      nypages=(dim2+ypmax-1)//ypmax
      nxpages=(dim3+xpmax-1)//xpmax
      print("number of pages: ",nypages*nxpages)
      for p in numpy.arange(nypages*nxpages):
         ppage[p]=numpy.empty((xpmax,ypmax),dtype='object')
      nxpage=0
      for i in numpy.arange(dim3):
         px=i-nxpage*xpmax
         if px >= xpmax:
            px=0
            nxpage=nxpage+1
         nypage=0
         for j in numpy.arange(dim2):
            py=j-nypage*ypmax
            if py >= ypmax:
               py=0
               nypage=nypage+1
            ppage[nypage+nxpage*nypages][px,py]=[args.index,j,i]

   if args.nonorm:
      prefix='SQUID_Vnoise'
   else:
      prefix='SQUID_Fluxnoise'

   if args.transresistance is not None:
      prefix='SQUID_Transresistance'

   noisefreq=None
   if args.noiseatfrequency is not None:
      prefix='SQUID_Noise_at_'+('%d' % (int(args.noiseatfrequency)) )+'khz'
      findx,=numpy.where( data['freq'] > args.noiseatfrequency )
      if findx.size == 0:
         sys.exit("Error, given frequency %11.3f (kHz) for noise flux does not exist" % args.noiseatfrequency)
      noisefreq=findx[0]
      if noisefreq > 0: 
         if (args.noiseatfrequency-data['freq'][noisefreq-1]) < \
            (data['freq'][noisefreq]-args.noiseatfrequency):
            noisefreq=noisefreq-1
      realnoisefreq=data['freq'][noisefreq]

   psfile=PdfPages('%s_%s_%3.3d.pdf' % (prefix,fnam,args.index) )
   matplotlib.rcParams.update({'font.size': 6})

   vrange=numpy.array(args.vrange,dtype=float)  # range for color scale
   if ( vrange[1] == 1E-5 ) and args.nonorm:
      vrange=numpy.array([1E-10,1E-8],dtype=float)  # default in case of 'nonorm'
 
   for ip in ppage.keys():                     # go through all pages

      print("Page %d:" % ip)
      fig=plt.figure()
      if dim1nam not in data:
         pptitle='%s %s=%6.2f $\mu$A' % (fnam,dim1txt,0)
      else:
         pptitle='%s %s=%6.2f $\mu$A' % (fnam,dim1txt,data[dim1nam][args.index])
      if args.noiseatfrequency is not None:
         pptitle=pptitle+' @ '+('%12.3f (kHz)' % realnoisefreq).strip()
      plt.suptitle(pptitle,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)

      probecurrent=1.0
      if args.transresistance is not None:
         probecurrent=args.transresistance

      for i in numpy.arange(ppage[ip][:,0].size):
         for j in numpy.arange(ppage[ip][0,:].size):
            if ppage[ip][i,j] is None:
               print("ppage is None")
               break
            lastplot=False
            if j <  (ppage[ip][0,:].size-1):
               if ppage[ip][i,(j+1)] is None:
                  lastplot=True
            print(ppage[ip][i,j])
            dim1i=ppage[ip][i,j][0]
            dim2i=ppage[ip][i,j][1]
            dim3i=ppage[ip][i,j][2]
            if imdim == '35':
               spec=data['spectrum'][dim2i,:,dim3i,0,:]
            else:
               spec=data['spectrum'][dim1i,dim2i,dim3i,:,:]

            spec_cal=spec.copy()/1E3	# data is stored in mV, convert back to V
            norm_cal=numpy.zeros(spec[:,0].size,dtype=float)			# buffer with spectrum norms 
            if noisefreq is not None:
               sqd_nois=spec[:,noisefreq]/1E3					# flux in V for noiseatfrequncy
            if not args.nonorm:         # normalize spectrum to intensity in cal-tone.
               nrows=spec[:,0].size
               for irow in numpy.arange(nrows):	# go through all spectra in image
                  itone=numpy.argmax(spec[irow,calindx])+calindx[0]	        # find cal-tone index
                  spec_cal[nrows-irow-1,:]=spec[irow,:]/spec[irow,itone]/calfactor	# normalize
                  norm_cal[nrows-irow-1]=spec[irow,itone]*calfactor*1E-3/probecurrent	# keep normalization constant
                  if noisefreq is not None:
                     sqd_nois[nrows-irow-1]=spec[irow,noisefreq]/spec[irow,itone]/calfactor # ssquid noise at given frequency
            plot=plt.subplot2grid((ypmax,xpmax),(j,i))
            xtitle=''
            if lastplot or (j == ppage[ip][0,:].size-1):
               xtitle=dim5txt
            if dim2nam not in data:
               ptitle=r'%s: %5.2f $\mu$A, %s=: %6.2f $\mu$A' % \
                   (dim2txt,0,dim3txt,data[dim3nam][dim3i])
            else:
               ptitle=r'%s: %5.2f $\mu$A, %s=: %6.2f $\mu$A' % \
                   (dim2txt,data[dim2nam][dim2i],dim3txt,data[dim3nam][dim3i])
            ytitle=''
            if i == 0:
               ytitle=dim4txt
            if args.transresistance is not None:
               ytit=''
               xtit=''
               if ytitle != '':
                  ytit=r'Trans resistance ($\Omega$)'
               if xtitle != '':
                  xtit=dim4txt
               pl=fillgraph(plot,norm_cal,xax=data[dim4nam],ytitle=ytit,\
                            xtitle=xtit,ptitle=ptitle) 
               if xtit == '':
                  plot.set_xticklabels([])
            elif args.noiseatfrequency is not None:
               ytit=''
               xtit=''
               if ytitle != '':
                  ytit=r'flux ($\phi_0/\sqrt{Hz}$)'
                  if args.nonorm:
                     ytit=r'noise spectral density (V)'
               if xtitle != '':
                  xtit=dim4txt
               pl=fillgraph(plot,sqd_nois,xax=data[dim4nam],ytitle=ytit,\
                            xtitle=xtit,ptitle=ptitle) 
               if xtit == '':
                  plot.set_xticklabels([])               
            else:
               im=fillplot(plot,spec_cal,xax=data[dim5nam],yax=data[dim4nam],\
                           xtitle=xtitle,ytitle=ytitle,ptitle=ptitle,vrange=vrange,\
                           compress=args.compress,xlim=xlim)
               if xtitle == '':
                  plot.set_xticklabels([])
            if lastplot:
               break

      print("finish page: ",ip," .................")

      plt.tight_layout()
      plt.subplots_adjust(bottom=0.100,top=0.940,left=0.100,right=0.850,hspace=0.20,wspace=0.15)

      if ( args.transresistance is None ) and ( args.noiseatfrequency is None ) :
         cax = fig.add_axes([0.905, 0.1, 0.03, 0.840])
         cbar=fig.colorbar(im, cax=cax)
         if not args.nonorm:
            cbar.set_label(r'flux noise spectral density ($\phi_0/\sqrt{Hz}$)')
         else:
            cbar.set_label(r'SQUID output voltage noise spectral density (V/$\sqrt{Hz}$)')

      fpage = plt.gcf()
      fpage.set_size_inches(8.27, 10.75)
      print("save to file......................")                   
      plt.savefig(psfile,format='pdf',papertype='a4')
      plt.close('all')

   print("close pdf file...................")      
   psfile.close()