#!/usr/bin/env python3

import time
import os
import tempfile
import shutil
import logging
from enum import Enum
from argparse import ArgumentParser, Namespace, FileType
from netCDF4 import Dataset, MFDataset
import geopandas as gpd
import numpy as np

domain_nodes_shp = "gis/ssm domain nodes.shp"
masked_nodes_txt = "gis/masked nodes.txt"

logger = logging.getLogger(__name__)

def get_node_ids(shps, masked):
    merged = None
    for i,shp in enumerate(shps):
        df = gpd.read_file(shp)
        df.set_index('node_id', inplace=True)
        logger.debug("Shapefile {0} has {1} nodes".format(shp, len(df)))
        if merged is None:
            merged = df.index
        else:
            merged = merged.union(df.index)
    logger.debug("get_node_ids found {0} nodes in {1} shapefiles".format(
        len(merged), len(shps)))

    masked_nodes = np.loadtxt(masked)
    merged = merged.difference(masked_nodes)
    logger.debug("{0} nodes left after masking".format(len(merged)))

    return merged.to_numpy()

DEFAULT_SIGLAYERS = [-0.01581139, -0.06053274, -0.12687974, -0.20864949,
                  -0.30326778, -0.40915567, -0.52520996, -0.65060186,
                  -0.78467834, -0.9269075 ]

def init_output(output_cdf, indata, nodes, **kwargs):
    args = Namespace(**kwargs)
    output = Dataset(output_cdf, "w")
    timeDim = output.createDimension('time', len(indata.dimensions['time']))
    nodeDim = output.createDimension('node', len(nodes))
    nodeVar = output.createVariable('node', "i4", ('node',))
    output['node'][:] = nodes
    timeVar = output.createVariable('time', "f4", ('time',))
    # Iterate over all output variables
    # If an extraction attribute is "all":
    # - add the 'siglay' dimension to the output if it's not already present
    # - include the 'siglay' dimension on the output variable
    # - add a 'zeta' output variable
    for var, attr in args.input_vars:
        if attr == InputAttr.ALL:
            siglayers = indata['siglay'][:] if 'siglay' in indata.variables else DEFAULT_SIGLAYERS
            output.createDimension('siglay', len(siglayers))
            output.createVariable('siglay', 'f4', ('siglay',))
            output['siglay'][:] = siglayers
            if 'zeta' in indata.variables:
                output.createVariable('zeta', 'f4', ('time','node'))
            break

    return output

def append_output(output_cdf):
    return Dataset(output_cdf, 'a')

def init_output_vars(output, **kwargs):
    args = Namespace(**kwargs)
    for var, attr in args.input_vars:
        out_name = args.outprefix + var
        if attr == InputAttr.BOTTOM:
            out_name += "_bottom"
        # TODO handle photic case
        dims = ('time','siglay','node') if attr == InputAttr.ALL else ('time','node')
        output.createVariable(out_name, 'f4', dims)

# Gotten from https://stackoverflow.com/questions/312443/how-do-you-split-a-list-or-iterable-into-evenly-sized-chunks
def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):
        yield lst[i:i+n]

class InputAttr(Enum):
    ALL = 0
    BOTTOM = 1
    # TODO add "photic" for the photic zone

attr_strings = {
    "all": InputAttr.ALL,
    "bottom": InputAttr.BOTTOM
}

# Expands an input variable argument into a variable name and an attribute
# describing the vertical extraction method.
def colon_meta(string):
    var, attr = string.split(':', 2)
    return (var, attr_strings[attr])

def main():
    script_home = os.path.dirname(os.path.realpath(__file__))
    parser = ArgumentParser(description="Extract data from SSM netcdf output files")
    parser.add_argument("incdf", nargs="+", help="each input CDF file")
    parser.add_argument("outcdf",
            help="the output CDF file (created if it doesn't exist)")
    parser.add_argument("outprefix",
            help="a prefix for the extracted variables in the output CDF")
    parser.add_argument("-d", dest="domain_node_shapefiles", action="append",
            help="Specify a domain node shapefile")
    parser.add_argument("-m", dest="masked_nodes_file", type=FileType('r'),
            help="Specify a different masked nodes text file")
    parser.add_argument("--invar", dest="input_vars", type=colon_meta,
            action="append",
            help="Extract the values of a different output variable")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
            help="Print progress messages during the extraction")
    parser.add_argument("-c", "--chunk-size", type=int, dest="chunk_size",
            help="Process this many CDF files at once")
    parser.add_argument("--cache", dest="cache", action="store_true",
            help="Use a read/write cache in a temporary directory")
    # Cannot include default values of lists here, see
    # https://bugs.python.org/issue16399
    parser.set_defaults(chunk_size=4, verbose=False,
            masked_nodes_file=os.path.join(script_home, masked_nodes_txt))
    args = parser.parse_args()
    # This is the workaround
    if not args.input_vars:
        args.input_vars = [("DOXG",InputAttr.BOTTOM)]
    if not args.domain_node_shapefiles:
        args.domain_node_shapefiles = [os.path.join(script_home, domain_nodes_shp)]

    logging.basicConfig(level=logging.INFO if args.verbose else logging.WARNING)
    #logger.setLevel(logging.DEBUG)

    if args.cache:
        with tempfile.TemporaryDirectory() as tmpdir:
            exist_cdfs = []
            logger.info("Caching input files...")
            for infile in args.incdf:
                newpath = os.path.join(tmpdir, os.path.basename(infile))
                shutil.copy(infile, newpath)
                exist_cdfs.append(newpath)
            output_cdf = os.path.join(tmpdir, os.path.basename(args.outcdf))
            if os.path.exists(args.outcdf):
                logger.info("Caching output file...")
                shutil.copy(args.outcdf, output_cdf)
            do_extract(exist_cdfs, output_cdf, **vars(args))
            # Copy the resulting output CDF back
            logger.info("Saving output file...")
            shutil.copy(output_cdf, args.outcdf)
            logger.info("Finished.")
    else:
        do_extract(args.incdf, args.outcdf, **vars(args))

def do_extract(exist_cdfs, output_cdf, **kwargs):
    args = Namespace(**kwargs)
    logger.info("Determining scope of work...")
    indata = MFDataset(exist_cdfs) if len(exist_cdfs) > 1 else Dataset(exist_cdfs[0])
    node_ids = get_node_ids(args.domain_node_shapefiles, args.masked_nodes_file)
    logger.info("Initializing output file...")
    if not os.path.exists(output_cdf):
        outdata = init_output(output_cdf, indata, node_ids, **vars(args))
        outdata['time'][:] = indata['time'][:] / 3600 / 24
    else:
        outdata = append_output(output_cdf)
    init_output_vars(outdata, **vars(args))

    # Attempts to use the entire MFDataset don't seem to scale well.
    # Instead, I'm resorting to a blocking approach where MFDatasets are
    # created for only a few netCDF files at a time
    indata.close()
    i = 0
    total = 0
    logger.info("Beginning extraction...")
    start_time = time.perf_counter()
    times_ct = outdata.dimensions['time'].size
    for cdfchunk in chunks(exist_cdfs, args.chunk_size):
        c = MFDataset(cdfchunk) if len(cdfchunk) > 1 else Dataset(cdfchunk[0])
        chunk_times = len(c.dimensions['time'])

        data = copy_data(c, outdata, i, node_ids, **vars(args))
        i += chunk_times
        c.close()

        elapsed = (time.perf_counter() - start_time)
        to_go = elapsed * (times_ct / i - 1)
        total += np.sum([d.size * d.itemsize for k,d in data.items()])
        logger.info("{0}/{1} ({2}s elapsed, {3}s to go, {4}KBps)".format(i,
            times_ct, int(elapsed), int(to_go), int(total/elapsed/1000)))
    logger.info("Extraction finished.")
    outdata.close()

def copy_data(cdfin, cdfout, timeidx, node_ids, **kwargs):
    args = Namespace(**kwargs)
    times_ct = len(cdfin.dimensions['time'])
    alldata = {}
    # Copy zeta if it's needed
    if 'zeta' in cdfout.variables:
        alldata['zeta'] = cdfin['zeta'][:, node_ids - 1]
        cdfout['zeta'][timeidx:timeidx + times_ct, :] = alldata['zeta']
    for var, attr in args.input_vars:
        out_name = args.outprefix + var
        if attr == InputAttr.ALL:
            slc = slice(None)
        elif attr == InputAttr.BOTTOM:
            slc = -1
            out_name += "_bottom"
        # TODO add "photic" case which will look rather different
        data = cdfin[var][:, slc, node_ids - 1]
        logger.debug("data is shape " + str(data.shape))
        if attr == InputAttr.ALL:
            cdfout[out_name][timeidx:timeidx+times_ct,:,:] = data
        else:
            cdfout[out_name][timeidx:timeidx+times_ct,:] = data
        alldata[out_name] = data
    return alldata

if __name__ == "__main__": main()

[((17, 9, 17, 36), 'logging.getLogger', 'logging.getLogger', ({(17, 27, 17, 35): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((32, 19, 32, 37), 'numpy.loadtxt', 'np.loadtxt', ({(32, 30, 32, 36): 'masked'}, {}), '(masked)', True, 'import numpy as np\n'), ((43, 11, 43, 30), 'argparse.Namespace', 'Namespace', ({}, {}), '(**kwargs)', False, 'from argparse import ArgumentParser, Namespace, FileType\n'), ((44, 13, 44, 37), 'netCDF4.Dataset', 'Dataset', ({(44, 21, 44, 31): 'output_cdf', (44, 33, 44, 36): '"""w"""'}, {}), "(output_cdf, 'w')", False, 'from netCDF4 import Dataset, MFDataset\n'), ((68, 11, 68, 35), 'netCDF4.Dataset', 'Dataset', ({(68, 19, 68, 29): 'output_cdf', (68, 31, 68, 34): '"""a"""'}, {}), "(output_cdf, 'a')", False, 'from netCDF4 import Dataset, MFDataset\n'), ((71, 11, 71, 30), 'argparse.Namespace', 'Namespace', ({}, {}), '(**kwargs)', False, 'from argparse import ArgumentParser, Namespace, FileType\n'), ((104, 13, 104, 84), 'argparse.ArgumentParser', 'ArgumentParser', (), '', False, 'from argparse import ArgumentParser, Namespace, FileType\n'), ((134, 4, 134, 80), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((158, 11, 158, 30), 'argparse.Namespace', 'Namespace', ({}, {}), '(**kwargs)', False, 'from argparse import ArgumentParser, Namespace, FileType\n'), ((177, 17, 177, 36), 'time.perf_counter', 'time.perf_counter', ({}, {}), '()', False, 'import time\n'), ((196, 11, 196, 30), 'argparse.Namespace', 'Namespace', ({}, {}), '(**kwargs)', False, 'from argparse import ArgumentParser, Namespace, FileType\n'), ((22, 13, 22, 31), 'geopandas.read_file', 'gpd.read_file', ({(22, 27, 22, 30): 'shp'}, {}), '(shp)', True, 'import geopandas as gpd\n'), ((103, 34, 103, 60), 'os.path.realpath', 'os.path.realpath', ({(103, 51, 103, 59): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((160, 13, 160, 34), 'netCDF4.MFDataset', 'MFDataset', ({(160, 23, 160, 33): 'exist_cdfs'}, {}), '(exist_cdfs)', False, 'from netCDF4 import Dataset, MFDataset\n'), ((160, 63, 160, 85), 'netCDF4.Dataset', 'Dataset', ({(160, 71, 160, 84): 'exist_cdfs[0]'}, {}), '(exist_cdfs[0])', False, 'from netCDF4 import Dataset, MFDataset\n'), ((163, 11, 163, 37), 'os.path.exists', 'os.path.exists', ({(163, 26, 163, 36): 'output_cdf'}, {}), '(output_cdf)', False, 'import os\n'), ((112, 61, 112, 74), 'argparse.FileType', 'FileType', ({(112, 70, 112, 73): '"""r"""'}, {}), "('r')", False, 'from argparse import ArgumentParser, Namespace, FileType\n'), ((126, 30, 126, 73), 'os.path.join', 'os.path.join', ({(126, 43, 126, 54): 'script_home', (126, 56, 126, 72): 'masked_nodes_txt'}, {}), '(script_home, masked_nodes_txt)', False, 'import os\n'), ((132, 39, 132, 82), 'os.path.join', 'os.path.join', ({(132, 52, 132, 63): 'script_home', (132, 65, 132, 81): 'domain_nodes_shp'}, {}), '(script_home, domain_nodes_shp)', False, 'import os\n'), ((138, 13, 138, 42), 'tempfile.TemporaryDirectory', 'tempfile.TemporaryDirectory', ({}, {}), '()', False, 'import tempfile\n'), ((146, 15, 146, 42), 'os.path.exists', 'os.path.exists', ({(146, 30, 146, 41): 'args.outcdf'}, {}), '(args.outcdf)', False, 'import os\n'), ((152, 12, 152, 48), 'shutil.copy', 'shutil.copy', ({(152, 24, 152, 34): 'output_cdf', (152, 36, 152, 47): 'args.outcdf'}, {}), '(output_cdf, args.outcdf)', False, 'import shutil\n'), ((180, 12, 180, 31), 'netCDF4.MFDataset', 'MFDataset', ({(180, 22, 180, 30): 'cdfchunk'}, {}), '(cdfchunk)', False, 'from netCDF4 import Dataset, MFDataset\n'), ((180, 58, 180, 78), 'netCDF4.Dataset', 'Dataset', ({(180, 66, 180, 77): 'cdfchunk[0]'}, {}), '(cdfchunk[0])', False, 'from netCDF4 import Dataset, MFDataset\n'), ((187, 19, 187, 38), 'time.perf_counter', 'time.perf_counter', ({}, {}), '()', False, 'import time\n'), ((143, 16, 143, 44), 'shutil.copy', 'shutil.copy', ({(143, 28, 143, 34): 'infile', (143, 36, 143, 43): 'newpath'}, {}), '(infile, newpath)', False, 'import shutil\n'), ((145, 46, 145, 75), 'os.path.basename', 'os.path.basename', ({(145, 63, 145, 74): 'args.outcdf'}, {}), '(args.outcdf)', False, 'import os\n'), ((148, 16, 148, 52), 'shutil.copy', 'shutil.copy', ({(148, 28, 148, 39): 'args.outcdf', (148, 41, 148, 51): 'output_cdf'}, {}), '(args.outcdf, output_cdf)', False, 'import shutil\n'), ((142, 47, 142, 71), 'os.path.basename', 'os.path.basename', ({(142, 64, 142, 70): 'infile'}, {}), '(infile)', False, 'import os\n')]