'''
Created on Oct 22, 2011

@author: crohr
'''

######################################################################
# standard
######################################################################
class standard:
    _labels = []
    _times = []
    _vals = []
    ######################################################################
    # read CSV
    ######################################################################
    def readCSV(self, filename, xDim, yDim):
        import csv, sys
        import numpy as np
       
        with open(filename, 'rU') as f:
            reader = csv.reader(f)
            rownum = 0
            try:
                for row in reader:
                    if(len(row) > 0):
                        if(rownum == 0):
                            self._labels = row[1::]
                        else:
                            self._times.append(row[0])
                            r = np.array(row[1::], float)
                            self._vals.append(r)
                    rownum+=1
            except csv.Error, e:
                sys.exit('file %s, line %d: %s' % (filename, reader.line_num, e))
        return
    ######################################################################
    # plot
    ######################################################################
    def plotCSV(self, xDim, yDim, pCols = 1, pRows = 1, pBegin=0, pEnd = 0):
        import numpy as np
        from matplotlib import rc
        rc('text', usetex=False)
        rc('font', family='serif')
        import matplotlib.pyplot as plt

        fig = plt.figure()
    
        if(pEnd <= 0 or pEnd > len(self._times)):
            pEnd = len(self._times)
    
        v = np.array(self._vals)
        for p in range(0,len(v[0])):
            plt.plot(self._times[pBegin:pEnd], v[pBegin:pEnd , p], '-', label = self._labels[p])

        plt.legend(loc=1)
        plt.show()
        return


######################################################################
# grid 2d
######################################################################
class grid2D:
    _labels = []
    _times = []
    _vals = []
    
    ######################################################################
    # read CSV
    ######################################################################
    def readCSV(self, filename, xDim, yDim, computeSum):
        import csv, sys
        import numpy as np
       
        with open(filename, 'rU') as f:
            reader = csv.reader(f)
            rownum = 0
            try:
                for row in reader:
                    if(rownum > 0 and len(row) > 0):
                        self._times.append(row[0])
                        r = np.zeros((xDim,yDim),float)
                        if(len(row) == (xDim*yDim+1)):
                            for i in range(xDim):
                                for j in range(yDim):
                                    r[i,j] = row[i*xDim+j+1]
                        elif(len(row) == (2*xDim*yDim+1)):
                            A = np.array(row[1::2], float)
                            B = np.array(row[2::2], float)
                            for i in range(xDim):
                                for j in range(yDim):
                                    a = A[i*xDim+j]
                                    b = B[i*xDim+j]
                                    if (computeSum == True):
                                        r[i,j] = (a + b)
                                    else:
                                        if(a>0 or b>0):
                                            r[i,j] = ( b / (a + b) )
                        else:
                            raise RuntimeError('wrong x or y dimension')
                        self._vals.append(r)
                    rownum+=1
            except csv.Error, e:
                sys.exit('file %s, line %d: %s' % (filename, reader.line_num, e))
            except IndexError, e:
                sys.exit('file %s, line %d: %s' % (filename, reader.line_num, e))
        return
    ######################################################################
    # plot
    ######################################################################
    def plotCSV(self, xDim, yDim, pCols = 1, pRows = 1, pBegin = 0, pEnd = 0, final = False):
        from matplotlib import rc
        rc('text', usetex=False)
        rc('font', family='serif')
        import matplotlib.pyplot as plt
        
        numPlots = pCols * pRows
        fig = plt.figure()
        
        if(pEnd <= 0 or pEnd > len(self._vals)):
            pEnd = len(self._vals)
        
        fStride = int((pEnd-pBegin) / numPlots)
        if (final == False):
            n = pBegin
        else:
            n = len(self._times) - 1
        
        f=1
        for i in range(pEnd):
            if(i == n and f <= numPlots):
                plt.subplot(pRows, pCols, f)
                plt.title('t='+self._times[i])
                plt.pcolormesh(self._vals[i], vmin=0, vmax=1, shading='flat')
                plt.axis([0, xDim-1, 0, yDim-1])
                cb = plt.colorbar()
                #cb.set_label('N')            
                n += fStride
                f += 1
        return plt

######################################################################
# helping class for making a recursive call on the files/directories 
######################################################################
class recursiveCall:
    ######################################################################
    # looks through the content of a directory
    ######################################################################
    def handleDirectory(self, directory, xDim, yDim, cols, rows, begin, end, final, computeSum):
	import matplotlib
	matplotlib.use('PDF')
        contents = os.listdir(directory) #contents of the current directory
        for filename in contents:
            if ((filename.endswith('.csv') == True or filename.endswith('.dat') == True) and os.path.isfile(directory + "/" + filename) == True):
                p = grid2D()
                p.readCSV(directory + "/" + filename, xDim, yDim, computeSum)
                plt = p.plotCSV(xDim, yDim, cols, rows, begin, end, final)
                plt.savefig(directory + "/" + filename + ".pdf", format='pdf')
            elif os.path.isdir(directory + "/" + filename) == True :
                if (recursive == True):
                    self.handleDirectory(self, directory + "/" + filename, xDim, yDim, cols, rows, begin, end, final)
        return

######################################################################
# main
######################################################################
import sys
import os
import optparse
from os.path import exists

if __name__=="__main__":
    parser = optparse.OptionParser("usage: %prog [options]")
    parser.add_option("-f", "--file", dest="filename",
                      #default="data.csv", 
                      type="string",
                      help="specify filename to read from")
    parser.add_option("-m", "--mode", dest="mode",
                      #default="standard",
                      help="data mode: standard, grid2D, ")
    parser.add_option("-x", "--xdim", dest="xdim", default=1,
                      type="int", help="dimension x of grid")
    parser.add_option("-y", "--ydim", dest="ydim", default=1,
                      type="int", help="dimension y of grid")
    parser.add_option("-c", "--columns", dest="cols", default=1,
                      type="int", help="number of columns to plot")
    parser.add_option("-r", "--rows", dest="rows", default=1,
                      type="int", help="number of rows to plot")
    parser.add_option("-b", "--begin", dest="begin", default=0,
                      type="int", help="number of data set to start")
    parser.add_option("-e", "--end", dest="end", default=0,
                      type="int", help="number of data set to end")
    parser.add_option("-d", "--directory", dest="dir",
                      default=None,
                      type="string",
                      help="directory in that csv files are searched for")
    parser.add_option("--recursive", dest="recursive",
                      action="store_true",
                      default=False,
                      help="search for files recursively")
    parser.add_option("-l", "--final", dest="final",
                       action="store_true",
                       default=False,
                       help="use the final snapshot only")
    parser.add_option("-s", "--sum", dest="sum",
                       action="store_true",
                       default=False,
                       help="compute the sum of A and B instead of their proportion. Only works with data mode 'grid2d'.")
    
    (options, args) = parser.parse_args()
    if len(args) > 0:
        parser.error("incorrect number of arguments")
    filename = options.filename
    mode = options.mode
    xDim = options.xdim
    yDim = options.ydim
    cols = options.cols
    rows = options.rows
    begin = options.begin
    end = options.end
    
    directory = options.dir
    recursive = options.recursive
    final = options.final
    computeSum = options.sum
    
    if (directory is None):
        if (exists(filename) is False):
            sys.exit('file %s not found' % (filename))
        
        if(mode == "standard"):
            p = standard()
            p.readCSV(filename, xDim, yDim)
            plt = p.plotCSV(xDim, yDim, cols, rows, begin, end)
            plt.show()
        elif(mode == "grid2D"):
            p = grid2D()
            p.readCSV(filename, xDim, yDim, computeSum)
            plt = p.plotCSV(xDim, yDim, cols, rows, begin, end, final)
            plt.show()
        else:
            sys.exit('wrong data mode')
    else:
        if (exists(directory) is False):
            sys.exit('directory %s not found' % (filename))
        
        call = recursiveCall()
        call.handleDirectory(directory, xDim, yDim, cols, rows, begin, end, final, computeSum)