"""PyARTS is designed as a wrapper for ARTS - ie an alternative to writing a
control file for each arts run.  The most useful class is arts_run.  Here
are some examples of its use.

To perform an arts run with the default arguments (currently a montecarlo limb
sounding simulation) just go

>>> import PyARTS
>>> a=PyARTS.arts_run()

This hasn't actually started the simulation. For montecarlo runs sometimes
its a good idea to check how long it will take before you begin.

>>> a.time_estimate()
Single processor time for current arguments (in s):
1039.34969902

this is a bit long for a demonstration. Check defaults...

>>> print PyARTS.defaults["maxiter"]
100000

Either start again with a=PyARTS.arts_run({"maxiter":1000}) or go

>>> a.params["maxiter"]=1000

Note that any settings not explicitly set like this revert to the defaults.

To see the complete list of available parameters and their default values
type...

>>> print PyARTS.defaults
{'scat_za_grid': {'start': 0, 'stop': 180, 'n': 30}, 'merid_angle_1d': 0, 'p_grid': {'start': 100000.0, 'stop': 1, 'n': 100}, 'stokes_dim': 2, 'lat_grid': {'start': -30, 'stop': 30, 'n': 101}, 'gas_species': '["H2O","O3","O2","N2"]', 'rte_los': {'aa': 180, 'za': 115.59999999999999}, 'pnd_field': 1000, 'lon_grid': {'start': -14, 'stop': 19, 'n': 50}, 'lat_1d': 45, 'cloud_box': {'p2': 10000.0, 'p1': 30000.0, 'lat1': -1.5, 'lat2': 1.5, 'lon1': 1, 'lon2': 4}, 'atm_basename': '"/home/cory/data/arts-xml-data/atmosphere/fascod/tropical"', 'maxiter': 100000, 'scat_data_file': '"/home/cory/data/scat_spher_data.xml"', 'freq': 318000000000.0,
'rte_pos': {'lat': 25.600000000000001, 'r_or_z': 705000.0, 'lon': 2.5}, 'scat_aa_grid': {'start': 0, 'stop': 360, 'n': 30}}

Now to actually perform the simulation use

>>> a.run()

The screen output of the run is now stored in a.output.  However the important
results are extracted and stored in numerical arrays:  the measurement Stokes vector
in a.y, the cloud box exit radiance in a.i_cloudbox_exit and the monte carlo error
in cloud box exit radiance in a.error.

Alternatives to the run() method include:
run_in_xterm() which does what it says but does not process the output,
run_parallel(number_of_processes) which is very useful for running monte carlo simulations
on machines with more than one processor.  Basically run_parallel runs several
independent arts runs each with maxiter = argdict["maxiter"]/number_of processes
and combines the results.

If you want to be really fancy you can go a.start(), which will let you continue
to use the Python interpreter while the simulation is running.  When you are
ready you can use a.process_out_stream() to process the results.

As well as montecarlo carlo scattering simulations you can also do "DOSO"
scattering simulations, Clearsky calculations or use arts to generate data
for use in Franklin Evans SHDOM code.  These are acheived by using a second
"run_type" argument when initialising an arts_run instance.  Valid strings
are "montecarlo"(default), "DOSO", "clear", or "SHDOM".  A third optional
string argument specifies the filename for the generated arts control file.

Regardless of the run_type. The method print_results() will print the processed
numerical results to the screen.

Some more examples:

>>> near_nadir_args={'rte_los': {'aa': 180, 'za': 179.9}, 'rte_pos': {'lat': 0,
'r_or_z': 705000.0, 'lon': 2.5}, "pnd_field":1000}
>>> a=PyARTS.arts_run(near_nadir_args,"DOSO")
>>> a.run()
>>> a.print_results()
Run type:
DOSO
Stoke's vector leaving cloudbox (in RJBT):
[ 255.72548427   35.55850833]
Measured Stoke's vector (in RJBT):
[ 255.331    35.4652]


>>> a=PyARTS.arts_run(near_nadir_args,"clear")
>>> a.run()
>>> a.print_results()
Run type:
clear
Measured Stoke's vector (in RJBT):
[ 266.936    0.   ]

FIXME:
handling of filenames is user specific
consider moving stuff to a new control_file_components file

"""
from os import *
from Numeric import *
from physics import c,k
from general import *
import artsXML
from arts_file_components import *
import time
import tempfile

# Base directory for arts installation
ARTS_PATH=os.getenv('ARTS_PATH',os.path.join(os.getenv('HOME'),'arts'))
ARTS_EXEC=os.path.join(ARTS_PATH,'src/arts')

ARTSRunError='Error running ARTS'

class arts_run:
    """A class representing a single arts simulation.  The initialisation
    arguments are

    params: a dictionary of parameters.  Any "keyword":value pairs not
    specified in argdict are given default values. default parameters are
    stored in PyARTS.defaults.  Its a good idea to refer to this dictionary
    for the valid keywords.

    run_type: "montecarlo", "DOSO", "clear", "clear3D",or "SHDOM"

    filename: the name of the generated arts control file.

    Some examples for using this class are given above in the docstring for
    the PyARTS module
    
    """
    def __init__(self,params={}, run_type = "montecarlo",
                 filename=tempfile.mktemp('.arts')):
	
        self.params=dict_combine_with_default(params,defaults)
	if run_type in ["montecarlo","DOSO","SHDOM"]:
	    tan_z2sensor_poslos_thru_cloud(self.params)
	else:
	    tan_z2sensor_poslos(self.params)
        self.run_type=run_type
        self.filename=filename
        self.arts_error=0#flag indicating unsuccessful arts run
	#so far only 1D DOSO runs are supported
	if run_type=="DOSO":
	    self.params["lat_grid"]["n"]=1
	    self.params["lon_grid"]["n"]=1
	    
        
    def file_gen(self):
        """Generates the arts control file, with name self.filename, according
        to the arguments in self.params"""
        if self.run_type == "montecarlo":
            self.command_list=arts_montecarlo_file(self.params,self.filename)
        elif self.run_type == "DOSO":
            self.command_list=arts_cloudy1D_file(self.params,self.filename)
        elif self.run_type == "clear":
            self.command_list=arts_clearsky_file(self.params,self.filename)
        elif self.run_type == "clear3D":
            self.command_list=arts_clear3D_file(self.params,self.filename)
        elif self.run_type == "SHDOM":
            self.command_list=arts2SHDOM_file(self.params,self.filename)
        else:
            raise "run_type must be \"montecarlo\", \"DOSO\", \"SHDOM\",  or \"clear\""

    def start(self):
        """Start the arts run. And initiates a Python file object self.out_stream
        to handle arts output. Note that arts output is also written to file
         self.filename + ".out"
         """
        #just to be sure. Make sure control file is up to date
        self.file_gen()
#        w,self.out_stream,self.err_stream = popen3("nice "+ARTS_EXEC+" "+\
#                                                  self.filename+" | tee " \
#                                                  + self.filename + ".out")
#        w.close()
        self.out_stream = popen("nice "+ARTS_EXEC+" "+\
                                                self.filename+" | tee " \
                                                + self.filename + ".out")
        

    def process_out_stream(self):
        """Extracts important numerical values
        from self.output"""
        if self.run_type == "montecarlo":
            #return self.output
            #Get RNG seed
            w,r=popen2('grep RNG | awk \'{ print $3 }\'')
            w.write(self.output)
            w.close()
            s=r.readlines()
            self.rng_seed=int(s[0])
                
            #Get i_cloudbox_exit values
            w,r=popen2('sed -n \'/i_rte/,/-/p\' | sed \'/i_rte/d\'')
            w.write(self.output)
            w.close()
            s=r.readlines()
	    self.i_cloudbox_exit=[]
	    if len(s)==0:
		print "Warning: cloudbox exit radiance error vector not found"
	    else:
		if s[0].find("nan")==-1:
		    for i in range(len(s[0].split())):
			self.i_cloudbox_exit.append(float(s[0].split()[i]))
		else:
		    print "nan encountered in error, error not processed."
	
            #Get i_error values
            w,r=popen2('sed -n \'/i_montecarlo_error/,/-/p\' | '+\
                          'sed \'/i_montecarlo_error/d\'')
            w.write(self.output)
            w.close()
            s=r.readlines()
	    self.error=[]
	    if len(s)==0:
		print "Warning: cloudbox exit radiance error vector not found"
	    else:
		if s[0].find("nan")==-1:
		    for i in range(len(s[0].split())):
			self.error.append(float(s[0].split()[i]))				
		else:
		    print "nan encountered in error, error not processed."
		    				
	    #convert to RJ brightness temperature
            freq=self.params.get("freq",defaults["freq"])
            self.i_cloudbox_exit = array(self.i_cloudbox_exit)*c**2/2/k/freq**2
            self.error = array(self.error)*c**2/2/k/freq**2
            
        elif self.run_type == "DOSO":
            #Get i_cloudbox_exit values
            w,r=popen2('sed -n \'/i_rte/,/-/p\' | sed \'/i_rte/d\'')
            w.write(self.output)
            w.close()
            s=r.readlines()
	    self.i_cloudbox_exit=[]
	    if len(s)==0:
		print "Warning: cloudbox exit radiance error vector not found"
	    else:
		if s[0].find("nan")==-1:
		    for i in range(len(s[0].split())):
			self.i_cloudbox_exit.append(float(s[0].split()[i]))				
		else:
		    print "nan encountered in error, error not processed."
	
            #convert to RJ brightness temperature
            freq=self.params.get("freq",defaults["freq"])
            self.i_cloudbox_exit = array(self.i_cloudbox_exit)*c**2/2/k/freq**2

        if self.run_type!= "SHDOM":
            #Get y values
            w,r=popen2('sed -n \'/\*y\*/,/-/p\' | '+\
                       'sed \'/\*y\*/d\'')
            w.write(self.output)
            w.close()
            s=r.readlines()
	    self.y=[]
	    if len(s)==0:
		print "Warning: cloudbox exit radiance vector not found"
	    else:
		if s[0].find("nan")==-1:
		    for i in range(len(s[0].split())):
			self.y.append(float(s[0].split()[i]))				
		else:
		    print "nan encountered in y, y not processed."
	

    def process_output(self):
	"""Not usually called by the user - normally called by run or run_parallel
	methods
	Extract string from self.out_stream to self.output. Destroys self.outstream
        so that arts_run object can be pickled.
	Extracts Numeric information for example measured radiance from arts
	output.  For monochromatic runs it makes sense to process the stdout
	stream, but for spectra it is easier to process the xml output.  Currently
	the latter option is only valid for clearsky spectra"""
#        errmsg=self.err_stream.read()
#        if len(errmsg)>0:
#            print errmsg
#            raise ARTSRunError
#        self.err_stream.close()
	self.output = self.out_stream.read()
        #remove self.out_stream so that an arts_run can be pickled
        self.out_stream.close()
        del self.out_stream
#        del self.err_stream
	freq=self.params.get("freq",defaults["freq"])
	if type(freq)==types.DictType:	# is argdict["freq"] a dictionary?
	    self.process_out_xml()
	else:
	    self.process_out_stream()

    def process_out_xml(self):
	"""intended as a private method.  Called by process_output in cases
	where radiances are best obtained from xml output. ie in multiple
	frequency calculations"""
	fname=self.filename
	yfilename=fname[:len(fname)-4]+'y.xml'
	print "Extracting radiance data from "+yfilename
	y=artsXML.load(yfilename)[0].array
	newshape=[self.params["freq"]["n"],self.params["stokes_dim"]]
	y.shape=newshape
	self.y=y
	
    def run(self):
        """Simply calls the start() and process_out_stream() methods"""        
	
	start_time=time.time()

        self.start()
        self.process_output()

	end_time=time.time()
	self.time_elapsed=end_time-start_time
        return self
        
    def run_in_xterm(self):
        """Runs arts in an xterm with no output processing"""
        self.file_gen()
        print "arts is running in an xterm, to return to python interpreter\n" + \
              "destroy xterm."
        r = popen("xterm -exec "+ARTS_EXEC+" "+self.filename+\
                  " | tee " + self.filename + ".out  -hold")

    def print_results(self):
        """Prints results of arts simulation"""
        print "Run type:"
        print self.run_type
        try:
            if (self.run_type == "montecarlo") | (self.run_type == "DOSO"):
                print "Stoke's vector leaving cloudbox (in RJBT):"
                print self.i_cloudbox_exit
                if (self.run_type == "montecarlo"):
                    print "Error in Stoke's vector leaving cloudbox (in RJBT):"
                    print self.error
            print "Measured Stoke's vector (in RJBT):"
            print self.y
        except:
            print "arts_run.print_results: Not all output arguments present"
	print "Time elapsed (s):"
	print self.time_elapsed

    def time_estimate(self):
        """Gives an estimate of run time. It only makes sense to use this for
        run_type = \"montecarlo\". An estimate of the expected error in the
	final calculation is also given"""
        import copy
        dummy_argdict=copy.deepcopy(self.params)
        dummy_argdict["maxiter"]=1000
        dummy_arts_run=arts_run(dummy_argdict,self.run_type)
        start_time=time.time()
        dummy_arts_run.run()
        end_time=time.time()
	maxiter=float(self.params.get("maxiter",defaults["maxiter"]))
        time_estimate=(end_time-start_time) * maxiter/1000.0
        print "Single processor time for current arguments (in s): "
        print time_estimate
	print "Error estimate for proposed calculation (K):"
	print dummy_arts_run.error/sqrt(maxiter/1000.0)

    def run_parallel(self,number_of_processes):
        """Only for monte carlo simulations.  Divides self.params["maxiter"]
        by number_of_processes and runs number_of_processes arts processes.  The
        results are then combined.  This is particularly worthwhile on multiple
        processor machines."""
	stokes_dim=self.params["stokes_dim"]
        self.i_cloudbox_exit=zeros(stokes_dim,Float)
        errorsquared = zeros(stokes_dim,Float)
        self.y=zeros(stokes_dim,Float)
        #create new instances with divided maxiter
        import copy

	start_time=time.time()
	
        dummy_argdict=copy.deepcopy(self.params)
        dummy_argdict["maxiter"]=self.params.get("maxiter",defaults["maxiter"])/ \
                                  number_of_processes
        self.run_list=[]
        for i in range(number_of_processes):
            self.run_list.append(arts_run(dummy_argdict,self.run_type,
					  self.filename))
            self.run_list[i].start()
            time.sleep(3)
        for i in range(number_of_processes):
            try:
                self.run_list[i].process_output()
                self.i_cloudbox_exit+=self.run_list[i].i_cloudbox_exit
                errorsquared+=self.run_list[i].error**2
                self.y+=self.run_list[i].y
            except:
                print "Error processing ARTS output"
                self.arts_error=1
        if (not self.arts_error):
            self.i_cloudbox_exit/=number_of_processes
            self.error=errorsquared**0.5/number_of_processes
            self.y/=number_of_processes
	
	end_time=time.time()
	self.time_elapsed=end_time-start_time
        return self

def limb_spectrum(f_grid,params={}):
    """A function that plots a clear sky limb spectrum with frequencies
    given in array f_grid"""
    
    radiance_vec=[]
    for freq in f_grid:
        params["freq"]=freq
        a_run=arts_run(params,"clear")
        a_run.run()
        print "f = "+str(freq)+", y = "+str(a_run.y)
        radiance_vec.append(a_run.y)
    y_data=array(radiance_vec)
    return y_data
    
########################################################################
#Some tests for the PyARTS module

import unittest
from general import *

class ArtsRun_TestCase(unittest.TestCase):
     def runTest(self):
	 params=quickunpickle(MYCODE_PATH+"/claudia_paramsMC11km.pickle")
	 params["maxiter"]=10000
         params["strat_sampling"]=0
         params["pfbCsca_file"]=DATA_PATH+"/pfbCsca_fine_sphere_gamma_68.5.xml"
	 run1=arts_run(params)
	 run1.run_parallel(2)
	 run1.print_results()
	 assert abs((run1.y[0]-154))<2,"answer too far from previous results"
	 quickpickle(run1,DATA_PATH+"/PyARTStest_suite_results.pickle")

class ArtsRun_ClearTestCase(unittest.TestCase):
     def runTest(self):
	 params=quickunpickle(MYCODE_PATH+"/claudia_paramsMC11km.pickle")
         params.update({"t_field" : "/home/cory/data/ral_data/ral.t_field_1D.xml",
                        "vmr_field" : "/home/cory/data/ral_data/ral.vmrs_field_1D.xml",
                        "z_field" : "/home/cory/data/ral_data/ral.z_field_1D.xml"})
	 run1=arts_run(params,"clear")
	 run1.run()
	 run1.print_results()
#	 assert abs((run1.y[0]-154))<2,"answer too far from previous results"
#	 quickpickle(run1,DATA_PATH+"/PyARTStest_suite_results.pickle")

#class ArtsRun_TestCase(unittest.TestCase):

#class p20_p30SphereCompTestCase(unittest.TestCase):
#    params=quickunpickle(MYCODE_PATH+"/claudia_paramsMC11km.pickle")
#    params["maxiter"]=10000
#    params["strat_sampling"]=0
    

def test_suite():
    Suite=unittest.TestSuite()
    Suite.addTest(ArtsRun_ClearTestCase())
    Suite.addTest(ArtsRun_TestCase())
    return Suite

def run_tests():
    TestRunner=unittest.TextTestRunner()
    TestRunner.run(test_suite())