"""
Module for testing the ``serpentTools`` package
"""
from functools import wraps
from unittest import TestCase, SkipTest
from logging import NOTSET

from numpy import stack
from numpy.testing import assert_allclose
from serpentTools.messages import (
    DictHandler, __logger__, removeHandler, addHandler,
)
from serpentTools.utils import checkScipy


def computeMeansErrors(*arrays):
    """
    Return the matrices of element-wise means and standard deviations

    This function assumes that:

    1. Each element in ``arrays`` is a numpy array
    2. Each element in ``arrays`` is of equal dimensionality

    The arrays are all stacked to create a new array with one
    additional axis at index 0. The mean and standard deviation are
    computed along this new axis so that the returned arrays
    are of equal dimensionality to the incoming arrays

    Parameters
    ----------
    arrays: iterable
        Arrays to be stacked

    Returns
    -------
    numpy.ndarray
        Element-wise mean of all incoming arrays
    numpy.ndarray
        Element-wise standard deviation of all incoming arrays
    """
    workMat = stack(arrays)
    return workMat.mean(axis=0), workMat.std(axis=0)


def compareDictOfArrays(expected, actual, fmtMsg=None, rtol=0, atol=0,
                        testCase=None):
    """
    Compare a dictionary of arrays.

    Parameters
    ----------
    expected: dict
        Dictionary of expected data
    actual: dict
        Dictionary of actual data.
    fmtMsg: str
        Message to be passed as the error message. Formatted with
        ``.format(key=key)``, where ``key`` is the specific key
        where the arrays were too different.
    rtol: float
        Relative tolerance for arrays
    atol: float
        Absolute tolerance for arrays
    testCase: None or :class:`unittest.TestCase`
        If given, use the ``testCase.assertSetEqual`` to compare keys

    Raises
    ------
    AssertionError:
        If the keys in both dictionaries differ, or if any
        one array in ``actual`` is too different from it's counterpart
        in ``expected``.
    """
    fmtMsg = fmtMsg or "Key: {key}"
    eKeys = set(expected.keys())
    aKeys = set(actual.keys())
    if isinstance(testCase, TestCase):
        testCase.assertSetEqual(eKeys, aKeys)
    else:
        in1Not2 = eKeys.difference(aKeys)
        in2Not1 = aKeys.difference(eKeys)
        errMsg = ''
        if any(in1Not2):
            errMsg += ('Keys in expected not actual: {}\n'
                       .format(', '.join(in1Not2)))
        if any(in2Not1):
            errMsg += ('Keys in actual not expected: {}\n'
                       .format(', '.join(in2Not1)))
        if errMsg:
            raise AssertionError(errMsg)
    for key, value in expected.items():
        actualValue = actual[key]
        assert_allclose(value, actualValue, rtol=rtol, atol=atol,
                        err_msg=fmtMsg.format(key=key))


class LoggerMixin(object):
    """
    Mixin class captures log messages

    Attributes
    ----------
    handler: :class:`serpentTools.messages.DictHandler`
        Logging handler that stores messages in a
        :attr:`serpentTools.messages.DictHandler.logMessages`
        dictionary according to level.
    """
    def __init__(self):
        self.__old = []
        self.handler = None

    def attach(self, level=NOTSET):
        """
        Attach the :class:`serpentTools.messages.DictHandler`

        Removes all :class:`logging.Handler` objects from the
        old logger, and puts them back when :class:`detach` is
        called

        Parameters
        ----------
        level: int
            Initial level to apply to handler
        """
        self.handler = DictHandler(level)
        self.__old = __logger__.handlers
        for handler in self.__old:
            removeHandler(handler)
        addHandler(self.handler)

    def detach(self):
        """Restore the original handers to the main logger"""
        if self.handler is None:
            raise AttributeError("Handler not set. Possibly not attached.")
        removeHandler(self.handler)
        for handler in self.__old:
            addHandler(handler)
        self.handler = None
        self.__old = []

    def msgInLogs(self, level, msg, partial=False):
        """
        Determine if the message is contained in the logs

        Parameters
        ----------
        level: str
            Level under which this message was posted.
            Must be a key in the
            :attr:`~serpentTools.messages.DictHandler.logMessages`
            on the :attr:`handler` for this class
        msg: str
            Message to be found in the logs.
        partial: bool
            If this evaluates to true, then search through each
            ``message`` in `logMessages` and return ``True`` if
            ``msg in message``. Otherwise, look for exact matches