"""
Utility function to facilitate testing.
"""

import os
import sys
import re
import operator
from nosetester import import_nose

__all__ = ['assert_equal', 'assert_almost_equal','assert_approx_equal',
           'assert_array_equal', 'assert_array_less', 'assert_string_equal',
           'assert_array_almost_equal', 'assert_raises', 'build_err_msg',
           'decorate_methods', 'jiffies', 'memusage', 'print_assert_equal',
           'raises', 'rand', 'rundocs', 'runstring', 'verbose']

verbose = 0

def rand(*args):
    """Returns an array of random numbers with the given shape.

    This only uses the standard library, so it is useful for testing purposes.
    """
    import random
    from numpy.core import zeros, float64
    results = zeros(args, float64)
    f = results.flat
    for i in range(len(f)):
        f[i] = random.random()
    return results

if sys.platform[:5]=='linux':
    def jiffies(_proc_pid_stat = '/proc/%s/stat'%(os.getpid()),
                _load_time=[]):
        """ Return number of jiffies (1/100ths of a second) that this
    process has been scheduled in user mode. See man 5 proc. """
        import time
        if not _load_time:
            _load_time.append(time.time())
        try:
            f=open(_proc_pid_stat,'r')
            l = f.readline().split(' ')
            f.close()
            return int(l[13])
        except:
            return int(100*(time.time()-_load_time[0]))

    def memusage(_proc_pid_stat = '/proc/%s/stat'%(os.getpid())):
        """ Return virtual memory size in bytes of the running python.
        """
        try:
            f=open(_proc_pid_stat,'r')
            l = f.readline().split(' ')
            f.close()
            return int(l[22])
        except:
            return
else:
    # os.getpid is not in all platforms available.
    # Using time is safe but inaccurate, especially when process
    # was suspended or sleeping.
    def jiffies(_load_time=[]):
        """ Return number of jiffies (1/100ths of a second) that this
    process has been scheduled in user mode. [Emulation with time.time]. """
        import time
        if not _load_time:
            _load_time.append(time.time())
        return int(100*(time.time()-_load_time[0]))
    def memusage():
        """ Return memory usage of running python. [Not implemented]"""
        raise NotImplementedError

if os.name=='nt' and sys.version[:3] > '2.3':
    # Code "stolen" from enthought/debug/memusage.py
    def GetPerformanceAttributes(object, counter, instance = None,
                                 inum=-1, format = None, machine=None):
        # NOTE: Many counters require 2 samples to give accurate results,
        # including "% Processor Time" (as by definition, at any instant, a
        # thread's CPU usage is either 0 or 100).  To read counters like this,
        # you should copy this function, but keep the counter open, and call
        # CollectQueryData() each time you need to know.
        # See http://msdn.microsoft.com/library/en-us/dnperfmo/html/perfmonpt2.asp
        # My older explanation for this was that the "AddCounter" process forced
        # the CPU to 100%, but the above makes more sense :)
        import win32pdh
        if format is None: format = win32pdh.PDH_FMT_LONG
        path = win32pdh.MakeCounterPath( (machine,object,instance, None, inum,counter) )
        hq = win32pdh.OpenQuery()
        try:
            hc = win32pdh.AddCounter(hq, path)
            try:
                win32pdh.CollectQueryData(hq)
                type, val = win32pdh.GetFormattedCounterValue(hc, format)
                return val
            finally:
                win32pdh.RemoveCounter(hc)
        finally:
            win32pdh.CloseQuery(hq)

    def memusage(processName="python", instance=0):
        # from win32pdhutil, part of the win32all package
        import win32pdh
        return GetPerformanceAttributes("Process", "Virtual Bytes",
                                        processName, instance,
                                        win32pdh.PDH_FMT_LONG, None)

def build_err_msg(arrays, err_msg, header='Items are not equal:',
                  verbose=True,
                  names=('ACTUAL', 'DESIRED')):
    msg = ['\n' + header]
    if err_msg:
        if err_msg.find('\n') == -1 and len(err_msg) < 79-len(header):
            msg = [msg[0] + ' ' + err_msg]
        else:
            msg.append(err_msg)
    if verbose:
        for i, a in enumerate(arrays):
            try:
                r = repr(a)
            except:
                r = '[repr failed]'
            if r.count('\n') > 3:
                r = '\n'.join(r.splitlines()[:3])
                r += '...'
            msg.append(' %s: %s' % (names[i], r))
    return '\n'.join(msg)

def assert_equal(actual,desired,err_msg='',verbose=True):
    """
    Raise an assertion if two objects are not equal.

    Given two objects (lists, tuples, dictionaries or numpy arrays), check
    that all elements of these objects are equal. An exception is raised at
    the first conflicting values.

    Parameters
    ----------
    actual : list, tuple, dict or ndarray
      The object to check.
    desired : list, tuple, dict or ndarray
      The expected object.
    err_msg : string
      The error message to be printed in case of failure.
    verbose : bool
      If True, the conflicting values are appended to the error message.

    Raises
    ------
    AssertionError
      If actual and desired are not equal.

    Examples
    --------
    >>> np.testing.assert_equal([4,5], [4,6]) # doctest:+ELLIPSIS
    ...
    <type 'exceptions.AssertionError'>:
    Items are not equal:
    item=1
    <BLANKLINE>
     ACTUAL: 5
     DESIRED: 6

    """
    if isinstance(desired, dict):
        if not isinstance(actual, dict) :
            raise AssertionError(repr(type(actual)))
        assert_equal(len(actual),len(desired),err_msg,verbose)
        for k,i in desired.items():
            if k not in actual :
                raise AssertionError(repr(k))
            assert_equal(actual[k], desired[k], 'key=%r\n%s' % (k,err_msg), verbose)
        return
    if isinstance(desired, (list,tuple)) and isinstance(actual, (list,tuple)):
        assert_equal(len(actual),len(desired),err_msg,verbose)
        for k in range(len(desired)):
            assert_equal(actual[k], desired[k], 'item=%r\n%s' % (k,err_msg), verbose)
        return
    from numpy.core import ndarray
    if isinstance(actual, ndarray) or isinstance(desired, ndarray):
        return assert_array_equal(actual, desired, err_msg, verbose)
    msg = build_err_msg([actual, desired], err_msg, verbose=verbose)
    if desired != actual :
        raise AssertionError(msg)

def print_assert_equal(test_string,actual,desired):
    import pprint
    try:
        assert(actual == desired)
    except AssertionError:
        import cStringIO
        msg = cStringIO.StringIO()
        msg.write(test_string)
        msg.write(' failed\nACTUAL: \n')
        pprint.pprint(actual,msg)
        msg.write('DESIRED: \n')
        pprint.pprint(desired,msg)
        raise AssertionError(msg.getvalue())

def assert_almost_equal(actual,desired,decimal=7,err_msg='',verbose=True):
    """ Raise an assertion if two items are not equal.

    I think this should be part of unittest.py

    The test is equivalent to abs(desired-actual) < 0.5 * 10**(-decimal)
    """
    from numpy.core import ndarray
    if isinstance(actual, ndarray) or isinstance(desired, ndarray):
        return assert_array_almost_equal(actual, desired, decimal, err_msg)
    msg = build_err_msg([actual, desired], err_msg, verbose=verbose)
    if round(abs(desired - actual),decimal) != 0 :
        raise AssertionError(msg)


def assert_approx_equal(actual,desired,significant=7,err_msg='',verbose=True):
    """ Raise an assertion if two items are not
        equal.  I think this should be part of unittest.py
        Approximately equal is defined as the number of significant digits
        correct
    """
    import math
    actual, desired = map(float, (actual, desired))
    if desired==actual:
        return
    # Normalized the numbers to be in range (-10.0,10.0)
    scale = float(pow(10,math.floor(math.log10(0.5*(abs(desired)+abs(actual))))))
    try:
        sc_desired = desired/scale
    except ZeroDivisionError:
        sc_desired = 0.0
    try:
        sc_actual = actual/scale
    except ZeroDivisionError:
        sc_actual = 0.0
    msg = build_err_msg([actual, desired], err_msg,
                header='Items are not equal to %d significant digits:' %
                                 significant,
                verbose=verbose)
    if math.fabs(sc_desired - sc_actual) >= pow(10.,-(significant-1)) :
        raise AssertionError(msg)

def assert_array_compare(comparison, x, y, err_msg='', verbose=True,
                         header=''):
    from numpy.core import asarray, isnan, any
    x = asarray(x)
    y = asarray(y)

    def isnumber(x):
        return x.dtype.char in '?bhilqpBHILQPfdgFDG'

    try:
        cond = (x.shape==() or y.shape==()) or x.shape == y.shape
        if not cond:
            msg = build_err_msg([x, y],
                                err_msg
                                + '\n(shapes %s, %s mismatch)' % (x.shape,
                                                                  y.shape),
                                verbose=verbose, header=header,
                                names=('x', 'y'))
            if not cond :
                raise AssertionError(msg)

        if (isnumber(x) and isnumber(y)) and (any(isnan(x)) or any(isnan(y))):
            # Handling nan: we first check that x and y have the nan at the
            # same locations, and then we mask the nan and do the comparison as
            # usual.
            xnanid = isnan(x)
            ynanid = isnan(y)
            try:
                assert_array_equal(xnanid, ynanid)
            except AssertionError:
                msg = build_err_msg([x, y],
                                    err_msg
                                    + '\n(x and y nan location mismatch %s, ' \
                                    '%s mismatch)' % (xnanid, ynanid),
                                    verbose=verbose, header=header,
                                    names=('x', 'y'))
            val = comparison(x[~xnanid], y[~ynanid])
        else:
            val = comparison(x,y)
        if isinstance(val, bool):
            cond = val
            reduced = [0]
        else:
            reduced = val.ravel()
            cond = reduced.all()
            reduced = reduced.tolist()
        if not cond:
            match = 100-100.0*reduced.count(1)/len(reduced)
            msg = build_err_msg([x, y],
                                err_msg
                                + '\n(mismatch %s%%)' % (match,),
                                verbose=verbose, header=header,
                                names=('x', 'y'))
            if not cond :
                raise AssertionError(msg)
    except ValueError:
        msg = build_err_msg([x, y], err_msg, verbose=verbose, header=header,
                            names=('x', 'y'))
        raise ValueError(msg)

def assert_array_equal(x, y, err_msg='', verbose=True):
    assert_array_compare(operator.__eq__, x, y, err_msg=err_msg,
                         verbose=verbose, header='Arrays are not equal')

def assert_array_almost_equal(x, y, decimal=6, err_msg='', verbose=True):
    from numpy.core import around
    def compare(x, y):
        return around(abs(x-y),decimal) <= 10.0**(-decimal)
    assert_array_compare(compare, x, y, err_msg=err_msg, verbose=verbose,
                         header='Arrays are not almost equal')

def assert_array_less(x, y, err_msg='', verbose=True):
    assert_array_compare(operator.__lt__, x, y, err_msg=err_msg,
                         verbose=verbose,
                         header='Arrays are not less-ordered')

def runstring(astr, dict):
    exec astr in dict

def assert_string_equal(actual, desired):
    # delay import of difflib to reduce startup time
    import difflib

    if not isinstance(actual, str) :
        raise AssertionError(`type(actual)`)
    if not isinstance(desired, str):
        raise AssertionError(`type(desired)`)
    if re.match(r'\A'+desired+r'\Z', actual, re.M): return
    diff = list(difflib.Differ().compare(actual.splitlines(1), desired.splitlines(1)))
    diff_list = []
    while diff:
        d1 = diff.pop(0)
        if d1.startswith('  '):
            continue
        if d1.startswith('- '):
            l = [d1]
            d2 = diff.pop(0)
            if d2.startswith('? '):
                l.append(d2)
                d2 = diff.pop(0)
            if not d2.startswith('+ ') :
                raise AssertionError(`d2`)
            l.append(d2)
            d3 = diff.pop(0)
            if d3.startswith('? '):
                l.append(d3)
            else:
                diff.insert(0, d3)
            if re.match(r'\A'+d2[2:]+r'\Z', d1[2:]):
                continue
            diff_list.extend(l)
            continue
        raise AssertionError(`d1`)
    if not diff_list:
        return
    msg = 'Differences in strings:\n%s' % (''.join(diff_list)).rstrip()
    if actual != desired :
        raise AssertionError(msg)


def rundocs(filename=None):
    """ Run doc string tests found in filename.
    """
    import doctest, imp
    if filename is None:
        f = sys._getframe(1)
        filename = f.f_globals['__file__']
    name = os.path.splitext(os.path.basename(filename))[0]
    path = [os.path.dirname(filename)]
    file, pathname, description = imp.find_module(name, path)
    try:
        m = imp.load_module(name, file, pathname, description)
    finally:
        file.close()
    if sys.version[:3]<'2.4':
        doctest.testmod(m, verbose=False)
    else:
        tests = doctest.DocTestFinder().find(m)
        runner = doctest.DocTestRunner(verbose=False)
        for test in tests:
            runner.run(test)
    return


def raises(*args,**kwargs):
    nose = import_nose()
    return nose.tools.raises(*args,**kwargs)

def assert_raises(*args,**kwargs):
    nose = import_nose()
    return nose.tools.assert_raises(*args,**kwargs)

def decorate_methods(cls, decorator, testmatch=None):
    ''' Apply decorator to all methods in class matching testmatch

    Parameters
    ----------
    cls : class
        Class to decorate methods for
    decorator : function
        Decorator to apply to methods
    testmatch : compiled regexp or string to compile to regexp
        Decorators are applied if testmatch.search(methodname)
        is not None.  Default value is
        re.compile(r'(?:^|[\\b_\\.%s-])[Tt]est' % os.sep)
        (the default for nose)
    '''
    if testmatch is None:
        testmatch = re.compile(r'(?:^|[\\b_\\.%s-])[Tt]est' % os.sep)
    else:
        testmatch = re.compile(testmatch)
    cls_attr = cls.__dict__

    # delayed import to reduce startup time
    from inspect import isfunction

    methods = filter(isfunction, cls_attr.values())
    for function in methods:
        try:
            if hasattr(function, 'compat_func_name'):
                funcname = function.compat_func_name
            else:
                funcname = function.__name__
        except AttributeError:
            # not a function
            continue
        if testmatch.search(funcname) and not funcname.startswith('_'):
            setattr(cls, funcname, decorator(function))
    return


def measure(code_str,times=1,label=None):
    """ Return elapsed time for executing code_str in the
    namespace of the caller for given times.
    """
    frame = sys._getframe(1)
    locs,globs = frame.f_locals,frame.f_globals

    code = compile(code_str,
                   'Test name: %s ' % label,
                   'exec')
    i = 0
    elapsed = jiffies()
    while i < times:
        i += 1
        exec code in globs,locs
    elapsed = jiffies() - elapsed
    return 0.01*elapsed