"""Utilities to ease handling of NetCDF files."""
import sys
import rgbimg
import os
import tempfile

from Scientific.IO.NetCDF import *
import Numeric as Nm




def view(inp):
    """Given a 2D Numeric Array, launch a viewer to display the array
    as a grayscale image."""
    
    viewer=os.getenv("IMGVIEWER","display")
    if (Nm.rank(inp)!=2):
        raise("Can only view 2-d arrays!")
    ar=Nm.transpose(inp)
    # Linearly scale input range to 0-256, and store as unsigned bytes
    mn=Nm.minimum.reduce(Nm.ravel(ar))
    mx=Nm.maximum.reduce(Nm.ravel(ar))
    nr=((Nm.ravel(ar)-mn)/(mx-mn)*255).astype('b')
    # Convert to 32-bit grayscale image. 
    u=nr*Nm.array(0x01000000,'u')+nr*0x00010100
    nm=tempfile.mktemp('.rgb')
    rgbimg.longstoimage((u).tostring(),ar.shape[1],ar.shape[0],1,nm)
    os.spawnlp(os.P_NOWAIT,viewer,"",nm)

    

class NCFile:
    def __init__(self,fname,rename_map=None):
        """Open a NetCDF file for reading. rename_map is an optional
        dictionary mapping netcdf variables to 'canonical' names, so that
        variables can be accessed via the canonical names. If a variable
        name does not appear as key in rename_map, then it can be accessed
        via its original name.
        """
        self.f=NetCDFFile(fname)
        self.datas={} # Cache for previously read datasets.
        self.strs={}
        if rename_map==None: rename_map={}
        # Fill in the inverse name lookup table.
        self.canon_to_local={'time':'time'}
        for var in self.f.variables:
            canon=rename_map.get(var,var)
            self.canon_to_local[canon]=var
        # Compute number of timesteps from the length of the 'time' array
        if self.canon_to_local['time'] in self.f.dimensions:
            self.num_timesteps=len(self.f.variables[self.canon_to_local['time']])
        else:
            # We use num_timesteps=None for 3D datasets to distinguish it from
            # 4D datasets with a single timestep
            self.num_timesteps=None
        self.timestep=0


    def __del__(self):
        self.f.close()
        
    def __getitem__(self,index):
        localname=self.canon_to_local[index]
        if self.datas.has_key(localname):
            # Retrieve from cache.
            return self.datas[localname]
        if self.f.variables.has_key(localname):
            # Read in variable from NetCDF file as a Numeric array.
            # If the variable has time as its first dimension, pick the right
            # timestep:
            ncvar=self.f.variables[localname]
            if self.canon_to_local['time'] == ncvar.dimensions[0]:
                result=self.f.variables[localname][self.timestep]
            else:
                result=self.f.variables[localname].getValue()
            self.datas[localname]=result
            return result
        raise("File contains no such variable: %s." %localname)

    def set_timestep(self,new_ts):
        if self.num_timesteps is None:
            if new_ts!=0:
                print >>sys.stderr,"Can't set timestep on 3D dataset. Ignoring."
            return
        if new_ts<0 or new_ts>=self.num_timesteps:
            print >>sys.stderr,"Timestep out of range. Clamping"
        self.timestep=min(max(new_ts,0),self.num_timesteps-1)
    
    def missing_value(self,varname):
        """ Returns the value of the missing_value of a variable if it exists,
        or None otherwise.
        """
        
        localname=self.canon_to_local[varname]
        if self.f.variables.has_key(localname):
            v=self.f.variables[localname]
            if hasattr(v,"missing_value"):
                return v.missing_value
        return None
        

    def parse_expression(self,expr,other_symbols={}):
        """Compute the result of an arithmetic expression involving
        variables from the NetCDF file, as well as any items passed in
        via other_symbols. All variables and other_symbols values must be
        convertible to Numeric arrays of the same shape. As we're using
        pythons eval here, this method should not be used with an expression
        from an untrusted source.
        """
        # Create a dictionary of all needed symbols
        arr_data=other_symbols.copy()
        # Only read in NetCDF variables which actually occur in the expression,
        # in order to avoid unneccessary disk access.
        for key in self.canon_to_local.keys():
            if expr.find(key)>=0:
                arr_data[key]=self[key]
        # Invoke python's interpreter to evaluate the expression. We
        # pass Numeric.__dict__ as the globals argument so that the
        # expression can contain Numeric functions (such as where, or identity)
        return eval(expr,Nm.__dict__,arr_data)