#!/usr/bin/env python
import os
import argparse
import numpy as np
from traceback import format_exc
from sys import float_info
from glob import glob
from GetIds import *
try:
    from parallelrange import parallelrange
except ImportError:  #we don't have issm code in path, just get it
    import devpath
    from parallelrange import parallelrange
from IdToName import IdToName
from arch import archread
from arch import archwrite


def runme(id=None, exclude=None, benchmark='nightly', procedure='check', output='none', rank=1, numprocs=1):
    """
    RUNME - test deck for ISSM nightly runs

        In a test deck directory (tests / Vertification / NightlyRun for example)
        The following command will launch all the existing tests:
    >>     runme()
        To run the tests 101 and 102:
    >>     runme(id = [101, 102])
        etc...

        Available options:
            'id'            followed by the list of ids or (parts of) test names requested
                            Note: runs all tests by default

            'exclude'       ids or (parts of) test names to be excluded (same format as id)
                            Note: exclude does nothing if 'id' is specified with different values

            'benchmark'     'all' (all of the tests)
                            'nightly' (nightly run / daily run)
                            'ismip'  : validation of ismip - hom tests
                            'eismint': validation of eismint tests
                            'thermal': validation of thermal tests
                            'mesh'   : validation of mesh tests
                            'adolc'  : validation of adolc tests
                            'slr'    : validation of slr tests

            'procedure'     'check' : run the test (default)
                            'update': update the archive

        Usage:
            runme(varargin)

        Examples:
            runme()
            runme(101)
            runme('SquareShelf')
            runme(exclude = 2001)
            runme(exclude = 'Dakota', benchmark = 'all')
            runme(id = [[101, 102], ['Dakota', 'Slr']])
        """
    #Get ISSM_DIR variable
    ISSM_DIR = os.environ['ISSM_DIR']

    #Process options
    #GET benchmark {{{
    if benchmark not in ['all', 'nightly', 'ismip', 'eismint',
                         'thermal', 'mesh', 'validation', 'tranforcing',
                         'adolc', 'slr', 'referential']:
        print(("runme warning: benchmark '{}' not supported, defaulting to test 'nightly'.".format(benchmark)))
        benchmark = 'nightly'
        # }}}
        #GET procedure {{{
    if procedure not in ['check', 'update']:
        print(("runme warning: procedure '{}' not supported, defaulting to test 'check'.".format(procedure)))
        procedure = 'check'
        # }}}
        #GET output {{{
    if output not in ['nightly', 'none']:
        print(("runme warning: output '{}' not supported, defaulting to test 'none'.".format(output)))
        output = 'none'
        # }}}
        #GET RANK and NUMPROCS for multithreaded runs {{{
    if (numprocs < rank):
        numprocs = 1
        # }}}
        #GET ids  {{{
    flist = glob('test*.py')  #File name must start with 'test' and must end by '.py' and must be different than 'test.py'
    list_ids = [int(file[4:-3]) for file in flist if not file == 'test.py']  #Keep test id only (skip 'test' and '.py')

    i1, i2 = parallelrange(rank, numprocs, len(list_ids))  #Get tests for this cpu only
    list_ids = list_ids[i1:i2 + 1]
    if np.size(id) > 0 and id is not None:
        test_ids = set(GetIds(id)).intersection(set(list_ids))
        benchmark = None
    else:
        # if no tests are specifically provided, do them all
        test_ids = set(list_ids)

    # }}}
    #GET exclude {{{
    exclude_ids = GetIds(exclude)
    test_ids = test_ids.difference(exclude_ids)
    # }}}
    #Process Ids according to benchmarks {{{
    if benchmark == 'nightly':
        test_ids = test_ids.intersection(set(range(1, 1000)))
    elif benchmark == 'validation':
        test_ids = test_ids.intersection(set(range(1001, 2000)))
    elif benchmark == 'ismip':
        test_ids = test_ids.intersection(set(range(1101, 1200)))
    elif benchmark == 'eismint':
        test_ids = test_ids.intersection(set(range(1201, 1300)))
    elif benchmark == 'thermal':
        test_ids = test_ids.intersection(set(range(1301, 1400)))
    elif benchmark == 'mesh':
        test_ids = test_ids.intersection(set(range(1401, 1500)))
    elif benchmark == 'tranforcing':
        test_ids = test_ids.intersection(set(range(1501, 1503)))
    elif benchmark == 'referential':
        test_ids = test_ids.intersection(set(range(1601, 1603)))
    elif benchmark == 'slr':
        test_ids = test_ids.intersection(set(range(2001, 2500)))
    elif benchmark == 'adolc':
        test_ids = test_ids.intersection(set(range(3001, 3200)))
    test_ids = list(test_ids)
    test_ids.sort()
    # }}}

    #Loop over tests and launch sequence
    root = os.getcwd()
    errorcount = 0
    erroredtest_list = []
    for id in test_ids:
        print(("----------------starting:{}----------------------- ".format(id)))
        try:
            #Execute test
            os.chdir(root)
            id_string = IdToName(id)
            exec(compile(open('test{}.py'.format(id)).read(), 'test{}.py'.format(id), 'exec'), globals())

            #UPDATE ARCHIVE?
            archive_name = 'Archive' + str(id)
            if procedure == 'update':
                archive_file = os.path.join('..', 'Archives', archive_name + '.arch')
                if os.path.isfile(archive_file):
                    os.remove(archive_file)
                for k, fieldname in enumerate(field_names):
                    field = np.array(field_values[k], dtype=float)
                    if len(field.shape) == 1:
                        if np.size(field):
                            field = field.reshape(np.size(field), 1)
                        else:
                            field = field.reshape(0, 0)
                    elif len(field.shape) == 0:
                        field = field.reshape(1, 1)
                        # Matlab uses base 1, so use base 1 in labels
                    archwrite(archive_file, archive_name + '_field' + str(k + 1), field)
                print(("File {} saved. \n".format(os.path.join('..', 'Archives', archive_name + '.arch'))))

                    #ELSE: CHECK TEST
            else:
                #load archive
                if os.path.exists(os.path.join('..', 'Archives', archive_name + '.arch')):
                    archive_file = os.path.join('..', 'Archives', archive_name + '.arch')
                else:
                    raise IOError("Archive file '../Archives/{}.arch' does not exist.".format(archive_name))

                for k, fieldname in enumerate(field_names):
                    try:
                        #Get field and tolerance
                        field = np.array(field_values[k])
                        if len(field.shape) == 1:
                            if np.size(field):
                                field = field.reshape(np.size(field), 1)
                            else:
                                field = field.reshape(0, 0)
                        tolerance = field_tolerances[k]

                        #compare to archive
                        # Matlab uses base 1, so use base 1 in labels
                        archive = np.array(archread(archive_file, archive_name + '_field' + str(k + 1)))
                        #Because np.array is weird (str(np.array(None)) becomes 'None' but np.array(None) is never equal to None, it basically becomes a type of string in an array):
                        if str(archive) == 'None':
                            raise NameError("Field name '" + archive_name + '_field' + str(k + 1) + "' does not exist in archive file.")
                        if np.shape(field) != np.shape(archive) and not np.shape(field) in [(1, 1), (0, 0), (1, 0), (0, 1)]:
                            field = field.T
                            if np.shape(field) != np.shape(archive):
                                raise RuntimeError("Field '{}'' from test is malformed; shape is {}, should be {} or {}".format(archive_name, np.shape(field.T), np.shape(archive), np.shape(archive.T)))

                        error_diff = np.amax(np.abs(archive - field), axis=0) / (np.amax(np.abs(archive), axis=0) + float_info.epsilon)
                        if not np.isscalar(error_diff):
                            error_diff = error_diff[0]

                            #disp test result
                        if (np.any(error_diff > tolerance) or np.isnan(error_diff)):
                            print(('ERROR   difference: {} > {} test id: {} test name: {} field: {}'.format(error_diff, tolerance, id, id_string, fieldname)))
                            errorcount += 1
                            erroredtest_list.append(id)
                        else:
                            print(('SUCCESS difference: {} < {} test id: {} test name: {} field: {}'.format(error_diff, tolerance, id, id_string, fieldname)))

                    except Exception as message:
                        #something went wrong, print failure message:
                        print((format_exc()))
                        if output == 'nightly':
                            fid = open(os.path.join(ISSM_DIR, 'nightlylog', 'pythonerror.log'), 'a')
                            fid.write('%s' % message)
                            fid.write('\n------------------------------------------------------------------    \n')
                            fid.close()
                            print(('FAILURE difference: N/A test id: {} test name: {} field: {}'.format(id, id_string, fieldname)))
                        else:
                            print(('FAILURE difference: N/A test id: {} test name: {} field: {}'.format(id, id_string, fieldname)))
                            raise RuntimeError(message)

        except Exception as message:
            #something went wrong, print failure message:
            print((format_exc()))
            if output == 'nightly':
                fid = open(os.path.join(ISSM_DIR, 'nightlylog', 'pythonerror.log'), 'a')
                fid.write('%s' % message)
                fid.write('\n------------------------------------------------------------------ \n')
                fid.close()
                print(('FAILURE difference: N/A test id: {} test name: {} field: {}'.format(id, id_string, 'N/A')))
            else:
                print(('FAILURE difference: N/A test id: {} test name: {} field: {}'.format(id, id_string, 'N/A')))
                raise RuntimeError(message)

        print(("----------------    finished:{}----------------------- ".format(id)))

    if errorcount > 0:
        print("{} errors were detected in test {}".format(errorcount, np.unique(erroredtest_list)))
    return


if __name__ == '__main__':
    if 'PYTHONSTARTUP' in os.environ:
        PYTHONSTARTUP = os.environ['PYTHONSTARTUP']
#print 'PYTHONSTARTUP = ', PYTHONSTARTUP
        if os.path.exists(PYTHONSTARTUP):
            try:
                exec(compile(open(PYTHONSTARTUP).read(), PYTHONSTARTUP, 'exec'))
            except Exception as e:
                print("PYTHONSTARTUP error: ", e)
        else:
            print(("PYTHONSTARTUP file '{}' does not exist.".format(PYTHONSTARTUP)))

        parser = argparse.ArgumentParser(description='RUNME - test deck for ISSM nightly runs')
        parser.add_argument('-i', '--id', nargs='*', type=int, help='followed by the list of ids requested', default=[])
        parser.add_argument('-in', '--include_name', nargs='*', type=str, help='followed by the list of test names requested', default=[])
        parser.add_argument('-e', '--exclude', nargs='+', type=int, help='ids to be excluded from the test', default=[])
        parser.add_argument('-en', '--exclude_name', nargs='+', type=str, help='test names to be excluded from the test', default=[])
        parser.add_argument('-b', '--benchmark', help='nightly/ismip/eismint/thermal/mesh/...', default='nightly')
        parser.add_argument('-p', '--procedure', help='check/update', default='check')
        parser.add_argument('-o', '--output', help='nightly/daily/none', default='none')
        parser.add_argument('-r', '--rank', type=int, help='rank', default=1)
        parser.add_argument('-n', '--numprocs', type=int, help='numprocs', default=1)
        args = parser.parse_args()

        md = runme([args.id, args.include_name], [args.exclude, args.exclude_name], args.benchmark, args.procedure, args.output, args.rank, args.numprocs)

        exit(md)
    else:
        print("PYTHONSTARTUP not defined in environment")
        raise RuntimeError()
