# ******************************************************
## Revision "$LastChangedDate: 2013-12-16 12:37:41 +0100 (Mon, 16 Dec 2013) $"
## Date "$LastChangedRevision: 427 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/generalcode/trunk/mouth.py $"
# ******************************************************

import write_dict
import ascraster
import my_sys

class Mouth:
    '''General functions for the classes NMouth and PMouth'''
    def set_index_grid(self,grid):
        if grid.masked:
            # Calculate index in an unmasked grid
            irow = grid.nrows - int((self.lat-float(grid.yllcorner))/float(grid.cellsize)) - 1
            icol = int((self.lon-float(grid.xllcorner))/float(grid.cellsize))
            ind = irow * int(grid.ncols) + icol
            try:
                self.index_grid = grid.mask.index(ind)
            except ValueError:
                # No grid cell found in the mask with this coordinates.
                print "No grid cell found in the mask with the coordinates "+ str(self.lat) + " " + str(self.lon)
                self.index_grid = -1
        else:
            self.index_grid = grid.get_index_from_coordin(self.lon,self.lat)
    
    def get_index_grid(self,grid):
        if grid.masked:
            return self.index_grid
        else:    
            return self.index_grid_full

    def weigh_avg(self,m,avg_attr,weigh_attr):
        '''
        Calculate the weighed average of two Mouth objects. avg_attr and weigh_attr must be attributes of Mouth.
        '''
        try:
            sum_avg = getattr(self,avg_attr) * getattr(self,weigh_attr) + getattr(m,avg_attr) * getattr(m,weigh_attr) 
            sum_weigh = getattr(self,weigh_attr) + getattr(m,weigh_attr) 
            setattr(self,avg_attr,sum_avg/sum_weigh)
        except ZeroDivisionError:
            setattr(self,avg_attr,0.0)
 
def read_file(filename,sep,my_object,keyitem="basinid"):
    '''
    Read the file and put everything in a dictionary with int(basinid) as key and a instance of a class (PMouth or NMouth) as value 
    '''
    dict = {}
    # Open file and read all the lines    
    lines = my_sys.my_readfile(filename,my_sys_ldebug=0,fatal=1)
    
    # The header gives the name of the column, which is the same as the item of the class.
    header = lines[0].split(sep)

    # Find key:
    ikey = -1
    for item in range(len(header)):
        if (header[item] == keyitem):
            ikey = item
    print ikey, header[ikey]

    for line in range(1,len(lines)):
        fields = lines[line].split(sep)
        try:
            key = int(fields[ikey])
        except ValueError:
            key = fields[ikey]
 
        dict[key] = my_object.copy()
        # Put key element in the class
        setattr(dict[key],keyitem,key)
        for item in range(min(len(fields),len(header))):
            if (item != ikey):
                # Put all the values in the Mouth object
                try:
                    dum = int(fields[item])
                    setattr(dict[key],header[item],int(fields[item]))
                except ValueError:
                    try:
                        dum = float(fields[item])
                        setattr(dict[key],header[item],float(fields[item]))
                    except ValueError: 
                        setattr(dict[key],header[item],str(fields[item]))
            
    return dict

def aggregate(dict,name_id_grid,mask,filename,sep,keyname="basinid"):    
    '''
    Aggreggate a dictionary of Mouth objects to another id grid and write the aggregation to output file.
    '''
    
    # Make a copy of the dictionary for the output
    new_dict={}

    # Read id_grid
    id_grid = ascraster.Asciigrid(ascii_file=name_id_grid,mask=mask,numtype=int)

    for key in dict.keys():
        ind = dict[key].index_grid
        if (ind == None):
            # There is no spatial location in isogrid found for this class.
            id = -9999
        else:
            id = int(id_grid.get_data(ind,-9999))
        try:
            new_dict[id].sum(dict[key])
        except KeyError:
            new_dict[id] = dict[key].copy()

    # Make a correction on the name
    for key in new_dict.keys():
        setattr(new_dict[key],keyname,key)
    
    # Aggregate to global level
    for key in new_dict.keys():
        try:
            new_dict["Total"].sum(new_dict[key])
        except KeyError:
            new_dict["Total"] = new_dict[key].copy()

    # Write all information of the river basin to output file:    
    fp = open(filename,"w")
    for key in new_dict.keys():
        new_dict[key].write_header(fp,sep=sep)
        break
    write_dict.write_dict(fp_in=fp,filename=None,dic=new_dict,headerkey="",sep=sep,lkey=0,ldebug=1,fatal=0)
    fp.close()
