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 sdplot(sdfile=None,telescope=None,fluxunit=None,specunit=None,frame=None,doppler=None,scanlist=None,field=None,iflist=None,scanaverage=None,timeaverage=None,polaverage=None,kernel=None,kwidth=None,stack=None,panel=None,flrange=None,sprange=None,linecat=None,linedop=None,plotfile=None):
	"""ASAP SD plotting task (prototype):

	Task Version 2007-03-06 STM
	Keyword arguments:
        sdfile -- name of input SD dataset
	        default: none - must input file name
		example: 'mysd.asap'
		         See sdcal for allowed formats.
	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 unit)
		WARNING: For GBT data, see description below.
	specunit -- units for spectral axis
	        options: (str) 'channel','km/s','GHz','MHz','kHz','Hz',''
	        default: '' (keep current specunit)
		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
	scanlist -- list of scan numbers to process
	        default: [] (use all scans)
		example: [21,22,23,24]
	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]
	scanaverage -- average integrations within scans
	        options: (bool) True,False
		default: False
		example: if True, this happens in read-in
	timeaverage -- average times for multiple scans
	        options: (bool) True,False
		default: True
	polaverage -- average polarizations
	        options: (bool) True,False
		default: True
	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)
	stack -- code for stacking on single plot
	        options: 'p','b','i','t','s' or
		         'pol', 'beam', 'if', 'time', 'scan'
		default: 'p'
		example: maximum of 25 stacked spectra
	                 stack by pol, beam, if, time, scan
	panel -- code for splitting into multiple panels
	        options: 'p','b','i','t','s' or
		         'pol', 'beam', 'if', 'time', 'scan'
		default: 'i'
		example: maximum of 25 panels
	                 panel by pol, beam, if, time, scan
	flrange -- range for flux axis of plot
	        options: (list) [min,max]
		default: [] (full range)
		example: flrange=[-0.1,2.0] if 'K'
		         assumes current fluxunit
	sprange -- range for spectral axis of plot
	        options: (list) [min,max]
		default: [] (full range)
		example: sprange=[42.1,42.5] if 'GHz'
		         assumes current specunit
	linecat -- control for line catalog plotting
	        options: (str) 'all','none' or by molecule
		default: 'none' (no lines plotted)
		example: linecat='SiO' for SiO lines
		         linecat='*OH' for alcohols
	                 uses sprange to limit catalog
		WARNING: specunit must be in frequency (*Hz)
			 to plot from the line catalog!
			 and must be 'GHz' or 'MHz' to use 
			 sprange to limit catalog
	linedop -- doppler offset for line catalog plotting
	        options: (float) doppler velocity (km/s)
		default: 0.0
		example: linedop=-30.0
	plotfile -- file name for hardcopy output
	        options: (str) filename.eps,.ps,.png
		default: '' (no hardcopy)
		example: 'specplot.eps','specplot.png'
	                 Note this autodetects the format from
	                 the suffix (.eps,.ps,.png).
	                 

        DESCRIPTION:

	Task sdplot displays single-dish spectra.  
	It assumes that the spectra have been calibrated in sdcal.
	It does allow selection of scans, IFs, polarizations, and
	some time and channel averaging/smoothing options also,
	but does not write out this data.

	Some plot options, like annotation and changing titles,
	legends, colors, fonts, and the like are not supported
	in this task.  You should use sd.plotter from the ASAP
	toolkit directly for this.

	This task uses the JPL line catalog as supplied by ASAP.
	If you wish to use a different catalog, or have it plot
	the line IDs from top or bottom (rather than alternating),
	then you will need to explore the sd toolkit also.	
	
	Note that multiple scans and IFs can in principle be handled
	through stacking and panelling, but this is fairly rudimentary
	at present and you have little control of what happens in
	individual panels.  We recommend that you use scanlist, field,
	and iflist to give a single selection for each run.

	Currently, setting specunit='GHz' fixes the x-axis span of
	each IF panel to be the same (an example of the limitations
	of ASAP plotting at present).

	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.  See sdcal description for information
	on fluxunit conversion.
	
	WARNING: be careful plotting otf data with lots of fields!

	"""
	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']='sdplot'
        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(sdplot.func_code.co_varnames,sdplot.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=['sdfile','telescope','fluxunit','specunit','frame','doppler',\
		   'scanlist','field','iflist',\
		   'scanaverage','timeaverage','polaverage',\
		   'kernel','kwidth','stack','panel','flprange','sprange',\
		   'linecat','linedop','plotfile']
	#
	# 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,(list,int),str,(int,list),\
		   bool,bool,bool,str,(int,float),str,str,list,list,str,\
		   (float,int),str]
	#parameter_printvalues(arg_names,arg_values,arg_types)
	#
	# Check allowed types and values
        try:
                parameter_checktype(arg_names,arg_values,arg_types)
		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('kernel',kernel,['hanning','gaussian','boxcar','none'])
		parameter_checkmenu('stack',stack,['p','i','t','b','s', 'pol', 'beam', 'if', 'time', 'scan'])
		parameter_checkmenu('panel',panel,['p','i','t','b','s', 'pol', 'beam', 'if', 'time', 'scan'])
        except TypeError, e:
                print "sdplot -- TypeError: ", e
		return
        except ValueError, e:
                print "sdplot -- OptionError: ", e
                return
        ###

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

	try:
#	    pdb.set_trace()
            #load the data with or without averaging
            s=sd.scantable(sdfile,scanaverage)
	    
	    # 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 spectral axis unit
	    if ( specunit != '' ):
		    print "Changing spectral axis to "+specunit
		    s.set_unit(specunit)

	    # reset frame and doppler if needed
	    if ( frame != '' ):
		    print "Changing frequency frame to "+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'

	    # Prepare a selection
	    sel=sd.selector()

	    # Scan selection
	    if ( type(scanlist) == list ):
		    # is a list
		    scans = scanlist
	    else:
		    # is a single int, make into list
		    scans = [ scanlist ]
	    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

	    # IF selection
	    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)

	    # Average in time if desired
	    if ( timeaverage ):
		    stave=sd.average_time(s,weight='tintsys')
		    del s
		    # 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 = s.npol()
			    if ( np > 1 ):
				    spave=s.average_pol(weight='tsys')
			    else:
				    # only single polarization
				    print "Single polarization data - no need to average"
				    spave=s.copy()
		    else:
			    spave=s.copy()
		    del s
	    
	    # Smooth the spectrum (if desired)
	    if ( kernel != 'none' and (not (kwidth<=0 and kernel!='hanning'))):
		    print "Smoothing spectrum with kernel "+kernel
		    spave.smooth(kernel,kwidth)

	    # Plot final spectrum
	    # each IF is separate panel, pols stacked
	    sd.plotter.plot(spave)
	    sd.plotter.set_mode(stacking=stack,panelling=panel)

	    # Plot red x-axis at y=0 (currently disabled)
	    # sd.plotter.set_histogram(hist=True)
	    # sd.plotter.axhline(color='r',linewidth=2)

	    # Set axis ranges (if requested)
	    if len(flrange)==1: print "flrange needs 2 limits - ignoring"
	    if len(sprange)==1: print "sprange needs 2 limits - ignoring"
	    if ( len(sprange) > 1 ):
		    if ( len(flrange) > 1 ):
			    sd.plotter.set_range(sprange[0],sprange[1],flrange[0],flrange[1])
		    else:
			    sd.plotter.set_range(sprange[0],sprange[1])
	    elif ( len(flrange) > 1 ):
		    sd.plotter.set_range(ystart=flrange[0],yend=flrange[1])
	    else:
		    # Set default range explicitly (in case range was ever set)
		    sd.plotter.set_range()

	    # Line catalog
	    dolinc=False
	    if ( linecat != 'none' and linecat != '' ):
		    # Use jpl catalog for now (later allow custom catalogs)
		    catpath=casapath+'/data/catalogs/lines'
		    catname=catpath+'/jpl_asap.tbl'
		    # TEMPORARY: hard-wire to my version
		    # catname='/home/sandrock/smyers/NAUG/Tasks/jpl.tbl'
		    # FOR LOCAL CATALOGS:
		    # catname='jpl.tbl'
		    try:
			    linc=sd.linecatalog(catname)
			    dolinc=True
		    except:
			    print "Could not find catalog at "+catname
			    dolinc=False
		    if ( dolinc ):
			    if ( len(sprange)>1 ):
				    if ( specunit=='GHz' or specunit=='MHz' ):
					    linc.set_frequency_limits(sprange[0],sprange[1],specunit)
				    else:
					    print "ERROR: sd.linecatalog.set_frequency_limits accepts only GHz and MHz"
					    print "continuing without sprange selection on catalog"
			    if ( linecat != 'all' and linecat != 'ALL' ):
				    # do some molecule selection
				    linc.set_name(linecat)
			    # Plot up the selected part of the line catalog
			    # use doppler offset
			    sd.plotter.plot_lines(linc,doppler=linedop)

	    # Hardcopy
	    if ( plotfile != '' ):
		    # currently no way w/o screen display first
		    sd.plotter.save(plotfile)

	    # Do some clean up
	    del spave
	    if dolinc: del linc
	    
	    # DONE
	
	except TypeError, e:
                print "sdplot -- TypeError: ", e
		return
        except ValueError, e:
                print "sdplot -- OptionError: ", e
                return
	except Exception, instance:
		print '***Error***',instance
		return

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

###End of sdplot

def sdplot_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['sdfile']=''
	myf['telescope']=''
	myf['fluxunit']=''
	myf['specunit']=''
	myf['frame']=''
	myf['doppler']=''
        myf['scanlist']=[]
        myf['field']=''
        myf['iflist']=[]
	myf['scanaverage']=False
	myf['timeaverage']=False
	myf['polaverage']=False
	myf['kernel']='none'
        myf['kwidth']=0
	myf['stack']='p'
	myf['panel']='i'
	myf['flrange']=[]
	myf['sprange']=[]
	myf['linecat']='none'
	myf['linedop']=0.0
	myf['plotfile']=''
	
def sdplot_description(key='sdplot'):
	desc={'sdplot': 'ASAP SD calibration task\n',
	'sdfile': 'name of input SD dataset',
	'telescope': 'name or param of telescope for flux conversion',
	'fluxunit': 'units for line flux (K,Jy)',
	'specunit': 'units for spectral axis (channel,km/s,GHz)',
	'frame': 'frequency reference frame, e.g. LSRK (''=current)',
	'doppler': 'doppler convention, e.g. RADIO (''=current)',
        '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',
	'stack': 'code for plot stacking (p,i,b,t,s)',
	'panel': 'code for plot panelling (p,i,b,t,s)',
	'flrange': 'flux axis range [min,max] ([]=all)',
	'sprange': 'spectral axis range [min,max] ([]=all)',
	'linecat': 'line catalog control (none,all,or by molecule)',
	'linedop': 'line catalog doppler offset (km/s)',
	'plotfile': 'file name for hardcopy (''=none,.eps,.ps,.png)',
	}
	return desc[key]

###End of task