#!/usr/bin/env python
################################################################################
# A test program to write large NetCDF files.
################################################################################

import sys
import getopt
import time
import numpy as num
import Scientific.IO.NetCDF as nc

######
# Various constants
######

# default modifier and number of files
DEF_MODIFIER = 'M'

# array slice definitions
XDIM = 1024*1024
SLICE_SIZE = XDIM*4	# 4 bytes / int
MBYTE = 1024*1024

# Dictionary for suffix modifiers
Suffixes = {'M': 1024*1024,
            'G': 1024*1024*1024
           }

######
# Globals
######

TimeFirstWrite = False
TimeVarWrite = False
TimeFileWrite = False

CloseAfterSlice = False
CloseAfterVar = False

Verbose = False

######
# mainline code
######

##
# @brief Write a NetCDF file with set number of variables of a defined size.
# @param variable_size Size of variable, in bytes.
# @param num_variables Number of required variables.
def write_file(variable_size, num_variables=1):
    # set file and variable name masks
    filename = 'test.nc'
    varname_mask = 'var_%04d'

    # create a data array slice
    slice_array = num.ones((XDIM,), 'i')

    # if timing file write, remember start time
    if TimeFileWrite:
        start_file_write = time.time()

    fid = nc.NetCDFFile(filename, 'w')
    fid.createDimension('y', None)
    fid.createDimension('x', XDIM)

    # num file bytes written counter
    file_bytes_written = 0

    for i in xrange(num_variables):
        varname = varname_mask % i
        fid.createVariable(varname, 'i', ('y', 'x'))

        # num variable bytes written counter
        var_bytes_written = 0

        if TimeVarWrite:
            start_var_write = time.time()

        index = 0
        while var_bytes_written < variable_size:
            var_bytes_written += SLICE_SIZE
            file_bytes_written += SLICE_SIZE

            if TimeFirstWrite and index == 0:
                start_first_var_write = time.time()

            if Verbose:
                print ('File %s, variable %s, writing slice %d: '
                       'var=%.1fMiB, file=%.1fMiB' %
                       (filename, varname, index, float(var_bytes_written)/MBYTE,
                        float(file_bytes_written)/MBYTE))

            fid.variables[varname][index,:] = slice_array

            if TimeFirstWrite and index == 0:
                stop_first_var_write = time.time()
                print ("Time to do first write for variable %s: %.2f sec" %
                       (varname,
                        (stop_first_var_write - start_first_var_write)))

            index += 1

            if CloseAfterSlice:
                fid.close()
                fid = nc.NetCDFFile(filename, 'a')
        if CloseAfterVar:
            fid.close()
            fid = nc.NetCDFFile(filename, 'a')

        if TimeVarWrite:
            stop_var_write = time.time()
            print ('Time to write variable %s: %.2f sec' %
                   (varname, (stop_first_var_write - start_first_var_write)))

    fid.close()

    if TimeFileWrite:
        stop_file_write = time.time()
        print ('Time to write file: %.2f sec' %
               (stop_file_write - start_file_write))

##
# @brief Provide help for the befuddled user.
# @return Doesn't, calls sys.exit().
def usage(msg=None):
    print "Usage: write_large_files <opts> <varsize> [<numvars>]"
    print ""
    print "where <varsize> is a number followed by an optional modifier:"
    print "                    1024M or 4G"
    print "                the assumed modifier is '%s' if none is given." \
          % DEF_MODIFIER
    print "  and <numvars> is the number of variables of the above size"
    print "                    to write.  If not supplied, 1 is assumed."
    print "  and <opts>    is zero or more of:"
    print "                    -c s[lice]  close & open the output file after"
    print "                                each variable slice is written,"
    print "                    -c v[ar]    close & open the output file after"
    print "                                each variable is written,"
    print "                    -t f[irst]  time the first var slice write,"
    print "                    -t v[ar]    time the complete var write,"
    print "                    -t w[hole]  time the complete file write,"

    if msg:
        print "\n%s" % msg

    sys.exit(10)

def main(argv=None):
    global TimeFirstWrite, TimeVarWrite, TimeFileWrite
    global CloseAfterSlice, CloseAfterVar
    global Verbose

    if argv is None:
        argv = sys.argv

    # parse command line args
    try:
        opts, args = getopt.getopt(argv[1:], "c:ht:v", ["help"])
    except getopt.error, msg:
        usage(msg)

    for (opt, optarg) in opts:
        if opt == '-c':
            optargchar = optarg[0].lower()
            if optargchar == 's':
                CloseAfterSlice = True
            elif optargchar == 'v':
                CloseAfterVar = True
            else:
                usage("Unrecognized -c suboption: %s" % optarg)
        elif opt == '-t':
            optargchar = optarg[0].lower()
            if optargchar == 'f':
                TimeFirstWrite = True
            elif optargchar == 'v':
                TimeVarWrite = True
            elif optargchar == 'w':
                TimeFileWrite = True
            else:
                usage("Unrecognized -c suboption: %s" % optarg)
        elif opt == '-v':
            Verbose = True
        else:
            usage("Unrecognized option: %s" % opt)

    if len(args) != 1 and len(args) != 2:
        usage()

    var_size = args[0][:-1]
    modifier =args[0][-1]

    if modifier in '0123456789':
        var_size = args[0]
        modifier = DEF_MODIFIER
    modifier = Suffixes.get(modifier, None)
    if modifier is None:
        usage()
        
    try:
        var_size = int(var_size) * modifier
    except:
        usage()

    num_vars = 1
    if len(args) == 2:
        try:
            num_vars = int(args[1])
        except:
            usage()

    # call production function
    write_file(var_size, num_vars)
        
if __name__ == "__main__":
    sys.exit(main())