#!/usr/bin/env python

import sys
import os
import argparse
import netCDF4 as nc
import numpy as np
import pandas as pd
import re
import datetime
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from lib_util import time_dim_to_pandas_periods


"""
What this script does:

Plot a full timeseries fields from mom ocean_scalar.nc

Example:

"""

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('input_files', nargs='+',
                        help='The MOM ocean_scalar.nc input data files.')
    parser.add_argument('--field', default='temp_global_ave',
                        help='The fields included in this plot.')
    parser.add_argument('--output_dir', default='./',
                        help='Directory where plots will be written.')

    args = parser.parse_args()

    title = None
    ylabel = None
    ts = pd.Series()

    # Go through input files one at a time building the timeseries as we go.
    for file in args.input_files:
        with nc.Dataset(file) as f:
            time_var = f.variables['time']
            data_var = f.variables[args.field]
            title = data_var.long_name
            ylabel = data_var.units

            # Calculate the times/dates, these will be our indices.
            periods = time_dim_to_pandas_periods(f.variables['time'])
            data = f.variables[args.field][:]
            assert(data.shape[1] == 1)
            data = data.flatten()

            new_ts = pd.Series(data, periods)
            ts = ts.append(new_ts)

    ts = ts.sort_index()

    plot = ts.plot()
    plt.xlabel('Time (years)')
    plt.ylabel(ylabel)
    plt.title(title)
    fig = matplotlib.pyplot.gcf()
    fig.set_size_inches(9,4.5)
    plt.savefig(os.path.join(args.output_dir, '{}.png'.format(args.field)))


if __name__ == '__main__':
    sys.exit(main())


