#!/usr/bin/env python
"""
A command line tool to calculate time and space aggregation statistics
"""
import numpy as np
import glob
import argparse

from ddice.dataset.netcdf4 import netCDF4Dataset
from ddice.field import CFField
from ddice.field import grouping
from ddice.field import functions

parser = argparse.ArgumentParser(description='Time and space aggregation statistics')

parser.add_argument('--target', '-t', required=True, help='target data uri')
parser.add_argument('--groupby', '-g', required=True,  help='groupby coordinate, function, and arguments using syntax: coordinate:function,arg1,arg2,..,argN')
parser.add_argument('--apply', '-a', required=False,  help='function(s) to apply to each group in the format func:arg1,arg2,...,argN', default='mean')
parser.add_argument('--post', '-p', required=False, help='post processing function to apply')
parser.add_argument('source', nargs=1, help='source data uri and fields in the form uri:fieldname1,fieldname2,..,fieldnameN')

args = parser.parse_args()
print(args)

# Process source argument of the form uri:
source_parts = args.source[0].split(':')

if len(source_parts) > 1:
	fieldfilter = set(source_parts[1].split(','))

else:
	fieldfilter = ''

# Get filenames from wildcard
filenames = glob.glob(source_parts[0])
filenames = sorted(filenames)

# Open the source
ds = netCDF4Dataset(filenames)


fieldselect = set(ds.variables.keys())

if fieldfilter:
	fieldselect = list(fieldselect & fieldfilter)
else:
	fieldselect = list(fieldselect)

print('selected fields: {}'.format(fieldselect))

# Construct the actual field we are going to process
field = CFField(ds.variables[fieldselect[0]])

# Try and get the apply function
apply_parts = args.apply.split(':')

try:
	apply_func = eval('functions.{}'.format(apply_parts[0]))
except:
	raise Exception('Cannot find function {} in functions modules'.format(apply_parts[0]))

apply_kwargs = {}

if len(apply_parts) > 1:
	apply_args = apply_parts[1:].split(',')

	for arg in apply_args:
		try:
			name, value = arg.split('=')
		except:
			pass

		apply_kwargs[name] = value

apply_kwargs_string = ''.join(['{}={}'.format(name, value) for name, value in apply_kwargs.items()])

print("applying function {}({})".format(apply_parts[0], apply_kwargs_string))

# Try and get grouping function
groupby_parts = args.groupby.split(':')
groupby_coordinate = groupby_parts[0]

try:
	function_parts = groupby_parts[1].split(',')
except:
	raise Exception('grouping argument must specify both coordinate and grouping function name. eg: time:yearmonth')

try:
	groupby_function = eval('grouping.{}'.format(function_parts[0]))
	groupby_args = function_parts[1:]
except:
	raise Exception('failed to identify grouping function: '.format(function_parts[0]))

print('grouping function: {}({}) on coordinate {}'.format(groupby_function, groupby_args, groupby_coordinate))

# Index the groups
groupby = field.groupby(groupby_coordinate, groupby_function)
print('{} groups ({}) indexed'.format(len(groupby.groups), function_parts[0]))


# Apply the function
ds, field = field.apply(groupby, apply_parts[0], **apply_kwargs)


# Construct the target filename and variables
target_parts = args.target.split(':')

if len(target_parts) > 1:
	target_varname_parts = target_parts[1].split(',')

	ds.variables[target_varname_parts[0]] = field.variable

	for expr in target_varname_parts[1:]:
		key, value = expr.split('=')
		ds.variables[target_varname_parts[0]].attributes[key] = value

	del(ds.variables[field.variable.name])


# Create the target dataset
target = netCDF4Dataset(uri=target_parts[0], dataset=ds)








