#!/usr/bin/env python
################################################################################
# A test program to write a large NetCDF file and see if we can read it.
################################################################################

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*8	# 8 bytes / float
MBYTE = 1024*1024

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

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

TimeFirstWrite = False
TimeVarWrite = False
TimeFileWrite = False

TimeFirstRead = False
TimeVarRead = False
TimeFileRead = False

CloseAfterSlice = False
CloseAfterVar = False

Verbose = False

# mask for variable names
VarnameMask = 'var_%04d'


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

##
# @brief Read a NetCDF file - see if it as we wrote it.
# @param filename The name of the file to read.
# @param variable_size Size of variable, in bytes.
# @param num_variables Number of variables.
def read_file(filename, variable_size, num_variables):
    # create a data array slice
    slice_array_mask = num.ones((XDIM,), 'd')

    # if timing file read, get start time
    if TimeFileRead:
        start_file_read = time.time()

    fid = nc.NetCDFFile(filename, 'r')

    # num file bytes read counter
    file_bytes_read = 0

    result = True

    for i in xrange(num_variables):
        varname = VarnameMask % i
        var = fid.variables[varname]

        # num variable bytes written counter
        var_bytes_read = 0

        # create expected slice array
        slice_array = slice_array_mask * (i + 1)
        for x in xrange(XDIM):
            slice_array[x] += x*1.0e-7

        if TimeVarRead:
            start_var_read = time.time()

        index = 0
        while var_bytes_read < variable_size:
            var_bytes_read += SLICE_SIZE
            file_bytes_read += SLICE_SIZE

            if TimeFirstRead and index == 0:
                start_first_var_read = time.time()

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

            var_array = fid.variables[varname][index,:]
            if num.any(var_array != slice_array):
                result = False
                if Verbose:
                    print 'Read variable %s, slice %d: got unexpected value' % \
                          (varname, index)

            if TimeFirstRead and index == 0:
                stop_first_var_read = time.time()
                print ("Time to do first read for variable %s: %.2f sec" %
                       (varname, (stop_first_var_read - start_first_var_read)))

            index += 1

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

        if TimeVarRead:
            stop_var_read = time.time()
            print ('Time to read variable %s: %.2f sec' %
                   (varname, (stop_var_read - start_var_read)))

    fid.close()

    if TimeFileRead:
        stop_file_read = time.time()
        print ('Time to read file: %.2f sec' %
               (stop_file_read - start_file_read))

    return result


##
# @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.
# @return The filename of the created file.
def write_file(variable_size, num_variables=1):
    # set file and variable name masks
    filename = 'test.nc'
    VarnameMask = 'var_%04d'

    # create a data array slice
    slice_array_mask = num.ones((XDIM,), 'd')

    # 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 = VarnameMask % i
        fid.createVariable(varname, 'd', ('y', 'x'))

        # num variable bytes written counter
        var_bytes_written = 0

        # create unique slice array
        slice_array = slice_array_mask * (i + 1)
        for x in xrange(XDIM):
            slice_array[x] = slice_array[x] + x*1.0e-7

        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_var_write - start_var_write)))

    fid.close()

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

    return filename


##
# @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 if none is given is '%s'." \
          % 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 read/written,"
    print "                    -c v[ar]    close & open the output file after"
    print "                                each variable is read/written,"
    print "                    -t rf[irst] time the first var slice read,"
    print "                    -t rv[ar]   time the complete var read,"
    print "                    -t rw[hole] time the complete file read,"
    print "                    -t wf[irst] time the first var slice write,"
    print "                    -t wv[ar]   time the complete var write,"
    print "                    -t ww[hole] time the complete file write,"

    if msg:
        derr, "\n%s" % msg

    sys.exit(10)

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

    if argv is None:
        argv = sys.argv

    # parse command line args
    try:
        opts, args = getopt.getopt(argv[1:], "c:t: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 == 'r':
                optargchar = optarg[1].lower()
                if optargchar == 'f':
                    TimeFirstRead = True
                elif optargchar == 'v':
                    TimeVarRead = True
                elif optargchar == 'w':
                    TimeFileRead = True
                else:
                    usage("Unrecognized -t suboption: %s" % optarg)
            elif optargchar == 'w':
                optargchar = optarg[1].lower()
                if optargchar == 'f':
                    TimeFirstWrite = True
                elif optargchar == 'v':
                    TimeVarWrite = True
                elif optargchar == 'w':
                    TimeFileWrite = True
                else:
                    usage("Unrecognized -t suboption: %s" % optarg)
            else:
                usage("Unrecognized -t 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()

    # write the required file
    filename = write_file(var_size, num_vars)
        
    # read the file to see if is as expected
    if not read_file(filename, var_size, num_vars):
        print "Didn't read data that we read!?"
        return 10

    print 'Read/write of NectCDF file was correct'
    return 0


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