#!/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

import numpy as np
from numpy import ma
from netCDF4 import Dataset

from maud import window_1Dbandpass
from cmaud import window_mean_2D_latlon


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

parser.add_option("-l", dest="l",
    help=".")

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")

parser.add_option("--interp", dest="interp",
    help="If selected fill the masked points if is there available data around.")

parser.add_option("--npes", dest="npes",
    help="Define the number of parallel processes.",
    default=None)

(options, args) = parser.parse_args()

options.l = float(options.l)
options.varout  = options.var + "_maud"
if hasattr(options, 'interp'):
    options.interp = True
else:
    options.interp = False

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

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


attributes = nc.variables[options.var].ncattrs()
#nc.variables[options.var].missing_value
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)

lat = nc.variables['latitude']
lon = nc.variables['longitude']

if (len(lat.dimensions)==1) & (len(lon.dimensions)==1):
    if (lat.dimensions[0] == nc.variables[options.var].dimensions[1]) & \
            (lon.dimensions[0] == nc.variables[options.var].dimensions[2]):
                Lon, Lat = np.meshgrid(lon[:], lat[:])
    T, I, J = nc.variables[options.var].shape
else:
    import sys; sys.exit()

data = nc.variables[options.var]

try:
    from progressbar import ProgressBar
    pbar = ProgressBar(maxval=T).start()
except:
    print "ProgressBar is not available"

try:
    import multiprocessing as mp
    if options.npes is None:
        npes = 2 * mp.cpu_count()
    else:
        #assert type(options.npes) == int
        npes = int(options.npes)
    print " Will work with %s npes" % npes
    data_smooth = ma.empty(data.shape)
    pool = mp.Pool(npes)
    results = []

    for nt in range(T):
        results.append( pool.apply_async( window_mean_2D_latlon, (Lat, Lon, data[nt], options.l, options.windowmethod, options.interp) ) )

    for nt, r in enumerate(results):
        try:
            pbar.update(nt)
        except:
            pass

        out[nt] = data[nt] - r.get()

except:
    print "Sorry, didn't work to filter in parallel"
    raise

#for nt in range(T):
#        try:
#            pbar.update(nt)
#        except:
#            pass
#
#        tmp = window_mean_2D_latlon(Lat, Lon, data[nt], options.l, method = options.windowmethod) 
#        out[nt] = tmp - data[nt]

nc.close()
