#!/usr/bin/env python
#
# NOTE: Switch to,
#
#   #!/usr/bin/python3
#
# when making Python 3 default
#
import argparse
from glob import glob
import os
from re import search, split
from sys import float_info
from traceback import format_exc

import numpy as np

try:
    from arch import archread
except:  # ISSM_DIR is not on path
    import devpath

from arch import archread
from arch import archwrite
from GetIds import *
from IdToName import IdToName
from parallelrange import parallelrange
from loadmodel import loadmodel
from solve import solve
from importlib import import_module


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 (for example, test/NightlyRun) the following
    command will launch all existing tests,

        ./runme.py

    To run tests 101 and 102,

        ./runme.py -i [101, 102]

    Options:
        -i/--id         Followed by the list of ids or (parts of) test names
                        requested
        -e/--exclude    Ids or (parts of) test names to be excluded (same
                        format as id). Does nothing if 'id' is specified with
                        different values.
        -b/--benchmark  'all'           : (all of the tests)
                        'nightly'       : (nightly run/daily run)
                        'validation'    : (validation)
                        'adolc'         : validation of adolc tests
                        'eismint'       : validation of eismint tests
                        'ismip'         : validation of ismip-hom tests
                        'mesh'          : validation of mesh tests
                        'qmu'           : validation of qmu tests
                        'referential'   : validation of referential tests
                        'slc'           : validation of slc tests
                        'thermal'       : validation of thermal tests
                        'tranforcing'   : validation of transient forcing tests
        -p/--procedure  'check'         : run the test (default)
                        'update'        : update the archive
                        'runFromNC'     : run from an existing nc file


    Usage:
        ./runme.py [option [args]]

    Examples:
        ./runme.py
        ./runme.py -i 101
        ./runme.py -i 'SquareShelf'
        ./runme.py -e 2001
        ./runme.py -e 'Dakota' --benchmark 'all'
        ./runme.py -i [[101, 102], ['Dakota', 'Slc']]

    NOTE:
    - Will only run test scripts whose names explicitly follow the convention,

        test<id>.py

    where <id> is any integer.

    TODO:
    - At '#disp test result', make sure precision of output matches that of
    MATLAB.
    - Check for failures that do not raise exceptions (for example, 'Standard
    exception'; see also jenkins/jenkins.sh). These should be counted as
    failures.
    """

    #Get ISSM_DIR variable
    ISSM_DIR = os.environ['ISSM_DIR']

    #Process options
    #GET benchmark {{{
    if benchmark not in ['all', 'nightly', 'validation', 'adolc', 'eismint', 'ismip', 'mesh', 'slc', 'thermal', 'tranforcing', 'qmu']:
        print(("runme warning: benchmark '{}' not supported, defaulting to test 'nightly'.".format(benchmark)))
        benchmark = 'nightly'
    # }}}
    #GET procedure {{{
    if procedure not in ['check', 'update', 'runFromNC']:
        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(search(r'\d+',file.split('.')[0]).group()) 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)
    # }}}
    if procedure == 'runFromNC':
        #That is a bamg test
        test_ids = test_ids.difference([119, 514])
        # that is smbGEMB format is weird for the test
        test_ids = test_ids.difference([243, 244, 252, 253])
        #those are amr runs where the index is missing from fieldnames
        test_ids = test_ids.difference([462, 463, 464, 465])
        #test247 solves for thermal and transient which makes it complex to check
        test_ids = test_ids.difference([247])
        #test 902 is running two models with different stepping
        test_ids = test_ids.difference([902])
        #I have a size issue in 517 needs investigation
        test_ids = test_ids.difference([517])

    #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 == 'slc':
        test_ids = test_ids.intersection(set(range(2001, 2500)))
    elif benchmark == 'adolc':
        test_ids = test_ids.intersection(set(range(3001, 3200)))
    elif benchmark == 'qmu':
        test_ids = test_ids.intersection(set((218, 234, 235, 417, 418, 420)).union(set(range(412, 414))))
    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)
            print(("----------------running-----------------------"))
            if procedure == 'runFromNC':
                Tmod = import_module('test{}'.format(id))
            else:
                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'))))
            elif procedure == 'runFromNC':
                print(("----------------loadingNC-----------------------"))
                mdl = loadmodel('test{}ma.nc'.format(id))
                for key in mdl.results.__dict__.keys():
                    if 'Solution' in key:
                        solvetype = split('Solution', key)[0]

                #we save the results, scrap them and solve.
                loaded_res = mdl.results
                mdl.results = []
                mdl = solve(mdl, solvetype)

                #we loop on the field_names from the nghtly test
                for k, fieldname in enumerate(Tmod.field_names):
                    try:
                        #first look for indexing
                        if search(r'\d+$', fieldname):
                            index = int(search(r'\d+$', fieldname).group()) - 1
                            fieldname = fieldname[:search(r'\d+$', fieldname).start()]
                        elif 'FirstStep' in fieldname:
                            index = 0
                            fieldname = fieldname[:search('FirstStep', fieldname).start()]
                        elif 'SecondStep' in fieldname:
                            index = 1
                            fieldname = fieldname[:search('SecondStep', fieldname).start()]
                        elif 'ThirdStep' in fieldname:
                            index = 2
                            fieldname = fieldname[:search('ThirdStep', fieldname).start()]
                        else:
                            index = 0

                        #Then check if the key exists in the loaded results
                        try:
                            reskeys = mdl.results.__dict__[solvetype + 'Solution'][index].__dict__.keys()
                        except TypeError:
                            # most probably a steady state so no subscripting
                            reskeys = mdl.results.__dict__[solvetype + 'Solution'].__dict__.keys()
                        if fieldname not in reskeys:
                            sufixes = ["P1bubble", "P1bubbleCondensed", "LliboutryDuval", "CuffeyTemperate", "SSA", "HO", "FS", "P1xP", "P2xP",
                                       'MINI', 'MINIcondensed', 'TaylorHood', 'XTaylorHood', 'LATaylorHood', 'CrouzeixRaviart', 'LACrouzeixRaviart']
                            namedifs = {'Misfits': 'J',
                                        'D': 'DamageDbar',
                                        'F': 'DamageF',
                                        'MaterialsRheologyB': 'MaterialsRheologyBbar',
                                        'SedimentWaterHead': 'SedimentHead',
                                        'EplWaterHead': 'EplHead',
                                        'SedimentWaterHeadSubstep': 'SedimentHeadSubstep',
                                        'EplWaterHeadSubstep': 'EplHeadSubstep',
                                        'Volume': 'IceVolume',
                                        'Bed': 'Base',
                                        'SMB': 'SmbMassBalance'}

                            if fieldname in namedifs.keys():
                                #Some fields are not consistent
                                fieldname = namedifs[fieldname]
                            elif any([suf in fieldname for suf in sufixes]):
                                #some test have loops that mess up with naming
                                try:
                                    sufix = sufixes[np.squeeze(np.where([suf in fieldname for suf in sufixes]))]
                                except TypeError:
                                    #probably severalmatches, we take the last one which should be the good one (Needs to be controled in the list above)
                                    sufix = sufixes[np.squeeze(np.where([suf in fieldname for suf in sufixes]))[-1]]
                                fieldname = fieldname[:search(sufix, fieldname).start()]
                            elif fieldname.endswith("P") and index == 1:
                                #we are looking for P2 but 2 as been considered as an index and so shifted by -1
                                fieldname = fieldname[:-1]
                            else:
                                # could be that the index selected above is part of the name
                                fieldname = fieldname + str(index + 1)
                        try:
                            field = mdl.results.__dict__[solvetype + 'Solution'][index].__dict__[fieldname]
                            loaded_field = loaded_res.__dict__[solvetype + 'Solution'][index].__dict__[fieldname]
                        except TypeError:
                            # most probably a steady state so no subscripting
                            try:
                                field = mdl.results.__dict__[solvetype + 'Solution'].__dict__[fieldname]
                                loaded_field = loaded_res.__dict__[solvetype + 'Solution'].__dict__[fieldname]
                            except KeyError:
                                print("WARNING: {}{} does not exist and checking will be skipped".format(fieldname, index + 1))
                                continue
                        except KeyError:
                            print("WARNING: {}{} does not exist and checking will be skipped".format(fieldname, index + 1))
                            continue

                        ref = Tmod.field_values[k]
                        #Get tolerance
                        tolerance = Tmod.field_tolerances[k]
                        #compute differences for the results computed from the nc file
                        error_diff = np.amax(np.abs(ref - field), axis=0) / (np.amax(np.abs(ref), axis=0) + float_info.epsilon)
                        if not np.isscalar(error_diff):
                            error_diff = error_diff[0]

                        #compute the differences for the results of the nc file
                        load_diff = np.amax(np.abs(np.squeeze(ref) - loaded_field), axis=0) / (np.amax(np.abs(np.squeeze(ref)), axis=0) + float_info.epsilon)
                        if not np.isscalar(load_diff):
                            load_diff = load_diff[0]

                        #disp test result
                        if (np.any(error_diff > tolerance) or np.isnan(error_diff)) and (np.any(load_diff > tolerance) or np.isnan(load_diff)):
                            if abs(error_diff - load_diff) < tolerance:
                                print(('WARNING difference: {:7.2g} > {:7.2g} test id: {} field: {}{} differs from computation but equal to saved results'.format(error_diff, tolerance, id, fieldname, index + 1)))
                            else:
                                print(('ERROR difference: {:7.2g} > {:7.2g} test id: {} field: {}{} is false in both loaded and computed results'.format(error_diff, tolerance, id, fieldname, index + 1)))
                                errorcount += 1
                                erroredtest_list.append(id)
                        elif (np.any(error_diff > tolerance) or np.isnan(error_diff)):
                            print(('ERROR   difference: {:7.2g} > {:7.2g} test id: {} test name: {} field: {}{}'.format(error_diff, tolerance, id, id_string, fieldname, index + 1)))
                            errorcount += 1
                            erroredtest_list.append(id)
                        elif (np.any(load_diff > tolerance) or np.isnan(load_diff)):
                            print(('SAVEERROR difference: {:7.2g} > {:7.2g} test id: {} test name: {} saved result : {}{}'.format(load_diff, tolerance, id, id_string, fieldname, index + 1)))
                            errorcount += 1
                            erroredtest_list.append(id)
                        else:
                            print(('SUCCESS difference: {:7.2g} < {:7.2g} test id: {} test name: {} field: {}{}'.format(error_diff, tolerance, id, id_string, fieldname, index + 1)))
                        #disp only if errors for the results

                    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)
            #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)))
                        # NOTE: 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: {:7.2g} > {:7.2g} test id: {} test name: {} field: {}'.format(error_diff, tolerance, id, id_string, fieldname)))
                            errorcount += 1
                            erroredtest_list.append(id)
                        else:
                            print(('SUCCESS difference: {:7.2g} < {:7.2g} 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']
        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=str, 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=str, 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()
