"""Provide estimator base class for information theoretic measures."""
import imp
import os
import importlib
import inspect
from pprint import pprint
from abc import ABCMeta, abstractmethod
import numpy as np
from . import idtxl_exceptions as ex

MODULE_EXTENSIONS = ('.py')  # ('.py', '.pyc', '.pyo')
ESTIMATOR_PREFIX = ('estimators_')


def _package_contents():
    # Return list of IDTxl modules containing estimators.
    file, pathname, description = imp.find_module(__package__)
    if file:
        raise ImportError('Not a package: %r', __package__)
    return [os.path.splitext(module)[0]
            for module in os.listdir(pathname)
            if (module.endswith(MODULE_EXTENSIONS) and
                module.startswith(ESTIMATOR_PREFIX))]


def list_estimators():
    """List all estimators available in IDTxl."""
    module_list = _package_contents()
    for m in module_list:
        module = importlib.import_module('.' + m, __package__)
        class_list = inspect.getmembers(module, inspect.isclass)
        if class_list:
            pprint(class_list)


def find_estimator(est):
    """Return estimator class.

    Return an estimator class. If input is a class, check if it implements
    methods 'estimate' and 'is_parallel' necessary for network analysis
    (see abstract class 'Estimator' for documentation). If input is a string,
    search for class with that name in IDTxl and return it.

    Args:
        est : str | Class
            name of an estimator class implemented in IDTxl or custom estimator
            class

    Returns
        Class
            Estimator class
    """
    if inspect.isclass(est):
        # Test if provided class implements the Estimator class. This
        # constraint may be relaxed in the future.
        if not np.issubclass_(est, Estimator):
            raise RuntimeError('Provided class should implement abstract class'
                               ' Estimator.')
        return est
    elif type(est) is str:
        module_list = _package_contents()
        estimator = None
        for m in module_list:
            try:
                module = importlib.import_module('.' + m, __package__)
                return getattr(module, est)
            except AttributeError:
                pass
        if not estimator:
            raise RuntimeError('Estimator {0} not found.'.format(est))
    else:
        raise TypeError('Please provide an estimator class or the name of an '
                        'estimator as string.')


class Estimator(metaclass=ABCMeta):
    """Abstract class for implementation of IDTxl estimators.

    Abstract class for implementation of IDTxl estimators. Child classes
    implement various estimators for information-theoretic measures, e.g.,
    mutual information (MI), conditional mutual information (CMI),
    active information storage (AIS), or transfer entropy (TE).

    Estimator classes implement a method 'estimate()' for the estimation from
    single or multiple data sets in parallel. Whether 'estimate()' provides
    parallel computing capabilities is indicated by the 'is_parallel()'
    method. The 'estimate_parallel()' method of this abstract class provides a
    common interface to parallel and serial methods (see docstring for
    details).

    The method 'is_analytic_null_estimator()' indicates whether the implemented
    estimator supports the generation of analytic surrogates (see docstring for
    details).
    """

    def __init__(self, settings=None):
        pass

    @abstractmethod
    def estimate(self, **vars):
        """Estimate measure for a single data set.

        The number of variables in data depends on the measure to be estimated,
        e.g., 2 for mutual information and 3 for a conditional mutual
        information.

        Each variable in vars should be a numpy array of realisations, where
        the first axis is assumed to represent realisations over samples and
        replications, while the second axis represents the variable dimension
        ([(n_samples * n_replications) x variable dimension]).

        For parallel estimators, the first axis of each variable is assumed to
        represent realisations for multiple chunks of realisations that are
        then handled in parallel. Each array has size
        [(n_samples * n_replications * n_chunks) x variable dimension]. The
        number of chunks has to be passed as an additional parameter (see for
        example the OpenCLKraskov() estimator).

        Args:
            self : Estimator class instance
                estimator
            vars: numpy arrays
                realisations of random variables

        Returns:
            float
                estimated value
        """
        pass

    @abstractmethod
    def is_parallel(self):
        """Indicate if estimator supports parallel estimation over chunks.

        Return true if the supports parallel estimation over chunks, where a
        chunk is one independent data set.

        Returns:
            bool
        """
        pass

    @abstractmethod
    def is_analytic_null_estimator(self):
        """Indicate if estimator supports analytic surrogates.

        Return true if the estimator implements estimate_surrogates_analytic()
        where data is formatted as per the estimate method for this estimator.

        Returns:
            bool
        """
        pass

    def _check_settings(self, settings=None):
        """Set default for settings dictionary.

        Check if settings dictionary is None. If None, initialise an empty
        dictionary. If not None check if type is dictionary. Function should be
        called before setting default values.
        """
        if settings is None:
            return {}
        elif type(settings) is not dict:
            raise TypeError('settings should be a dictionary.')
        else:
            return settings

    def _check_number_of_points(self, n_points):
        """Sanity check for number of points going into the estimator."""
        if (n_points - 1) <= int(self.settings['kraskov_k']):
            raise RuntimeError('Insufficient number of points ({0}) for the '
                               'requested number of nearest neighbours '
                               '(kraskov_k: {1}).'.format(
                                        n_points, self.settings['kraskov_k']))
        if (n_points - 1) <= (int(self.settings['kraskov_k']) +
                              int(self.settings['theiler_t'])):
            raise RuntimeError('Insufficient number of points ({0}) for the '
                               'requested number of nearest neighbours '
                               '(kraskov_k: {1}) and Theiler-correction '
                               '(theiler_t: {2}).'.format(
                                                n_points,
                                                self.settings['kraskov_k'],
                                                self.settings['theiler_t']))

    def _ensure_one_dim_input(self, var):
        """Make sure input arrays have one dimension.

        Check dimensions of input to AIS and TE estimators. JIDT expects one-
        dimensional arrays for these estimators, while it expects two-
        dimensional arrays for MI and CMI estimators. To make usage of all
        estimator types easier, allow both 1D- and 2D inputs for all
        estimators. Squeeze 2D-arrays if their second dimension is 1 when
        calling AIS and TE estimators (assuming that this array dimension
        represents the variable dimension).
        """
        if len(var.shape) == 2:
            if var.shape[1] == 1:
                var = np.squeeze(var)
            else:
                raise TypeError('2D input arrays must have shape[1] == 1.')
        elif len(var.shape) > 2:
            raise TypeError('Input arrays must be 1D or 2D with shape[1] == '
                            '1.')
        return var

    def _ensure_two_dim_input(self, var):
        """Make sure input arrays have two dimension.

        Check dimensions of input to MI and CMI estimators. JIDT expects two-
        dimensional arrays for these estimators, while it expects one-
        dimensional arrays for MI and CMI estimators. To make usage of all
        estimator types easier allow both 1D- and 2D inputs for all estimators.
        Add an extra dimension to 1D-arrays when calling MI and CMI estimators
        (assuming that this array dimension represents the variable dimension).
        """
        if len(var.shape) == 1:
            var = np.expand_dims(var, axis=1)
        elif len(var.shape) > 2:
            raise TypeError('Input arrays must be 1D or 2D')
        return var

    def estimate_parallel(self, n_chunks=1, re_use=None, **data):
        """Estimate measure for multiple data sets (chunks).

        Test if the estimator used provides parallel capabilities; if so,
        estimate measure for multiple data sets ('chunks') in parallel.
        Otherwise, iterate over individual chunks.

        The number of variables in data depends on the measure to be estimated,
        e.g., 2 for mutual information and 3 for a conditional mutual
        information.

        Each entry in data should be a numpy array with realisations, where the
        first axis is assumed to represent realisations over samples and
        replications and chunks, while the second axis represents the variable
        dimension ([(samples * replications) * chunks x variable dimension]).

        Each entry in data should be a numpy array with realisations, where the
        first axis is assumed to represent realisations (over chunks), while
        the second axis is the variable dimension.

        Each numpy array with realisations can hold either the realisations for
        multiple chunks or can hold the realisation for a single chunk, which
        gets replicated for parallel estimation and gets re-used for iterative
        estimation, in order to save memory. The variables for re-use are
        provided in re-use as list of dictionary keys indicating entries in
        data for re-use.

        Args:
            self : Estimator class instance
                estimator
            n_chunks : int [optional]
                number of data chunks (default=1)
            re_use : list of keys [optional}
                realisatins to be re-used (default=None)
            data: numpy arrays
                realisations of random variables

        Returns:
            numpy array
                estimated values for each chunk

        Raises:
            ex.AlgorithmExhaustedError
                Raised from self.estimate() when calculation cannot be made
        """
        assert n_chunks > 0, 'n_chunks must be positive.'
        if re_use is None:
            re_use = []

        # If the estimator supports parallel estimation, pass the variables
        # and number of chunks on to the estimator.
        if self.is_parallel():
            for k in re_use:  # multiply data for re-use
                if data[k] is not None:
                    data[k] = np.tile(data[k], (n_chunks, 1))
            return self.estimate(n_chunks=n_chunks, **data)

        # If estimator does not support parallel estimation, loop over chunks
        # and estimate iteratively for individual chunks.
        else:
            # Find arrays that have to be cut up into chunks because they are
            # not re-used.
            slice_vars = list(set(data.keys()).difference(set(re_use)))
            if not slice_vars:
                # If there is nothing to slice, we only have one chunk and can
                # return the estimate directly.
                return [self.estimate(**data)]

            # Find the number of samples, check that all data to be sliced into
            # chunks has the same number of samples.
            n_samples_total = []
            for v in slice_vars:
                if data[v] is not None:
                    n_samples_total.append(data[v].shape[0])
            assert (np.array(n_samples_total) == n_samples_total[0]).all(), (
                'No. realisations should be the same for all variables: '
                '{0}'.format(n_samples_total))
            n_samples_total = n_samples_total[0]
            assert n_samples_total is not None, (
                'All variables provided for estimation are empty.')
            assert n_samples_total % n_chunks == 0, (
                'No. chunks ({0}) does not match data length ({1}). Remainder:'
                ' {2}.'.format(n_chunks,
                               data[slice_vars[0]].shape[0],
                               data[slice_vars[0]].shape[0] % n_chunks))

            # Cut data into chunks and call estimator serially on each chunk.
            chunk_size = int(n_samples_total / n_chunks)
            idx_1 = 0
            idx_2 = chunk_size
            results = np.empty((n_chunks))
            for i in range(n_chunks):
                chunk_data = {}
                # Slice data into single chunks
                for v in slice_vars:  # NOTE: I am consciously not creating a deep copy here to save memory
                    if data[v] is not None:
                        chunk_data[v] = data[v][idx_1:idx_2, :]
                    else:
                        chunk_data[v] = data[v]
                # Collect data that is reused over chunks.
                for v in re_use:
                    if data[v] is not None:
                        assert data[v].shape[0] == chunk_size, (
                            'No. samples in variable {0} ({1}) is not equal '
                            'to chunk size ({2}).'.format(
                                v, data[v].shape[0], chunk_size))
                    chunk_data[v] = data[v]
                results[i] = self.estimate(**chunk_data)
                idx_1 = idx_2
                idx_2 += chunk_size

            return results
