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 sdstat(sdfile=None,telescope=None,fluxunit=None,specunit=None,frame=None,doppler=None,scanlist=None,field=None,iflist=None,masklist=None,invertmask=None,xstat=None):
	"""ASAP SD spectral region statistics 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 fluxunit)
		WARNING: For GBT data, see description below.
	specunit -- units for spectral axis
	        options: 'channel','km/s','GHz','MHz','kHz','Hz',''
		default: '' (keep current specunit)
		example: make sure this is the same units of masklist
	frame -- frequency frame for spectral axis
	        options: (str) 'LSRK','REST','TOPO','LSRD','BARY',
	                 'GEO','GALACTO','LGROUP','CMB'
		default: currently set frame in scantable
	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]
	masklist -- list of mask regions to INCLUDE in stats
	        default: [] (whole spectrum)
		example: [4000,4500] for one region
		         [[1000,3000],[5000,7000]]
			 these must be pairs of [lo,hi] boundaries
	invertmask -- invert mask (EXCLUDE masklist instead)
	        options: (bool) True,False
		default: false

        -------------------------------------------------------------------
	xstat -- RETURN ONLY: a Python dictionary of line statistics
	           keys: 'rms','stddev','max','min','sum','median','mean',
		         'eqw'
		example: print "rms = ",xstat['rms']
		         these can be used for testing in scripts or
			 for regression
			 
			 'eqw' is equivalent width (sum/mag) where mag
			 is either max or min depending on which has
			 greater magnitude.


        DESCRIPTION:

	Task sdstat computes basic statistics (rms,mean,median,sum)
	for single-dish spectra.  It assumes that the spectra have
	been calibrated in sdcal.  Furthermore, it assumes that any
	selection of scans, IFs, polarizations, and time and channel
	averaging/smoothing has also already been done (in sdcal) as
	there are no controls for these.  Note that you can run sdcal
	with calmode='none' and do selection, writing out a new
	scantable.
	
	Note that multiple scans and IFs can in principle be handled, but
	we recommend that you use scanlist, field, and iflist to give a
	single selection for each run.

	WARNING: If you do have multiple scantable rows, then xstat
	values will be lists.

	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.
	
	"""
	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']='sdstat'
        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(sdstat.func_code.co_varnames,sdstat.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
        #
	# LIST PARAMETER NAMES HERE
        arg_names=['sdfile','telescope','fluxunit','specunit','frame','doppler','scanlist','field','iflist','masklist','invertmask','xstat']
	#
	# 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),list,bool,dict]
	#parameter_printvalues(arg_names,arg_values,arg_types)
        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',''])
        except TypeError, e:
                print "sdstat -- TypeError: ", e
		return
        except ValueError, e:
                print "sdstat -- OptionError: ", e
                return
        ###

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

	try:
#	    pdb.set_trace()
            #load the data without averaging
            s=sd.scantable(sdfile,False)
	    
	    # 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'

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

	    # set the mask region
	    if ( len(masklist) > 0):
		    msk=s.create_mask(masklist,invert=invertmask)
		    
		    maxl=s.stats('max',msk)
		    minl=s.stats('min',msk)
		    suml=s.stats('sum',msk)
		    meanl=s.stats('mean',msk)
		    medianl=s.stats('median',msk)
		    rmsl=s.stats('rms',msk)
		    stdvl=s.stats('stddev',msk)
		    del msk
	    else:
		    # Full region
		    print 'Using full region'
		    
		    maxl=s.stats('max')
		    minl=s.stats('min')
		    suml=s.stats('sum')
		    meanl=s.stats('mean')
		    medianl=s.stats('median')
		    rmsl=s.stats('rms')
		    stdvl=s.stats('stddev')

	    # Put into output dictionary
	    ns = len(maxl)
	    if ( ns > 1 ):
		    # User selected multiple scans,ifs
		    # put into lists
		    xstat['rms']=list(rmsl)
		    xstat['stddev']=list(stdvl)
		    xstat['max']=list(maxl)
		    xstat['min']=list(minl)
		    xstat['sum']=list(suml)
		    xstat['median']=list(medianl)
		    xstat['mean']=list(meanl)

		    eqw=[]
		    for i in range(ns):
			    if ( maxl[i] != 0.0 or minl[i] != 0.0 ):
				    if ( abs(maxl[i]) >= abs(minl[i]) ):
					    eqw = eqw + [ suml[i]/maxl[i] ]
				    else:
					    eqw = eqw + [ suml[i]/minl[i] ]
			    else:
				    eqw = eqw + [0.0]
					    
		    xstat['eqw']=eqw
	    else:
		    # Single scantable only
		    # put into scalars
		    xstat['rms']=rmsl[0]
		    xstat['stddev']=stdvl[0]
		    xstat['max']=maxl[0]
		    xstat['min']=minl[0]
		    xstat['sum']=suml[0]
		    xstat['median']=medianl[0]
		    xstat['mean']=meanl[0]

		    # Construct equivalent width (=sum/max)
		    if ( maxl[0] != 0.0 or minl[0] != 0.0 ):
			    if ( abs(maxl[0]) >= abs(minl[0]) ):
				    eqw = suml[0]/maxl[0]
			    else:
				    eqw = suml[0]/minl[0]
		    else:
			    eqw = 0.0
		    xstat['eqw']=eqw

	    # Assign variables for return
	    myf['xstat']=xstat
	    
	    # Final clean up
	    del s
	    
	    # DONE
	
	except TypeError, e:
                print "sdstat -- TypeError: ", e
		return
        except ValueError, e:
                print "sdstat -- OptionError: ", e
                return
	except Exception, instance:
		print '***Error***',instance
		return

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

###End of sdstat

def sdstat_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['masklist']=[]
	myf['invertmask']=False
	myf['xstat']={}
	
def sdstat_description(key='sdstat'):
	desc={'sdstat': 'ASAP SD spectrum statistics task\n',
	'sdfile': 'name of input SD dataset',
	'telescope': 'name or param of telescope for flux conversion',
	'fluxunit': 'units for line flux (K,Jy) (def=current)',
	'specunit': 'units for spectral axis (def=current)',
	'frame': 'frequency reference frame, e.g. LSRK (def=current)',
	'doppler': 'doppler convention, e.g. RADIO (def=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])',
        'masklist': 'list of channel [lo,hi] regions for line fitting',
	'invertmask': 'invert masklist (exclude regions) (True,False)',
	'xstat': '(RETURN ONLY) Python dictionary with line fitting output',
	}
	return desc[key]

###End of task