#!/work/ROMO/anaconda_envs/basic38/bin/python
import PseudoNetCDF as pnc
import numpy as np
import argparse
import os


def dailymass(outv, average, factor):
    """
    Arguments
    ---------
    outv : array-like
        Variable to sum. Assumes 24 steps is averageinstantaneous  and anything elhours of s**-1.

    Returns
    -------
    total : float
        Sum of mass over the day
    """
    axes = tuple(list(range(len(outv.shape)))[1:])
    diurnal = outv[:].sum(axes)
    if average or diurnal.shape[0] == 0:
        out = diurnal.sum()
    else:
        out = (diurnal[1:] + diurnal[:-1]).sum()
    
    return float(out)


def merge(outf, path, args):
    """
    Arguments
    ---------
    outf : PseudoNetCDF file
        Add data from path to this file
    path : str
        Path to file to add

    Returns
    -------
    None

    """
    if args.verbose > 0:
        print(path)
    othf = pnc.pncopen(path, format=args.format)
    hist = getattr(outf, 'HISTORY', '')
    hist += '\n' + path
    for k, othv in othf.variables.items():
        if k in args.coords:
            continue
        if k in outf.variables:
            outv = outf.variables[k]
        else:
            outv = outf.copyVariable(othv, key=k, withdata=False)

        outunit = outv.units.strip()
        if outunit.endswith('/s') and args.factor == 3600:
            fileunit = outunit[:-2] + '/file'
        elif outunit.endswith('/h') and args.factor == 1:
            fileunit = outunit[:-2] + '/file'
        else:
            fileunit = outunit + '*{}/file'.format(args.factor)

        start = dailymass(outv, args.average, args.factor)
        toadd = dailymass(othv, args.average, args.factor)
        outv[:, :othv.shape[1], :, :] += othv[:]
        new = dailymass(outv, args.average, args.factor)
        added = new - start
        if args.verbose > 1:
            print(outv.units.strip(), othv.units.strip())
            print(
                (
                    '{} Start: {:.2g} {}; ToAdd: {:.2g}; ' +
                    'Added: {:.2g}; New: {:.2g}'
                ).format(k, start, fileunit, toadd, added, new)
            )
        outv.daily_total_grams = new
        hist += '\n + {}: {:4g} {}'.format(k, added, fileunit)
        setattr(outf, 'HISTORY', hist)


def process(args):
    if not args.overwrite and os.path.exists(args.outpath):
        raise IOError(
            args.outpath + 
            ' exists; remove it or use the overwrite option'
        )

    # Requires starting with the most layers.
    idx = np.argmax([
        len(pnc.pncopen(path, format=args.format).dimensions[args.layer])
        for path in args.inpaths
    ])

    tmplpath = args.inpaths[idx]
    outf = pnc.pncopen(tmplpath, format=args.format).copy()

    skipkeys = ['TFLAG'] + list(outf.dimensions)
    args.coords += skipkeys
    if args.format == 'ioapi':
        TSTEP = outf.TSTEP

    # Empty data
    for k, v in outf.variables.items():
        if k in args.coords:
            continue
        v[:] = 0.

    for path in args.inpaths:
        merge(outf, path, args)

    if args.format == 'ioapi':
        del outf.variables['TFLAG']
        outf.TSTEP = TSTEP
        outf.updatemeta()

    outf.save(
        args.outpath, format='NETCDF4_CLASSIC',
        complevel=1, verbose=args.verbose
    )


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-c', '--coords', type=lambda x: x.split(','), default=[],
        help=(
            'Comma delimited list of variables to be treated as ' + 
            'coordinates in addition to dimensions.'
        )
    )
    parser.add_argument(
        '--factor', default=3600,
        help=(
            'Used to convert rates to time step integrals ' +
            '(default: 3600; e.g., hourly moles/s * 3600s/h)'
        )
    )
    parser.add_argument(
        '-a', '--average', default=False, action='store_true',
        help=(
            'For use in calculating rates. Default (average=False) assumes' +
            'instantaneous, but -a toggles to assume averaged'
        )
    )
    parser.add_argument(
        '-f', '--format', default='ioapi', help='Format ioapi (default), netcdf, etc'
    )
    parser.add_argument(
        '-z', '--layer', default='LAY', help='Name of the Z dimension (default LAY)'
    )
    parser.add_argument(
        '-O', '--overwrite', default=False, action='store_true',
        help=(
            'By default, exiting files will not be overwritten. Use this ' +
            'option to overwrite existing files'
        )
    )
    parser.add_argument(
        '-v', '--verbose', action='count', default=0,
        help=(
            'Increasing verbosity prints file-level progress at level 1, ' +
            'variable level progress at level 2'
        )
    )
    parser.add_argument(
        'inpaths', nargs='+', metavar='inpath',
        help='Takes one or many input files'
    )
    parser.add_argument('outpath')
    parser.description = """
Will add variables from files together and document the masses from each
in the HISTORY attribute. Files must either all have the same number of times
or a single time, which will be treated as constant. Files can have different
layers, and different variables. All other dimensions must be shared except
for time. The number of times must either be the same or length 1. If length
1, the files rate will be assumed constant and applied to all times.

Summary process:
1. Finds the input file with the most layers.
2. Creates a template output file.
3. For each input file, it add each varible to the template output file
   in the appropriate layers. If a new variable is encountered, it is created
   in the output file.
4. The amount of mass added is documented in the HISTORY attribute.
"""
    args = parser.parse_args()
    process(args)