#!/usr/bin/env python

""" This is a first prototype of MAUD script to filter a NetCDF file

    Far away from an ideal solution. It's just to resolve the Bia's
      problem for now.
"""

from optparse import OptionParser

from numpy import ma
from netCDF4 import Dataset

from maud import window_1Dbandpass


# ==== Parsing the options on command line
parser = OptionParser()

parser.add_option("--lowpasswindowlength", dest="lowpasswindowlength",
    help="Length of the filter window. Must be on the same scale of the scalevar")

parser.add_option("--highpasswindowlength", dest="highpasswindowlength",
    help="Length of the filter window. Must be on the same scale of the scalevar")

parser.add_option("--scalevar", dest="scalevar",
    help="The scale on the dimension to be filtered.", default="time")

parser.add_option("--var", dest="var",
    help="Variable to be filtered")

parser.add_option("-w", dest="windowmethod",
    help="Type of window [hamming, hann, boxcar, triangle, lanczos]",
    default="hamming")


(options, args) = parser.parse_args()

options.lowpasswindowlength = float(options.lowpasswindowlength)
options.highpasswindowlength = float(options.highpasswindowlength)
options.varout  = options.var + "_maud"

nc = Dataset(args[0], 'a')

#if var not in nc.variables.keys():
#    import sys; sys.exit()


attributes = nc.variables[options.var].ncattrs()
try:
    out = nc.createVariable(options.varout, 
              nc.variables[options.var].dtype, 
              nc.variables[options.var].dimensions, 
              fill_value=nc.variables[options.var]._FillValue)
    attributes.remove('_FillValue')
except:
    out = nc.createVariable(options.varout, 
              nc.variables[options.var].dtype, 
              nc.variables[options.var].dimensions)

for a in attributes:
    setattr(out, a, getattr(nc.variables[options.var], a))

#window_1Dmean(data, l, t=None, method='hann', axis=0, parallel=True)

axis = nc.variables[options.var].dimensions.index(nc.variables[options.scalevar].dimensions[0])

assert axis==0
assert len(nc.variables[options.var].dimensions)==3
    
I, J = nc.variables[options.var].shape[1:]
try:
    from progressbar import ProgressBar
    pbar = ProgressBar(maxval=I*J).start()
except:
    pass

data = nc.variables[options.var]

for i in range(I):
    for j in range(J):
        try:
            pbar.update((i+1)*(j+1))
        except:
            pass
        #tmp = window_1Dmean(data=ma.array(data[:, i, j]), 
        #                t = nc.variables[options.scalevar][:],
        #                l=options.lowpasswindowlength, axis=axis, 
        #                parallel=False)

        #out[:, i, j] = tmp - window_1Dmean(data = tmp, 
        #                t = nc.variables[options.scalevar][:],
        #                l = options.highpasswindowlength, 
        #                axis = axis, 
        #                parallel=False)

        out[:, i, j] = window_1Dbandpass(data[:, i, j], 
                lshorterpass = options.highpasswindowlength, 
                llongerpass = options.lowpasswindowlength, 
                t = nc.variables[options.scalevar][:], 
                method = options.windowmethod, 
                axis=0, 
                parallel = False)

nc.close()
