#############################################################
#copyright CSIRO
#ben baldwinson
#
# script provided as-is with no warranty
# not for distribution to any third party.
#
#python script for disdrometer data
#
#takes input of disdrometer data file (contains "rd" file, 
    #no need to read "drops" file)
#function:
#    expands ODM 470 data file and outputs dat file
#    of rain rate for every minute in 24hrs
#    then creates rain rate, culmative rain graph, 
#    rain particle distribution for the highest rain period for that 
#    day,
#    and waterfall chart
#
#some detail:
#
#note that drop sizes less than 0.35mm or class 12 and under are ignored, so 
#will have 0 value (ie not recorded)
#it is noted in the datasheet as a requirement for use on ships as 
#vibration can be mis-interpreted for drops
#   
#the data file contains minute aggregated data. ie sum of drops over last minute
#if there are no particles for any particular minute, then no data is recorded.
#
#postprocessing is normally done in hamburg, germany..
#and normally requires data for air temp, humidity, and ship navigation
#
##########################################################3

from matplotlib import pyplot as plt
from matplotlib import dates
import numpy as np
import os

#from matplotlib.mlab import bivariate_normal
import time
import math

#functions
def fileGetStr(fileObj,delimeter,delimeter2):
    temp=fileObj.read(1) #1st char
    #check if delimiter
    if ((temp==delimeter) or (temp==delimeter2)):
        return
    else:
        temp2=fileObj.read(1) #2nd char
        if ((temp2==delimeter) or (temp==delimeter2)):
            return str(temp)
        else:
            tempStr=temp + temp2
            temp=fileObj.read(1)
            while ((temp!=delimeter) and (temp!=delimeter2)): #next char
                tempStr+=temp;
                temp=fileObj.read(1)
    return str(tempStr)

#parameters
#details of disdrometer
disdroDiam=21.0 #[mm]
disdroLength=120.0 #[mm]
disdroVol=(disdroDiam/2000) ** 2*math.pi*disdroDiam/1000 #[m^3]
disdroArea=disdroDiam*disdroLength #[mm^2]

#main
print '***Script for disdrometer***'

#get disdrometer size bins
fo_bins = open("bins_ave.csv",'r')

#clear first title line
dummy=fo_bins.readline()

bins=[0.0 for i in range(128)]
binVol=[0.0 for i in range(128)]
binMin=[0.0 for i in range(128)]
binMax=[0.0 for i in range(128)]

for i in range(0,128):
    dummy=int(fileGetStr(fo_bins,',','\n')) # (don't need bin numbers)
    binMin[i]=float(fileGetStr(fo_bins,',','\n')) # (don't need min size)
    binMax[i]=float(fileGetStr(fo_bins,',','\n')) # (don't need max size)
    bins[i]=float(fileGetStr(fo_bins,',','\n')) # (need bin ave size)
    binVol[i]=(bins[i]/2) ** 3 * math.pi * 4/3000 #assume that particle is sphere

#close bin file    
fo_bins.close()

#get disdrometer data file name

#fileNameIn = raw_input("Enter disdrometer data file name:")
#fileNameIn = "ODM0931683_145_rd_20160325"
allFileList=os.listdir('.')

for i in range(len(allFileList)):
    
    if (allFileList[i].find("ODM")!=-1) and (allFileList[i].find("rd")!=-1) and (len(allFileList[i])==26):
        fileNameIn=allFileList[i];
        fo_IN = open(fileNameIn,"r")
        print "translating disdrometer data from " + fo_IN.name + ".."
        
        lineNum=0
        #get first line
        dummy=fo_IN.readline() #read a whole line
        
        #work out how many lines
        while dummy!='':
            dummy=fo_IN.readline() #read a whole line
            lineNum+=1
        
        #array size is the number of lines divided by 8
            # (this is how many lines per disdrometer entry)
        totalEntries=int(lineNum//8) #floor division
        
        #reset line num
        lineNum=0
        
        #make empty arrays for the lines to be entered.
        dateTimeStr=['0' for i in range(totalEntries)] #string
        refList=[0.0 for i in range(totalEntries)] #float
        windList=[0.0 for i in range(totalEntries)] #float
        timeArray=[0.0 for i in range(totalEntries)] #float
        snowBinsList=[[0 for j in range(128)] for i in range(totalEntries)] #int
        rainBinsList=[[0 for j in range(128)] for i in range(totalEntries)] #int
        
        #reset to start of file
        #close files    
        fo_IN.close()
        fo_IN = open(fileNameIn,"r")
        
        for i in range(0,totalEntries):    
            
            #first line
            dateTimeStr[i] = fo_IN.read(15) #read the date and time string
            dummy=fo_IN.read(1) #read whitespace
            refList[i] = float(fo_IN.read(4)) #read the reference voltage
            dummy=fo_IN.read(1) #read whitespace
            windList[i] = float(fo_IN.read(5))
           
            dummy=fo_IN.read(1) #read whitespace
            snowNum = int(fo_IN.read(4)) #this is the total number of particles of snow of this entry 
            dummy=fo_IN.read(1) #read whitespace
            snowNumBins = int(fo_IN.read(3)) #this is the number of bins of snow of this entry
            dummy=fo_IN.read(1) #read whitespace
            rainNum = int(fo_IN.read(4)) #this is the total number of particles of rain for this entry
            dummy=fo_IN.read(1) #read whitespace
            rainNumBins = int(fo_IN.read(3)) #this is the number of bins for this entry 
            
            #second line
            #start some blank arrays for this line
            tempSnowClassList=[0 for j in range(snowNumBins)] #int
            tempSnowNumList=[0 for j in range(snowNumBins)] #int
            tempRainClassList=[0 for j in range(rainNumBins)] #int
            tempRainNumList=[0 for j in range(rainNumBins)] #int
        
            #SNOW
            #retrieve next line, (do need this line)
            #get "CS", class list for snow
            dummy=fo_IN.read(3) #if this is not CS then error?
            #know that we only need to read a certain number of classes
            for j in range(0,snowNumBins):
                dummy=fo_IN.read(1) #read whitespace
                tempSnowClassList[j]=int(fo_IN.read(3))
            
            #retrieve next line, (do need this line)
            #get "DS"
            dummy=fo_IN.read(4) #if this is not DS then error?
            #know that we only need to read a certain number of classes
            for j in range(0,snowNumBins):
                dummy=fo_IN.read(1) #read whitespace
                tempSnowNumList[j]=int(fo_IN.read(3))
            
            #retrieve next line (don't need this line)
            #get "TS"
            dummy=fo_IN.read(4) #if this is not TS then error?
            #know that we only need to read a certain number of classes
            for j in range(0,snowNumBins):
                dummy=fo_IN.read(1) #read whitespace
                dummy=int(fo_IN.read(6))
                
            #RAIN    
            #retrieve next line, (do need this line)
            #get "CR"
            dummy=fo_IN.read(4) #if this is not CR then error?
            #know that we only need to read a certain number of classes
            for j in range(0,rainNumBins):
                dummy=fo_IN.read(1) #read whitespace
                tempRainClassList[j]=int(fo_IN.read(3))
            
            #retrieve next line, (do need this line)
            #get "DR"
            dummy=fo_IN.read(4) #if this is not DR then error?
            #know that we only need to read a certain number of classes
            for j in range(0,rainNumBins):
                dummy=fo_IN.read(1) #read whitespace
                tempRainNumList[j]=int(fo_IN.read(3))
            
            #retrieve next line (don't need this line)
            #get "TR"
            dummy=fo_IN.read(4) #if this is not TR then error?
            #know that we only need to read a certain number of classes
            for j in range(0,rainNumBins):
                dummy=fo_IN.read(1) #read whitespace
                dummy=int(fo_IN.read(6))
             
            #convert the time string to tupleTime
            tupleTimeTemp=time.strptime(str(dateTimeStr[i]), "%d%m%Y %H%M%S")
            timeArray[i]=time.mktime(tupleTimeTemp)
            
            #write SNOW
            #now write the values to the bins
            k=0
            for j in range(0,128):
                if k>=snowNumBins: #this stops overindexing into the bins
                    snowBinsList[i][j]=0
                elif j==tempSnowClassList[k]:
                    snowBinsList[i][j]=tempSnowNumList[k]
                    k+=1
                else:
                    snowBinsList[i][j]=0
            
            #RAIN
            #now write the values to the bins
            k=0
            for j in range(0,128):
                if k>=rainNumBins:
                    rainBinsList[i][j]=0
                elif j==tempRainClassList[k]:
                    rainBinsList[i][j]=tempRainNumList[k]
                    k+=1
                else:
                    rainBinsList[i][j]=0
        
            #characters for newline
            dummy=fo_IN.read(3)
                        
            #now can go to next set    
        #end (for)
        
        #close files    
        fo_IN.close()
        
        #now can make 30minute and/or running totals across the file time range
        aveTime=1 #5 #30 #[minutes] sums rain bin nums over this time
        timeIncrement=1 #5 #30 #[minutes] this is the increment in time for the array,
                            #and also the amount of time for distribution graph
        fileTotalTime=timeArray[len(timeArray)-1]-timeArray[0] #this is the total number 
            #of seconds from the first entry to the last entry, should be 24hrs..
        
        #make empty result array
        snowBinsFinalList=[[0 for j in range(128)] for i in range((24*60)/timeIncrement)]
        rainBinsFinalList=[[0 for j in range(128)] for i in range((24*60)/timeIncrement)]
        
        #make empty vol array
        snowVolArray=[0.0 for i in range((24*60)/timeIncrement)]
        rainVolArray=[0.0 for i in range((24*60)/timeIncrement)]
        timeArrayFinal=[0.0 for i in range((24*60)/timeIncrement)]
        refListFinal=[0.0 for i in range((24*60)/timeIncrement)]
        #make new time array
        #start at midnight0:00 (start) of the day listed in the filename.
        startTimeStr=fileNameIn[18:26]+' 000000' #eg "ODM0931683_145_rd_20160324"
        #convert the time string to tupleTime
        tupleTimeTemp1=time.strptime(str(startTimeStr), "%Y%m%d %H%M%S")
        #convert to seconds since epoch    
        timeArrayFinal[0]=time.mktime(tupleTimeTemp1)
    
        #fill in rest of 24 hrs
        for i in range(1,(24*60)/timeIncrement): 
            #(at the moment the above must result in whole numbers)
            
            #make time array for this output
            timeArrayFinal[i]=timeArrayFinal[i-1]+timeIncrement*60
        #end (for)
        
        #make new array, for each interval from above over an assumed 24 hours
        for i in range(0,(24*60)/timeIncrement): 
            #(at the moment the above must result in whole numbers)
            
            #search the time array for entries in the previous aveTime interval, and 
            #add the bin counts for those times
            for j in range(0,len(timeArray)):
                #if the times are within the last aveTime interval then add them to 
                #the arraySums
                if ((timeArray[j]>(timeArrayFinal[i]-aveTime*60))&(timeArray[j]<timeArrayFinal[i])):
                    for k in range(0,128): 
                        tempInteger=snowBinsList[j][k]
                        snowBinsFinalList[i][k]+=tempInteger
                    
                        tempInteger=rainBinsList[j][k]
                        rainBinsFinalList[i][k]+=tempInteger
              
                        refListFinal[i]=refList[j]
                
            #calc snow volume for the interval
            snowVolArrayTemp2=0
            for k in range(0,128):
                snowVolArrayTemp2+=snowBinsFinalList[i][k]*binVol[k]
            snowVolArray[i]=snowVolArrayTemp2 #saves snow volume in [mL]
        
            #calc rain volume for the interval
            rainVolArrayTemp2=0
            for k in range(0,128):
                rainVolArrayTemp2+=rainBinsFinalList[i][k]*binVol[k]
            rainVolArray[i]=rainVolArrayTemp2 #saves rain volume in [mL]
        
            
        #end (for)
        
        #work out time array for display.
        timeArrayFinalPlot=[0.0 for i in range((24*60)/timeIncrement)]
                
        for i in range(0,len(timeArrayFinal)):
            timeArrayFinalPlot[i]=dates.epoch2num(timeArrayFinal[i]-13*60*60) #adjusted by 10hrs
        
        #new blank rainRate array
        rainRate=[0.0 for i in range(len(rainVolArray))] #float
        
        #get rain Rate array
        for i in range(0,len(rainVolArray)):
            rainRate[i]=(rainVolArray[i]*1000/disdroArea)/(float(aveTime)/60) #this gives mm of rain per hour
        
        #new blank culmative array
        culmRain=[0.0 for i in range(len(rainRate))] #float
        
        culmRain[0]=0.0
        for i in range(1,len(rainVolArray)):
            culmRain[i]=culmRain[i-1]+rainRate[i]*(float(aveTime)/60)
            
        #plot the 30min rain/snow sum
        fig1=plt.figure(1)
        ax1 = fig1.add_subplot(111)
        plt.plot(timeArrayFinalPlot,rainRate, color='blue') #color='red')
        plt.plot(timeArrayFinalPlot,culmRain, color='red')
        plt.xlabel('Time of day UTC [hr]', fontsize=10, color='black')
        plt.ylabel('Rain Rate [mm/hr] (blue), culmative [mm] (red)', fontsize=10, color='black')
        plt.title('Rain Rate from UTC '+time.ctime(timeArrayFinal[0]) + ' to ' + time.ctime(timeArrayFinal[len(timeArrayFinal)-1]), fontsize=10 )
        yChartmax=max([(max(culmRain)+1),(max(rainRate)+1)])
        plt.axis([min(timeArrayFinalPlot), max(timeArrayFinalPlot), 0,yChartmax ])
        plt.grid(True)
        ax1.xaxis.set_major_locator(dates.HourLocator())
        hfmt = dates.DateFormatter('%H:%M')
        ax1.xaxis.set_major_formatter(hfmt)
        plt.xticks(rotation='vertical')
        plt.subplots_adjust(bottom=.15)
        #plt.show()
        
        #find the time of the maximum rain event
        maxVol=0.0
        maxVolIndex=0
        for i in range(0,len(snowVolArray)):
            if snowVolArray[i]>maxVol:
                maxVol=snowVolArray[i]
                maxVolIndex=i
        
        #plot distribution for the 30minutes with the max rain
        binNum=int(maxVolIndex)
        print 'printing distrbution of line=' + str(binNum)
        
        #calc the distribution
        rainParticlepm=[0 for i in range(len(rainBinsFinalList[binNum]))]
        
        #get rain Rate array
        for i in range(1,len(rainParticlepm)):
            rainParticlepm[i]=rainBinsFinalList[binNum][i]*bins[i]/1000/disdroVol #gives #drops/m^3/m
        
        #plot a rain/snow histogram sum
        fig2=plt.figure(2)
        ax2 = fig2.add_subplot(111)
        for i in range(0,128):
            plt.bar(bins[i],rainParticlepm[i], color='black',width=(binMax[i]-binMin[i])) #color='red')
        plt.xlabel('Particle Diameter [mm]', fontsize=10, color='black')
        plt.ylabel('Number of Particles per m^3 * Diameter', fontsize=10, color='black')
        plt.title('Rain Particle Distribution over '+ str(timeIncrement)+ ' minutes at '+ time.ctime(timeArrayFinal[binNum]), fontsize=10)
        plt.xscale('log') #log scale of the bin size
        plt.axis([0.35, 20, 0, max(rainParticlepm)+10])
        ticksXaxis=[0.4,0.5,0.6,0.7,0.8,0.9,1,2,3,4,5,6,7,8,9,10]
        ax2.set_xticks(ticksXaxis)
        plt.grid(True)
        
        #output rainRate to file
        fileNameOut = fileNameIn + "rainRate 1min.dat"
        fo_OUT = open(fileNameOut,"w")
        print "writing out disdrometer data to " + fo_OUT.name + ".."
        
        for i in range(0,len(timeArrayFinalPlot)): 
            fo_OUT.write(str(timeArrayFinal[i])) #send time since epoch [s]
            fo_OUT.write(',')
            fo_OUT.write(str(rainRate[i])) #send rainRate
            fo_OUT.write(',')
            fo_OUT.write(str(culmRain[i])) #send rainRate
            fo_OUT.write('\n')
        
        #close output file
        fo_OUT.close()
        
        print 'plotting waterfall..'
        #attempt to make waterfall plot
        fig3=plt.figure(3)
        ax3 = fig3.add_subplot(111)
               
        timeArrayFinalPlotA=[[0 for j in range(len(bins))] for i in range(len(timeArrayFinalPlot))]
        #make 3D array for pcolor plot
        for i in range(0,len(timeArrayFinalPlot)):
            for j in range(0,len(bins)):
                timeArrayFinalPlotA[i][j]=timeArrayFinalPlot[i]
        
        
        binsA=[[0 for j in range(len(bins))] for i in range(len(timeArrayFinalPlot))]
        #make 3D array for pcolor plot
        for i in range(0,len(bins)):
            for j in range(0,len(timeArrayFinalPlot)):
                binsA[j][i]=bins[i]
        
        binsA=np.array(binsA)
        timeArrayFinalPlotA=np.array(timeArrayFinalPlotA)
        rainBinsFinalListA=np.array(rainBinsFinalList)
        pcm=ax3.pcolormesh(timeArrayFinalPlotA, binsA, rainBinsFinalListA, vmin=0, vmax=100,cmap='cool')
        
        plt.xlabel('Time of day UTC [hr]', fontsize=10, color='black')
        plt.ylabel('Bin size [mm]', fontsize=10, color='black')
        plt.title('Rain Particle distribution vs time from UTC '+time.ctime(timeArrayFinal[0]) + ' to ' + time.ctime(timeArrayFinal[len(timeArrayFinal)-1]), fontsize=10 )
        
        #format x axis for time
        ax3.xaxis.set_major_locator(dates.HourLocator())
        hfmt = dates.DateFormatter('%H:%M')
        ax3.xaxis.set_major_formatter(hfmt)
        plt.xticks(rotation='vertical')
        plt.subplots_adjust(bottom=.15)
        
        fig3.colorbar(pcm, ax=ax3, extend='max')
        plt.yscale('log') #log scale of the bin size
        plt.axis([min(timeArrayFinalPlot), max(timeArrayFinalPlot), bins[14],20])
        plt.grid(True)
        ticksYaxis=[0.4,0.5,0.6,0.7,0.8,0.9,1,2,3,4,5,6,7,8,9,10]
        ax3.set_yticks(ticksYaxis)
        
        #translate original time array to date format for plotting
        timeArrayOrigPlot=[0.0 for i in range(len(timeArray))]
        for i in range(0,len(timeArray)):
            timeArrayOrigPlot[i]=dates.epoch2num(timeArray[i]-13*60*60) #adjusted by 10hrs
        
        fig4=plt.figure(4)
        ax4 = fig4.add_subplot(111)
        plt.plot(timeArrayOrigPlot,refList, color='blue') #color='red')
        plt.xlabel('Time of day UTC [hr]', fontsize=10, color='black')
        plt.ylabel('reference [V]', fontsize=10, color='black')
        plt.title('Reference voltage for day starting ' +time.ctime(timeArrayFinal[0]), fontsize=10 )
        plt.grid(True)
        
        #format x axis for time
        ax4.xaxis.set_major_locator(dates.HourLocator())
        hfmt = dates.DateFormatter('%H:%M')
        ax4.xaxis.set_major_formatter(hfmt)
        plt.xticks(rotation='vertical')
        plt.subplots_adjust(bottom=.15)
        
        #save images of plot to file
        fig1.savefig(fileNameIn+"rainRate.png", dpi=300) 
        fig2.savefig(fileNameIn+"distribution.png", dpi=300) 
        fig3.savefig(fileNameIn+"waterfall.png", dpi=300)
        fig4.savefig(fileNameIn+"reference.png", dpi=300)
        
        #plt.show() #have to do this after saving?
        
        #close figures so fresh for next time
        plt.close(fig1)
        plt.close(fig2)
        plt.close(fig3)
        plt.close(fig4)
        
#for debug
print 'Completed.'