import sys
import os
import casac
import string
import inspect
import pdb

from parameter_check import *

casapath=os.environ['AIPSPATH']
import asap as sd
os.environ['AIPSPATH']=casapath

def sdcal(infile=None,telescope=None,fluxunit=None,specunit=None,frame=None,doppler=None,calmode=None,scanlist=None,field=None,iflist=None,scanaverage=None,timeaverage=None,polaverage=None,kernel=None,kwidth=None,tau=None,blmode=None,blpoly=None,interactive=None,masklist=None,sdfile=None,outform=None,plotlevel=None,thresh=None,avg_limit=None,edge=None):
	"""ASAP SD calibration task (prototype):

	Task Version 2007-03-06 STM
	Keyword arguments:
        infile -- name of input SD dataset
	        options: (str) file name
		default: '' (none set) REQUIRED
		example: 'mopra-2005-05-08_0350.rpf'
		         Supported formats: ASAP,MS,RPFITS,SDFITS
	telescope -- the telescope name or characteristics
	        options: (str) name or (list) list of gain info
		default: '' (none set)
		example: telescope='GBT' for the GBT
		         telescope='AT' for one of the default scopes
			 known to ASAP.
			 telescope=[104.9,0.43] diameter(m), ap.eff.
			 telescope=[0.743] gain in Jy/K
			 telescope='FIX' to change default fluxunit
			 see description below
	fluxunit -- units for line flux
	        options: 'K','Jy',''
		default: '' (keep current fluxunit)
		WARNING: For GBT data, see description below.
	specunit -- units for spectral axis
	        options: (str) 'channel','km/s','GHz','MHz','kHz','Hz'
	        default: 'channel'
		example: this will be the units for masklist
	frame -- frequency frame for spectral axis
	        options: (str) 'LSRK','REST','TOPO','LSRD','BARY',
	                 'GEO','GALACTO','LGROUP','CMB'
		default: currently set frame in scantable
		WARNING: frame='REST' not yet implemented
	doppler -- doppler mode
	        options: (str) 'RADIO','OPTICAL','Z','BETA','GAMMA'
	        default: currently set doppler in scantable
	calmode -- calibration mode
	        options: 'ps','nod','fs','fsotf','quotient','none'
	        default: 'none'
		example: choose mode 'none' if you have
		         already calibrated and want to
	                 try baselines or averaging
	scanlist -- list of scan numbers to process
	        default: [] (use all scans)
		example: [21,22,23,24]
		         this selection is in addition to field
			 and iflist
	field -- selection string for selecting scans by name
	        default: '' (no name selection)
		example: 'FLS3a*'
		         this selection is in addition to scanlist
			 and iflist
	iflist -- list of IF id numbers to select
	        default: [] (use all IFs)
		example: [15]
		         this selection is in addition to scanlist
			 and field
	scanaverage -- average integrations within scans
	        options: (bool) True,False
		default: False
		example: if True, this happens in read-in
	                 For GBT, set False!
	timeaverage -- average times for multiple scan cycles
	        options: (bool) True,False
		default: False
		example: if True, this happens after calibration
	polaverage -- average polarizations
	        options: (bool) True,False
		default: False
	kernel -- type of spectral smoothing
	        options: 'hanning','gaussian','boxcar','none'
	        default: 'none'
	kwidth -- width of spectral smoothing kernel
	        options: (int) in channels 
	        default: 5
		example: 5 or 10 seem to be popular for boxcar
		         ignored for hanning (fixed at 5 chans)
                         (0 will turn off gaussian or boxcar)
	tau -- atmospheric optical depth
	        default: 0.0 (no correction)
	blmode -- mode for baseline fitting
	        options: (str) 'auto','list','none'
		default: 'none'
		example: blmode='none' turns off baseline fitting
		         blmode='auto' uses AUTOPARS (see below)
		         in addition to blpoly to run linefinder
			 to determine line-free regions
			 USE WITH CARE! May need to tweak AUTOPARS.
	blpoly -- order of baseline polynomial
	        options: (int) (<0 turns off baseline fitting)
	        default: 5
		example: typically in range 2-9 (higher values
		         seem to be needed for GBT)
	interactive -- interactive mode for baseline fitting
	        options: (bool) True,False
		default: False
		WARNING: Currently this just asks whether you accept
		         the displayed fit and if not, continues
	                 without doing any baseline fit.
	masklist -- list of mask regions to INCLUDE in BASELINE fit
	        default: [] (entire spectrum)
		example: [[1000,3000],[5000,7000]]
		         if blmode='auto' then this mask will be applied
			 before fitting
	sdfile -- Name of output file
	        default: '' (<infile>_cal)
		example: note that sdfile is the OUTPUT of sdcal and
	                 is the INPUT filename param for the other
			 sdtasks to steamline processing
	        WARNING: output file will be overwritten
	outform -- format of output file
	        options: 'ASCII','SDFITS','MS','ASAP'
		default: 'ASAP'
                example: the ASAP format is easiest for further sd
	                 processing; use MS for CASA imaging.
			 If ASCII, then will append some stuff to
			 the sdfile name
	plotlevel -- control for plotting of results
	        options: (int) 0=none, 1=some, 2=more, <0=hardcopy
		default: 0 (no plotting)
		example: plotlevel<0 as abs(plotlevel), e.g.
		         -1 => hardcopy of final plot (will be named
		         <sdfile>_calspec.eps)
		WARNING: be careful plotting in fsotf mode!

	AUTOPARS: the following parameters are used for blmode='auto' ONLY
	-------------------------------------------------------------------
	thresh -- S/N threshold for linefinder
	        default: 5
		example: a single channel S/N ratio above which the channel is
	                 considered to be a detection
	avg_limit -- channel averaging for broad lines
	        default: 4
		example: a number of consequtive channels not greater than
                         this parameter can be averaged to search for broad lines
	edge -- channels to drop at beginning and end of spectrum
	        default: 0
		example: [1000] drops 1000 channels at beginning AND end
		         [1000,500] drops 1000 from beginning and 500 from end

                Note: For bad baselines threshold should be increased,
                and avg_limit decreased (or even switched off completely by
                setting this parameter to 1) to avoid detecting baseline
                undulations instead of real lines.


        DESCRIPTION:

	Task sdcal performs data selection, calibration, and/or spectral
	baseline fitting for single-dish spectra.  By setting calmode='none'
	one can run sdcal on already calibrated data, for further selection
	or baseline fitting.  Likewise, one can set blmode='none' to bypass
	baseline fitting.

	If you give multiple IFs in iflist, then your scantable will have
	multiple IFs.  This can be handled, but there can be funny interactions
	later on.  We recommend you split each IF out into separate files
	by re-running sdcal with each IF in turn.

	ASAP recognizes the data of the "AT" telescopes, but currently
	does not know about the GBT or any other telescope. This task
	does know about GBT.  Therefore, if you wish to change the
	fluxunit (see below), then you need to tell it what to do.  If
	you set telescope='AT' it will use its internal defaults.  If
	you set telescope 'GBT', it will use an approximate aperture
	efficiency conversion.  If you give it a list instead of a
	string, then if the list has a single float it is assumed to
	be the gain in Jy/K, if two or more elements they are assumed
	to be telescope diameter (m) and aperture efficiency
	respectively.

	Note that sdcal assumes that the fluxunit is set correctly in
	the data already.  If not, then set telescope='FIX' and it
	will set the default units to fluxunit without conversion.
	WARNING: If the data in infile is an ms from GBT, it will
	currently say its in 'Jy' but its really 'K', so set
	telescope='FIX' and fluxunit='K' to fix this.
	
	"""
	a=inspect.stack()
	stacklevel=0
	for k in range(len(a)):
   	  if (string.find(a[k][1], 'ipython console') > 0):
       		stacklevel=k
       		break
	myf=sys._getframe(stacklevel).f_globals
	myf['__last_taskname']='sdcal'
        tb=myf['tb']
        im = myf['im']
	qa = myf['qa']
        me = myf['me']
	sm = myf['sm']
        ia = myf['ia']
        cl = myf['cl']
        pl = myf['pl']
	rg = myf['rg']
	default=myf['default']
	plotxy=myf['plotxy']
        ###
        ###
        #Handle globals or user over-ride of arguments
        #
        function_signature_defaults=dict(zip(sdcal.func_code.co_varnames,sdcal.func_defaults))
        for item in function_signature_defaults.iteritems():
		key,val = item
                keyVal = eval(key)
                if (keyVal == None):
                        #user hasn't set it - use global/default
                        pass
                else:
                        #user has set it - use over-ride
			#        sys._getframe(1).f_globals[key]=keyVal
			myf[key]=keyVal
			
	###
	#
        #Add type/menu/range error checking here
#	if type(mask)==str: mask=[mask]
        #
	# LIST PARAMETER NAMES HERE
        arg_names=['infile','telescope','fluxunit','specunit','frame','doppler',\
		   'calmode','scanlist','field','iflist',\
		   'scanaverage','timeaverage','polaverage',\
		   'kernel','kwidth','tau','blmode','blpoly','interactive',\
		   'masklist','sdfile','outform','plotlevel',\
		   'thresh','avg_limit','edge']
	#
	# Assign parameter values
	arg_values = []
	for ptst in arg_names:
		stst = myf[ptst]
		if ( type(stst) == str ):
			# Strip off leading and trailing whitespace
			stst = stst.strip()
			myf[ptst] = stst
			sex = ptst+"='"+stst+"'"
		else:
			sex = ptst+"="+str(stst)
		exec(sex)
		arg_values = arg_values + [stst]
	#
	# LIST ALLOWED PARAMETER TYPES HERE
        arg_types=[str,(str,list),str,str,str,str,str,(list,int),str,(list,int),bool,bool,bool,str,(int,float),(float,int),str,int,bool,list,str,str,int,(float,int),(int,float),(list,int,float)]
	#parameter_printvalues(arg_names,arg_values,arg_types)
	#
	# Check allowed types and values
        try:
                parameter_checktype(arg_names,arg_values,arg_types)
		parameter_checkmenu('outform',outform,['ASCII','SDFITS','MS','MS2','ASAP','ascii','sdfits','ms','ms2','asap'])
		parameter_checkmenu('calmode',calmode,['ps','nod','fs','fsotf','quotient','none'])
		parameter_checkmenu('fluxunit',fluxunit,['K','k','Jy','JY','jy',''])
		parameter_checkmenu('specunit',specunit,['channel','km/s','GHz','MHz','kHz','Hz'])
		parameter_checkmenu('frame',frame,['REST','TOPO','LSRD','LSRK','BARY','GEO','GALACTO','LGROUP','CMB',''])
		parameter_checkmenu('doppler',doppler,['RADIO','OPTICAL','Z','BETA','GAMMA','radio','optical','z',''])
		parameter_checkmenu('blmode',blmode,['auto','list','none'])
		parameter_checkmenu('kernel',kernel,['hanning','gaussian','boxcar','none'])
        except TypeError, e:
                print "sdcal -- TypeError: ", e
		return
        except ValueError, e:
                print "sdcal -- OptionError: ", e
                return
        ###

        ###
        ### Now the actual task code
        ###

	try:
#	    pdb.set_trace()
            # strip off leading and trailing whitespace
	    datfile = infile.strip()

            # load the data with or without averaging
            s=sd.scantable(datfile,scanaverage)
	    if ( abs(plotlevel) > 1 ):
		    # print summary of input data
		    print "Initial Raw Scantable:"
		    print s
	    
	    # Default file name
	    if ( sdfile == '' ):
		    project = datfile + '_cal'
	    else:
		    sdfile = sdfile.strip()
		    project = sdfile

	    # determine current fluxunit
	    fluxunit_now = s.get_fluxunit()
	    print "Current fluxunit = "+fluxunit_now
	    
	    # set flux unit string (be more permissive than ASAP)
	    if ( fluxunit == 'k' ):
		    fluxunit = 'K'
	    elif ( fluxunit == 'JY' or fluxunit == 'jy' ):
		    fluxunit = 'Jy'
	    
	    # fix the fluxunit if necessary
	    # No way to query scantable to find what telescope is
	    # so rely on user input
	    if ( telescope == 'FIX' or telescope == 'fix' ):
			    if ( fluxunit != '' ):
				    if ( fluxunit == fluxunit_now ):
					    print "No need to change default fluxunits"
				    else:
					    s.set_fluxunit(fluxunit)
					    print "Reset default fluxunit to "+fluxunit
					    fluxunit_now = s.get_fluxunit()
			    else:
				    print "Warning - no fluxunit for set_fluxunit"
		    
	    elif ( fluxunit=='' or fluxunit==fluxunit_now ):
		    print "No need to convert fluxunits"
	    elif ( telescope == 'GBT' or telescope == 'gbt' ):
		    print "Convert fluxunit to "+fluxunit
		    # THIS IS THE CHEESY PART
		    # Calculate ap.eff eta at rest freq
		    # Use Ruze law
		    #   eta=eta_0*exp(-(4pi*eps/lambda)**2)
		    # with
		    print "Using GBT parameters"
		    eps = 0.390  # mm
		    eta_0 = 0.71 # at infinite wavelength
		    # Ideally would use a freq in center of
		    # band, but rest freq is what I have
		    rf = s.get_restfreqs()[0]*1.0e-9 # GHz
		    eta = eta_0*pl.exp(-0.001757*(eps*rf)**2)
		    print "Calculated ap.eff. eta = %5.3f " % (eta)
		    print "At rest frequency %5.3f GHz" % (rf)
		    D = 104.9 # 100m x 110m
		    print "Assume phys.diam D = %5.1f m" % (D)
		    s.convert_flux(eta=eta,d=D)
		    
		    print "Successfully converted fluxunit to "+fluxunit
	    elif ( telescope == 'AT' or telescope == 'at' ):
		    s.convert_flux()
	    elif ( type(telescope) == list ):
		    # User input telescope params
		    if ( len(telescope) > 1 ):
			    D = telescope[0]
			    eta = telescope[1]
			    print "Use phys.diam D = %5.1f m" % (D)
			    print "Use ap.eff. eta = %5.3f " % (eta)
			    s.convert_flux(eta=eta,d=D)
		    elif ( len(telescope) > 0 ):
			    jypk = telescope[0]
			    print "Use gain = %6.4f Jy/K " % (jypk)
			    s.convert_flux(jyperk=jypk)
		    else:
			    print "Empty telescope list"
	    else:
		    # Unknown telescope type
		    print "Unknown telescope - cannot convert"

	    # set default spectral axis unit
	    if ( specunit != '' ):
		    s.set_unit(specunit)

	    # reset frame and doppler if needed
	    if ( frame != '' ):
		    s.set_freqframe(frame)
	    else:
		    print 'Using current frequency frame'

	    if ( doppler != '' ):
		    if ( doppler == 'radio' ):
			    ddoppler = 'RADIO'
		    elif ( doppler == 'optical' ):
			    ddoppler = 'OPTICAL'
		    elif ( doppler == 'z' ):
			    ddoppler = 'Z'
		    else:
			    ddoppler = doppler
		    
		    s.set_doppler(ddoppler)
	    else:
		    print 'Using current doppler convention'

	    # Select scan and field
	    sel = sd.selector()
	    
	    # Set up scanlist
	    if ( type(scanlist) == list ):
		    # is a list
		    scans = scanlist
	    else:
		    # is a single int, make into list
		    scans = [ scanlist ]
	    # Now select them
	    if ( len(scans) > 0 ):
		    sel.set_scans(scans)
	    
	    # Select source names
	    if ( field != '' ):
		    sel.set_name(field)
		    # NOTE: currently can only select one
		    # set of names this way, will probably
		    # need to do a set_query eventually

	    # Select IFs
	    if ( type(iflist) == list ):
		    # is a list
		    ifs = iflist
	    else:
		    # is a single int, make into list
		    ifs = [ iflist ]
	    if ( len(ifs) > 0 ):
		    # Do any IF selection
		    sel.set_ifs(ifs)

	    # Apply the selection (if any)
	    s.set_selection(sel)

	    del sel

	    scanns = s.getscannos()
	    sn=list(scanns)
	    print "Number of scans to be processed:", len(sn)

	    if ( calmode == 'nod' ):
		    # nod data
		    print "Calibrating nod data"
		    scal=sd.calnod(s,sn)
	    elif ( (calmode=='fs') or (calmode=='fsotf') ):
		    # frequency-switched data    
		    print "Calibrating frequency-switched data"
		    scal=sd.calfs(s,sn)
	    elif ( calmode == 'ps' ):
		    # position switch (default)
		    print "Calibrating position-switched data"
		    scal=sd.calps(s,sn)
	    elif ( calmode == 'quotient' ):
	            # quotient (for Mopra etc.)
		    print "Calibrating using quotient"
		    scal=s.auto_quotient()
	    else:
		    # no calibration
		    print "No calibration"
		    scal=s.copy()

	    # Done with scantable s - clean up to free memory
	    del s
	    
	    # do opacity (atmospheric optical depth) correction
	    if ( tau > 0.0 ):
		    # recalculate az/el (needed for GBT data)
		    scal.recalc_azel()
		    scal.opacity(tau)

	    # Average in time if desired
	    if ( timeaverage ):
		    stave=sd.average_time(scal,weight='tintsys')
		    del scal
		    # Now average over polarizations;Tsys-weighted (1/Tsys**2) average
		    if ( polaverage ):
			    np = stave.npol()
			    if ( np > 1 ):
				    spave=stave.average_pol(weight='tsys')
			    else:
				    # only single polarization
				    print "Single polarization data - no need to average"
				    spave=stave.copy()
		    else:
			    spave=stave.copy()
		    del stave
	    else:
		    if ( polaverage ):
			    np = scal.npol()
			    if ( np > 1 ):
				    spave=scal.average_pol(weight='tsys')
			    else:
				    # only single polarization
				    print "Single polarization data - no need to average"
				    spave=scal.copy()
		    else:
			    spave=scal.copy()
		    del scal
	    
	    if ( abs(plotlevel) > 1 ):
		    # print summary of calibrated data
		    print "Final Calibrated Scantable:"
		    print spave

	    # Smooth the spectrum (if desired)
	    if ( kernel != 'none' and (not (kwidth<=0 and kernel!='hanning'))):
		    if ( abs(plotlevel) > 1 ):
			    # plot spectrum before smoothing
			    # each IF is separate panel, pols stacked
			    sd.plotter.plot(spave)
			    sd.plotter.set_mode(stacking='p',panelling='i')
			    if ( plotlevel < -1 ):
				    # Hardcopy - currently no way w/o screen display first
				    pltfile=project+'_rawspec.eps'
				    sd.plotter.save(pltfile)
		    print "Smoothing spectrum with kernel "+kernel
		    spave.smooth(kernel,kwidth)

	    # Polynomial baseline (if requested, e.g. blpoly=5)
	    if ( blmode == 'auto' ):
		    # Auto-fit baseline
		    if( blpoly >= 0 ):
			    if (len(masklist) > 0):
				    # Use baseline mask for regions to INCLUDE in baseline fit
				    # Create mask using list, e.g. masklist=[[500,3500],[5000,7500]]
				    basemask=spave.create_mask(masklist)
				    spave.auto_poly_baseline(mask=basemask,order=blpoly,edge=edge,
							     threshold=thresh,chan_avg_limit=avg_limit,
							     plot=interactive)
				    del basemask
			    else:
				    # Automatic polynomial baseline without mask
				    spave.auto_poly_baseline(order=blpoly,edge=edge,threshold=thresh,
							     chan_avg_limit=avg_limit,plot=interactive)
	    elif ( blmode == 'list' ):
		    # Fit in specified region
		    if( blpoly >= 0 ):
			    if (len(masklist) > 0):
				    # Use baseline mask for regions to INCLUDE in baseline fit
				    # Create mask using list, e.g. masklist=[[500,3500],[5000,7500]]
				    basemask=spave.create_mask(masklist)
				    spave.poly_baseline(mask=basemask,order=blpoly,plot=interactive)
				    del basemask
			    else:
				    # Without mask
				    spave.poly_baseline(order=blpoly,plot=interactive)
	    
	    # Plot final spectrum
	    if ( abs(plotlevel) > 0 ):
		    # each IF is separate panel, pols stacked
		    sd.plotter.plot(spave)
		    sd.plotter.set_mode(stacking='p',panelling='i')
		    sd.plotter.set_histogram(hist=True)
		    sd.plotter.axhline(color='r',linewidth=2)
		    if ( plotlevel < 0 ):
			    # Hardcopy - currently no way w/o screen display first
			    pltfile=project+'_calspec.eps'
			    sd.plotter.save(pltfile)
			    
	    
	    # Now save the spectrum and write out final ms
	    if ( (outform == 'ASCII') or (outform == 'ascii') ):
		    outform = 'ASCII'
		    spefile = project + '_'
	    elif ( (outform == 'ASAP') or (outform == 'asap') ):
		    outform = 'ASAP'
		    spefile = project
	    elif ( (outform == 'SDFITS') or (outform == 'sdfits') ):
		    outform = 'SDFITS'
		    spefile = project
	    elif ( (outform == 'MS') or (outform == 'ms') or (outform == 'MS2') or (outform == 'ms2') ):
		    outform = 'MS2'
		    spefile = project
	    else:
		    outform = 'ASAP'
		    spefile = project
		    
		    
	    spave.save(spefile,outform,True)
	    print "Wrote output "+outform+" file "+spefile

	    # Clean up scantable
	    del spave
	    
	    # DONE
	
	except TypeError, e:
                print "sdcal -- TypeError: ", e
		return
        except ValueError, e:
                print "sdcal -- OptionError: ", e
                return
	except Exception, instance:
		print '***Error***',instance
		return

        saveinputs=myf['saveinputs']
        saveinputs('sdcal','sdcal.last')

###End of sdcal

def sdcal_defaults():
	a=inspect.stack()
	stacklevel=0
	for k in range(len(a)):
   	  if (string.find(a[k][1], 'ipython console') > 0):
       		stacklevel=k
       		break
	myf=sys._getframe(stacklevel).f_globals
	myf['infile']=''
	myf['telescope']=''
	myf['fluxunit']=''
	myf['specunit']='channel'
	myf['frame']=''
	myf['doppler']=''
	myf['calmode']='none'
        myf['scanlist']=[]
        myf['field']=''
        myf['iflist']=[]
	myf['scanaverage']=False
	myf['timeaverage']=False
	myf['polaverage']=False
        myf['kernel']='none'
        myf['kwidth']=5
        myf['tau']=0.0
        myf['blmode']='none'
        myf['blpoly']=5
        myf['interactive']=False
        myf['masklist']=[]
        myf['sdfile']=''
        myf['outform']='ASAP'
	myf['plotlevel']=0
	myf['thresh']=5
	myf['avg_limit']=4
	myf['edge']=[0]
	
def sdcal_description(key='sdcal'):
	desc={'sdcal': 'ASAP SD calibration task\n',
	'infile': 'name of input SD dataset',
	'telescope': 'name or param of telescope for flux conversion',
	'fluxunit': 'units for line flux (K,Jy) (''=current)',
	'specunit': 'units for spectral axis (channel,km/s,GHz)',
	'frame': 'frequency reference frame, e.g. LSRK (''=current)',
	'doppler': 'doppler convention, e.g. RADIO (''=current)',
	'calmode': 'SD calibration mode (ps,nod,fs,fsotf,none)',
        'scanlist': 'list of scans to use (e.g. [1,2,3,4])',
        'field': 'string for selection by source name (e.g. FLS3a*)',
        'iflist': 'list of IF ids to select (e.g. [0,1])',
	'scanaverage': 'average integs within scans (True,False)',
	'timeaverage': 'average scans over time (True,False)',
	'polaverage': 'average over polarizations (True,False)',
        'kernel': 'type of spectral smoothing (e.g. hanning)',
	'kwidth': 'width of spectral kernel in channels',
        'tau': 'atmospheric optical depth for correction',
	'blmode': 'baseline fitting mode (auto,list,none)',
        'blpoly': 'order of polynomial for baseline fitting (<0=none)',
        'interactive': 'interactive baseline fitting (True/False)',
        'masklist': 'list of channel regions for baseline fitting ([]=auto)',
        'sdfile': 'output file name',
        'outform': 'output file format (ASCII,MS,SDFITS,ASAP)',
	'plotlevel': 'plot results (0=none,1+=some,<0=hardcopy)',
	'thresh': '(Auto) S/N ratio limit for linefinder',
	'avg_limit': '(Auto) max number of chans to average for broad lines',
	'edge': '(Auto) channels to drop from beginning and end',
	}
	return desc[key]

###End of task