'''Utilities for test calculations'''

#  Copyright (C) 2017 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

filename_status = 'test_status'

# YL 10/11/2021 not in use
def checkExternal(interface):
    '''Check for binary executable or linked-in library'''

    from sys    import exit, stdout
    from shutil import which
    from ..base import errors
    
    chk = interface()

    if(isLinked(interface.__name__)):
        # code is directly linked in
        return True
    elif(which(chk.path) or which(chk._syscmd)): 
        # code is in path and can be executed
        return True
    else:
        # set test result to unsupported
        with open(filename_status, 'w') as fp:
            fp.write('unsupported')
        stdout.flush()
        exit(errors.UNSUPPORTED)


def setUnsupported():
    '''Set final status to pass'''

    with open(filename_status, 'w') as fp:
        fp.write('unsupported')


def setPassed():
    '''Set final status to pass'''

    print('\n '+'-'*70)
    print(' Test passed')
    print(' '+'-'*70)
    print('\n')

    with open(filename_status, 'w') as fp:
        fp.write('pass')


def setFailed():
    '''Set final status to fail'''

    from sys    import exit, stdout
    from ..base import errors

    print('\n '+'-'*70)
#    print(' Test failed')
    print(" Test failed. Exiting Py-ChemShell...")
    print(' '+'-'*70)
    print('\n')

    with open(filename_status, 'w') as fp:
        fp.write('fail')

    # TWK exit immediately on fail so result cannot be overwritten
    stdout.flush()
    exit(errors.FAILED)


def validate(calc, ref, tol=None, summarise=20):
    '''Compare the result energy with reference value'''

    from numpy import array, allclose, floating, inf, integer, ndarray, nonzero, set_printoptions

    if str(summarise).lower() in [ 'inf' ]:
        summarise = inf

    set_printoptions(threshold=summarise)

    print('\n '+'='*70)
    print(' ChemShell Test Utility')
    print(' '+'='*70+'\n')

    selectcases = { int     :('%d'    , 1),
                    integer :('%d'    , 1),
                    float   :('%16.9f', 1.0e-6),
                    floating:('%16.9f', 1.0e-6),
                    ndarray :('%16.9f', 1.0e-6),
                    bool    :('%r'    , True),
                  }

    # YL 20/02/2024: added support for a list of reference
    if type(ref) is list:
        ref = array(ref)

    # array reference
    try:
        for t in selectcases.keys():
            if issubclass(ref.dtype.type, t):
                format = selectcases[t][0]
    # scalar reference
    except:
        format = selectcases[type(ref)][0]

    # default tolerance
    try:
        if tol is None:
            # no tolerance given
            print(" No tolerance specified so default used.")
            tol = selectcases[type(ref)][1]
        elif not issubclass(ref.dtype.type, type(tol)):
            # if given tolerance doesn't match the data type
            print(" Reference type %s inconsistent with tolerance type %s. Default tolerance is used:"%(type(ref), type(tol)), selectcases[type(ref)][1], '\n')
            tol = selectcases[type(ref)][1]
    except:
        if type(tol) != type(ref):
            print(" Reference type %s inconsistent with tolerance type %s. Default tolerance is used:"%(type(ref), type(tol)), selectcases[type(ref)][1], '\n')
            tol = selectcases[type(ref)][1]
        if tol is None:
            tol = selectcases[type(ref)][1]


    def validateScalars(_calc, _ref, _format=format, _tol=tol):

        # YL 18/06/2023: _calc could typically be the results.energy object which is a size-one array and it causes warning
        #                DeprecationWarning: Conversion of an array with ndim > 0 to a scalar is deprecated, and will error in future. Ensure you extract a single element from your array before performing this operation. (Deprecated NumPy 1.25.)
        if hasattr(_calc, '__len__'):
            if len(_calc) == 1:
                _calc = _calc[0]
        print(' Tolerance : '+_format%_tol)
        print(' Calculated: '+_format%_calc)
        print(' Reference : '+_format%_ref)
        print(' Deviation : '+_format%(_calc-_ref))
        try:
            assert abs(_calc - _ref) < _tol
            return 0
        except:
            return 1


    def validateArrays(_calc, _ref, _format=format, _tol=tol, print_full=False):

        _devs = _calc - _ref

        print('\n Tolerance : '+_format%_tol)
        print('\n Calculated:\n', _calc)
        print('\n Reference :\n', _ref)
        print('\n Deviation :\n', _devs)
        try:
            assert allclose(_calc, _ref, rtol=0.0, atol=_tol)
            return 0
        except:
            print('\n Deviated elements:')
            print(' calculated                  reference                  indices')
            indices = array(nonzero(_devs))
            _c = _calc[tuple(indices)]
            _r = _ref [tuple(indices)]
            for i, index in enumerate(indices.T):
                # YL 18/05/2024: print only the deviated elements by default
                if abs(_c[i]-_r[i]) <= tol and not print_full:
                    continue
                print(' {:24s}    {:24s}  '.format(str(_c[i]), str(_r[i])), '['+','.join((len(index)*'{} ').format(*index).split())+']')
            return 1

    selectcases = {
                    bool   : validateScalars,
                    int    : validateScalars,
                    float  : validateScalars,
                    ndarray: validateArrays,
                    list   : validateArrays,
                  }

    if selectcases[type(ref)](calc, ref, _tol=tol):
        setFailed()
    else:
        setPassed()


