#!/usr/bin/env python3
try:
    from mpi4py import MPI
    rank = MPI.COMM_WORLD.Get_rank()
    n_ranks = MPI.COMM_WORLD.Get_size()
except ImportError:
    rank = 0
    n_ranks = 1
from teca import *
import argparse
import numpy as np
import sys
import os

teca_profiler.initialize()
teca_profiler.start_event('temporal_reduction')

# parse the command line
parser = argparse.ArgumentParser(
    description='Reduce the time axis of a NetcCDF CF2 dataset '
                'using a predfined interval and reduction operator',
    formatter_class=lambda prog: argparse.ArgumentDefaultsHelpFormatter(
                        prog, max_help_position=4, width=100))

parser.add_argument('--input_file', type=str, required=False,
                    help='a teca_multi_cf_reader configuration file identifying'
                         ' the set of NetCDF CF2 files to process. When present'
                         ' data is read using the teca_multi_cf_reader. Use one'
                         ' of either --input_file or --input_regex.')

parser.add_argument('--input_regex', type=str, required=False,
                    help='a teca_cf_reader regex identifying the set of NetCDF'
                         ' CF2 files to process. When present data is read'
                         ' using the teca_cf_reader. Use one of either'
                         ' --input_file or --input_regex.')

parser.add_argument('--interval', type=str, default='monthly',
                    help='interval to reduce the time axis to. One of '
                         'daily, monthly, or seasonal')

parser.add_argument('--operator', type=str, default='average',
                    help='reduction operator to use. One of minimum, '
                         'maximum, or average')

parser.add_argument('--point_arrays', nargs='+', required=True,
                    help='list of point centered arrays to process.')

parser.add_argument('--fill_value', type=float, default=None,
                    help='A value that identifies missing or invalid data. '
                         'Specifying the fill value on the command line '
                         'overrides array specific fill values stored in '
                         'the file.')

parser.add_argument('--ignore_fill_value',
                    default=False, action='store_true',
                    help='Boolean flag that enables missing or invalid '
                         'value handling. When enabled NetCDF CF conventions '
                         'are used to determine fill value. Alternativley one '
                         'can explicitly provide a fill value on the command '
                         'line via the --fill_value argument.')

parser.add_argument('--output_file', type=str, required=True,
                    help='A path and file name pattern for the output NetCDF'
                         ' files. %%t%% is replaced with a human readable date'
                         ' and time corresponding to the time of the first time'
                         ' step in the file. Use --date_format to change'
                         ' the formatting')

parser.add_argument('--file_layout', type=str, default='yearly',
                    help='Selects the size and layout of the set of output'
                         ' files. May be one of number_of_steps, daily,'
                         ' monthly, seasonal, or yearly. Files are structured'
                         ' such that each file contains one of the selected'
                         ' interval. For the number_of_steps option use'
                         ' --steps_per_file.')

parser.add_argument('--steps_per_file', type=int, default=128,
                    help='number of time steps to write to each output '
                         'file')

parser.add_argument('--x_axis_variable', type=str, default='lon',
                    help='name of the variable to use for x-coordinates')

parser.add_argument('--y_axis_variable', type=str, default='lat',
                    help='name of the variable to use for y-coordinates')

parser.add_argument('--z_axis_variable', type=str, default='',
                    help='name of z coordinate variable. When processing 3D set this to'
                         ' the variable containing vertical coordinates. When empty the'
                         ' data will be treated as 2D.')

parser.add_argument('--t_axis_variable', type=str, default='time',
                    help='name of the variable to use for t-coordinates')

parser.add_argument('--n_threads', type=int, default=2,
                    help='Number of threads to use when streaming the '
                         'reduction')

parser.add_argument('--verbose', type=int, default=0,
                    help='enable verbose mode.')

# prevent spew when running under mpi
try:
    args = parser.parse_args()
except Exception:
    if rank == 0: raise

in_files = args.input_regex
out_files = args.output_file
layout = args.file_layout
steps_per_file = args.steps_per_file
n_threads = args.n_threads
interval = args.interval
operator = args.operator
point_arrays = args.point_arrays
fill_value = args.fill_value
ignore_fill_value = args.ignore_fill_value
x_axis_var = args.x_axis_variable
y_axis_var = args.y_axis_variable
z_axis_var = args.z_axis_variable
t_axis_var = args.t_axis_variable
verbose = args.verbose

if verbose and rank == 0:
    sys.stderr.write('running on %d ranks' % (n_ranks))
    sys.stderr.write('n_threads=%d\n' % (n_threads))
    sys.stderr.write('file_layout=%s\n'%(layout))
    sys.stderr.write('steps_per_file=%d\n' % (steps_per_file))
    sys.stderr.write('interval=%s\n' % (interval))
    sys.stderr.write('operator=%s\n' % (operator))
    sys.stderr.write('point_arrays=%s\n' % (str(point_arrays)))
    sys.stderr.write('ignore_fill_value=%d\n' % (ignore_fill_value))
    sys.stderr.write('fill_value=%s\n' % (str(fill_value)))

if args.input_regex:
    cfr = teca_cf_reader.New()
    cfr.set_files_regex(args.input_regex)
elif args.input_file:
    cfr = teca_multi_cf_reader.New()
    cfr.set_input_file(args.input_file)
else:
    if rank == 0:
        raise RuntimeError('Missing one of --input_file or --input_regex')

cfr.set_x_axis_variable(x_axis_var)
cfr.set_y_axis_variable(y_axis_var)
cfr.set_z_axis_variable(z_axis_var)
cfr.set_t_axis_variable(t_axis_var)

mav = teca_temporal_reduction.New()
mav.set_input_connection(cfr.get_output_port())
mav.set_interval(interval)
mav.set_operator(operator)
mav.set_fill_value(fill_value)
mav.set_use_fill_value( 0 if ignore_fill_value else 1 )
mav.set_point_arrays(point_arrays)
mav.set_verbose(verbose)
mav.set_thread_pool_size(n_threads)
mav.set_stream_size(2)

cfw = teca_cf_writer.New()
cfw.set_input_connection(mav.get_output_port())
cfw.set_verbose(verbose)
cfw.set_thread_pool_size(1)
cfw.set_layout(layout)
cfw.set_steps_per_file(steps_per_file)
cfw.set_file_name(out_files)
cfw.set_point_arrays(point_arrays)
cfw.update()

teca_profiler.end_event('temporal_reduction')
teca_profiler.finalize()
