#   modified mexican hat wavelet test.py
#   spectral analysis for RADAR and WRF patterns
import os, shutil
import time
import pickle
import numpy as np
from scipy import signal, ndimage
import matplotlib.pyplot as plt
from armor import defaultParameters as dp
from armor import pattern
from armor import objects4 as ob
#from armor import misc as ms
dbz = pattern.DBZ
testScriptsFolder = dp.root + 'python/armor/tests/'
testName    = "modifiedMexicanHatTest"
timeString  = str(int(time.time()))
outputFolder = dp.root + 'labLogs/%d-%d-%d-%s/' % \
                (time.localtime().tm_year, time.localtime().tm_mon, time.localtime().tm_mday, testName)
if not os.path.exists(outputFolder):
    os.makedirs(outputFolder)
shutil.copyfile(testScriptsFolder+testName+".py", outputFolder+ timeString + testName+".py")

kongreywrf  = ob.kongreywrf
kongreywrf.fix()
kongrey     = ob.kongrey
monsoon     = ob.monsoon
monsoon.list= [v for v in monsoon.list if '20120612' in v.dataTime]
march2014   = ob.march2014
march2014wrf11  = ob.march2014wrf11
march2014wrf12  = ob.march2014wrf12

summaryFile = open(outputFolder + timeString + "summary.txt", 'a')

for ds in [kongrey, kongreywrf, march2014wrf11, march2014wrf12, march2014, monsoon, ]:
    summaryFile.write("\n===============================================================\n\n\n")
    streamMean  = 0.
    dbzCount    = 0
    for a in ds:
        print "-------------------------------------------------"
        print testName
        print
        print a.name
        a.load()
        a.setThreshold(0)
        a.saveImage(imagePath=outputFolder+a.name+".png")

        L   = []
        for sigma in [1, 2, 4, 8 ,16, 32, 64, 128, 256, 512]:
            print "sigma:", sigma
            a.load()
            a.setThreshold(0)
            arr0 = a.matrix
            #arr1 = signal.convolve2d(arr0, mask_i, mode='same', boundary='fill')
            arr1    = ndimage.filters.gaussian_laplace(arr0, sigma=sigma, mode="constant", cval=0.0) * sigma**2 #2014-04-29
            a1 = dbz(matrix=arr1.real, name=a.name + "_" + testName + "_sigma" + str(sigma))
            L.append({  'sigma'     : sigma,
                        'a1'   :  a1,
                        'abssum1': abs(a1.matrix).sum(),
                        'sum1'  : a1.matrix.sum(),
                      }) 
            print "abs sum", abs(a1.matrix.sum())
            #a1.show()
            #a2.show()
            plt.close()
            a1.histogram(display=False, outputPath=outputFolder+a1.name+"_histogram.png")
        #pickle.dump(L, open(outputFolder+ a.name +'_test_results.pydump','w'))     # no need to dump if test is easy
        x = [v['sigma'] for v in L]
        y1 = [v['abssum1'] for v in L]
        plt.close()
        plt.plot(x,y1)
        plt.title(a1.name+ '\n absolute values against sigma')
        plt.savefig(outputFolder+a1.name+"-spectrum-histogram.png")
        plt.close()

        #   now update the mean
        streamMeanUpdate = np.array([v['abssum1'] for v in L])
        dbzCount    += 1
        streamMean  = 1.* ((streamMean*(dbzCount -1)) + streamMeanUpdate ) / dbzCount
        sigmas      =[v['sigma'] for v in L]
        print "Stream Count and Mean so far:", dbzCount, streamMean
        #   now save the mean and the plot
        summaryText = '\n---------------------------------------\n'
        summaryText += str(int(time.time())) + '\n'
        summaryText += "dbzStream Name:" + ds.name + '\n'
        summaryText += "dbzCount:\t" + str(dbzCount) + '\n'
        summaryText +="sigma:\t\t" + str(sigmas) + '\n'
        summaryText += "streamMean:\t" + str(streamMean.tolist()) +'\n'
        print summaryText
        print "saving..."
        #   release the memory

        a.matrix = np.array([0])
        summaryFile.write(summaryText)
        plt.close()
        plt.plot(sigmas, streamMean)
        plt.title(ds.name + '- average laplacian-of-gaussian spectrum for ' +str(dbzCount) + ' DBZ patterns')
        plt.savefig(outputFolder + ds.name + "_average_LoG_spectrum.png")
        plt.close()
        logText = "sigmas:\n    " + str(sigmas)
        logText+= "\nstreamMean of absolute values of Laplacian-of-gaussian filter responses\n    "
        logText+= str(streamMean) + '\n'
        open(outputFolder+ds.name+'_average_LoG_spectrum.txt').write(logText)
summaryFile.close()

