import collections
from enum import Enum
from functools import reduce
from operator import mul
from typing import Iterable, Tuple, TypeVar, Dict, NamedTuple, Optional, Union
import warnings
import logging

import numpy as np
import six
from scipy.linalg import block_diag
from scipy.optimize import OptimizeResult
from collections import abc

from autofit.mapper.variable import Variable, VariableData

def try_getitem(value, index, default=None):
    try:
        return value[index]
    except TypeError:
        return default

class LogWarnings(warnings.catch_warnings):
    def __init__(self, *, module=None, messages=None, action=None, logger=logging.warning):
        super().__init__(record=True, module=module)
        self.messages = [] if messages is None else messages
        self.log = []
        self.action = action 
        self.logger = logger

    def log_warning(self, warn):
        self.log.append(warn)
        warn_message = f"{warn.filename}:{warn.lineno}: {warn.message}"
        self.messages.append(warn_message)
        self.logger(warn_message)

    def __enter__(self):
        self.log = super().__enter__()
        self._module._showwarnmsg_impl = self.log_warning
        if self.action:
            warnings.simplefilter(self.action)

        return self
        

def is_variable(v, *args):
    return isinstance(v, Variable)


def is_iterable(arg):
    return isinstance(arg, abc.Iterable) and not isinstance(
        arg, six.string_types
    )


def nested_filter(func, *args):
    """ Iterates through a potentially nested set of list, tuples and dictionaries, 
    recursively looping through the structure and returning the arguments
    that func return true on, 

    Example
    -------
    >>> list(nested_filter(
    ...     lambda x, *args: x==2,
    ...     [1, (2, 3), [3, 2, {1, 2}]]
    ... ))
    [(2,), (2,), (2,)]

    >>> list(nested_filter(
    ...     lambda x, *args: x==2,
    ...     [1, (2, 3), [3, 2, {1, 2}]],
    ...     [1, ('a', 3), [3, 'b', {1, 'c'}]]
    ... ))
    [(2, 'a'), (2, 'b'), (2, 'c')]
    """
    out, *_ = args
    if isinstance(out, dict):
        for k in out:
            yield from nested_filter(func, *(out[k] for out in args))
    elif is_iterable(out):
        for elems in zip(*args):
            yield from nested_filter(func, *elems)
    else:
        if func(*args):
            yield args


def nested_update(out, to_replace: dict, replace_keys=False):
    """
    Given a potentially nested set of list, tuples and dictionaries, recursively loop through the structure and
    replace any values that appear in the dict to_replace
    can set to replace dictionary keys optionally,

    Example
    -------
    >>> nested_update([1, (2, 3), [3, 2, {1, 2}]], {2: 'a'})
    [1, ('a', 3), [3, 'a', {1, 'a'}]]

    >>> nested_update([{2: 2}], {2: 'a'})
    [{2: 'a'}]

    >>> nested_update([{2: 2}], {2: 'a'}, True)
    [{'a': 'a'}]
    """
    try:
        return to_replace[out]
    except KeyError:
        pass

    if isinstance(out, dict):
        if replace_keys:
            return type(out)(
                {
                    nested_update(k, to_replace, replace_keys): nested_update(
                        v, to_replace, replace_keys
                    )
                    for k, v in out.items()
                }
            )
        else:
            return type(out)(
                {k: nested_update(v, to_replace, replace_keys) for k, v in out.items()}
            )
    elif is_iterable(out):
        return type(out)(nested_update(elem, to_replace, replace_keys) for elem in out)

    return out


class StatusFlag(Enum):
    FAILURE = 0
    SUCCESS = 1
    NO_CHANGE = 2
    BAD_PROJECTION = 3

    @classmethod
    def get_flag(cls, success, n_iter):
        if success:
            if n_iter > 0:
                return cls.SUCCESS
            else:
                return cls.NO_CHANGE

        return cls.FAILURE


class Status(NamedTuple):
    success: bool = True
    messages: Tuple[str, ...] = ()
    updated: bool = True
    flag: StatusFlag = StatusFlag.SUCCESS

    def __bool__(self):
        return self.success

    def __str__(self):
        if self.success:
            return "Optimisation succeeded"
        return f"Optimisation failed: {self.messages}"


class FlattenArrays(dict):
    """
    >>> shapes = FlattenArrays(a=(1, 2), b=(2, 3))
    >>> shapes
    FlattenArrays(a=(1, 2), b=(2, 3))
    >>> shapes.flatten(
        a = np.arange(2).reshape(1, 2),
        b = np.arange(6).reshape(2, 3)**2)
    array([ 0,  1,  0,  1,  4,  9, 16, 25])
    >>> shapes.unflatten(
        [ 0,  1,  0,  1,  4,  9, 16, 25])
    {'a': array([[0, 1]]), 'b': array([[ 0,  1,  4],
        [ 9, 16, 25]])}
    """

    def __init__(self, dict_: Dict[Variable, Tuple[int, ...]]):
        super().__init__()

        self.update(dict_)
        self.splits = np.cumsum([np.prod(s) for s in self.values()], dtype=int)
        self.inds = [
            slice(i0, i1)
            for i0, i1 in
            # np.arange(i0, i1, dtype=int) for i0, i1 in
            zip(np.r_[0, self.splits[:-1]], self.splits)
        ]
        self.sizes = {k: np.prod(s, dtype=int) for k, s in self.items()}
        self.k_inds = dict(zip(self, self.inds))

    @classmethod
    def from_arrays(cls, arrays: Dict[str, np.ndarray]) -> "FlattenArrays":
        return cls({k: np.shape(arr) for k, arr in arrays.items()})

    def flatten(self, arrays_dict: Dict[Variable, np.ndarray]) -> np.ndarray:
        assert all(np.shape(arrays_dict[k]) == shape for k, shape in self.items())
        return np.concatenate([np.ravel(arrays_dict[k]) for k in self.keys()])

    def extract(self, key, flat, ndim=None):
        if ndim is None:
            ndim = len(flat.shape)

        ind = self.k_inds[key]
        return flat[(ind,) * ndim]

    def unflatten(self, arr: np.ndarray, ndim=None) -> Dict[str, np.ndarray]:
        arr = np.asanyarray(arr)
        if ndim is None:
            ndim = arr.ndim
        arrays = [arr[(ind,) * ndim] for ind in self.inds]
        arr_shapes = [arr.shape[ndim:] for arr in arrays]
        return VariableData({
            k: arr.reshape(shape * ndim + arr_shape)
            if shape or arr_shape
            else arr.item()
            for (k, shape), arr_shape, arr in zip(self.items(), arr_shapes, arrays)
        })

    def flatten2d(self, values: Dict[Variable, np.ndarray]) -> np.ndarray:
        assert all(np.shape(values[k]) == shape * 2 for k, shape in self.items())

        return block_diag(
            *(np.reshape(values[k], (n, n)) for k, n in self.sizes.items())
        )

    unflatten2d = unflatten

    def __repr__(self):
        shapes = ", ".join(map("{0[0]}={0[1]}".format, self.items()))
        return f"{type(self).__name__}({shapes})"

    @property
    def size(self):
        return self.splits[-1]


class OptResult(NamedTuple):
    mode: Dict[Variable, np.ndarray]
    hess_inv: Dict[Variable, np.ndarray]
    log_norm: float
    full_hess_inv: np.ndarray
    result: OptimizeResult
    status: Status = Status()


def gen_subsets(n, x, n_iters=None, rng=None):
    """
    Generates random subsets of length n of the array x, if the elements of
    x are unique then each subset will not contain repeated elements. Each 
    element is guaranteed to reappear after at most 2*len(x) new elements. 

    If `x` is a multi-dimensional array, it is only shuffled along its
first index.

    if x is an integer, generate subsets of ``np.arange(x)``.

    generates n_iters subsets before stopping. If n_iters is None then
    generates random subsets for ever

    rng is an optionally passed random number generator

    Examples
    --------
    >>> list(gen_subsets(3, 5, n_iters=3))
    [array([0, 2, 3]), array([1, 4, 0]), array([2, 3, 4])]
    >>> list(gen_subsets(3, [1,10,5,3], n_iters=3))
    [array([ 5, 10,  1]), array([3, 5, 1]), array([10,  3,  5])]
    """
    rng = rng or np.random.default_rng()
    x_shuffled = rng.permutation(x)
    tot = len(x_shuffled)

    i = 0 
    stop = tot - n + 1
    iters = iter(int, 1) if n_iters is None else range(n_iters)
    for j in iters:
        if i < stop:
            yield x_shuffled[i : i + n]
            i += n
        else:
            x_shuffled = np.r_[x_shuffled[i:], rng.permutation(x_shuffled[:i])]
            yield x_shuffled[:n]
            i = n

def gen_dict(dict_gen):
    """
    Examples
    --------
    >>> list(gen_dict({1: gen_subsets(3, 4, 3), 2: gen_subsets(2, 5, 3)}))
    [{1: array([2, 1, 3]), 2: array([2, 0])},
     {1: array([0, 3, 1]), 2: array([3, 1])},
     {1: array([2, 0, 1]), 2: array([4, 2])}]
    """
    keys = tuple(dict_gen.keys())
    for val in zip(*dict_gen.values()):
        yield dict(zip(keys, val))


_M = TypeVar("_M")


def prod(iterable: Iterable[_M], *arg: Tuple[_M]) -> _M:
    """calculates the product of the passed iterable,
    much like sum, if a second argument is passed,
    this is the initial value of the calculation

    Examples
    --------
    >>> prod(range(1, 3))
    2

    >>> prod(range(1, 3), 2.)
    4.
    """
    iterable = list(iterable)
    return reduce(mul, iterable, *arg)


def r2_score(y_true, y_pred, axis=None):
    y_true = np.asanyarray(y_true)
    y_pred = np.asanyarray(y_pred)

    mse = np.square(y_true - y_pred).mean(axis=axis)
    var = y_true.var(axis=axis)

    return 1 - mse / var


def propagate_uncertainty(cov: np.ndarray, jac: np.ndarray) -> np.ndarray:
    """Propagates the uncertainty of a covariance matrix given the
    passed Jacobian

    If the variable arrays are multidimensional then will output in
    the shape of the arrays

    see https://en.wikipedia.org/wiki/Propagation_of_uncertainty
    """
    cov = np.asanyarray(cov)

    var_ndim = cov.ndim // 2
    det_ndim = jac.ndim - var_ndim
    det_shape, var_shape = jac.shape[:det_ndim], jac.shape[det_ndim:]
    assert var_shape == cov.shape[:var_ndim] == cov.shape[var_ndim:]

    var_size = np.prod(var_shape, dtype=int)
    det_size = np.prod(det_shape, dtype=int)

    cov2d = cov.reshape((var_size, var_size))
    jac2d = jac.reshape((det_size, var_size))

    det_cov2d = np.linalg.multi_dot((jac2d, cov2d, jac2d.T))
    det_cov = det_cov2d.reshape(det_shape + det_shape)
    return det_cov


def rescale_to_artists(artists, ax=None):
    import matplotlib.pyplot as plt

    ax = ax or plt.gca()
    while True:
        r = ax.figure.canvas.get_renderer()
        extents = [
            t.get_window_extent(renderer=r).transformed(ax.transData.inverted())
            for t in artists
        ]
        min_extent = np.min([e.min for e in extents], axis=0)
        max_extent = np.max([e.max for e in extents], axis=0)
        min_lim, max_lim = zip(ax.get_xlim(), ax.get_ylim())

        # Sometimes the window doesn't always rescale first time around
        if (min_extent < min_lim).any() or (max_extent > max_lim).any():
            extent = max_extent - min_extent
            max_extent += extent * 0.05
            min_extent -= extent * 0.05
            xlim, ylim = zip(
                np.minimum(min_lim, min_extent), np.maximum(max_lim, max_extent)
            )
            ax.set_xlim(*xlim)
            ax.set_ylim(*ylim)
        else:
            break

    return xlim, ylim


# These may no longer be needed?
def add_arrays(*arrays: np.ndarray) -> np.ndarray:
    """Sums over broadcasting multidimensional arrays
    whilst preserving the total sum

    a = np.arange(10).reshape(1, 2, 1, 5)
    b = np.arange(8).reshape(2, 2, 2, 1)

    >>> add_arrays(a, b).sum()
    73.0
    >>> add_arrays(a, b).shape
    (2, 2, 2, 5)
    >>> a.sum() + b.sum()
    73
    """
    b = np.broadcast(*arrays)
    return sum(a * np.size(a) / b.size for a in arrays)


Axis = Optional[Union[bool, int, Tuple[int, ...]]]


def aggregate(array: np.ndarray, axis: Axis = None, **kwargs) -> np.ndarray:
    """
    aggregates the values of array

    if axis is False then aggregate returns the unmodified array

    otherwise aggrate returns np.sum(array, axis=axis, **kwargs)
    """
    if axis is False:
        return array
        
    return np.sum(array, axis=axis, **kwargs)


def diag(array: np.ndarray, *ds: Tuple[int, ...]) -> np.ndarray:
    array = np.asanyarray(array)
    d1 = array.shape
    if ds:
        ds = (d1,) + ds
    else:
        ds = (d1, d1)

    out = np.zeros(sum(ds, ()))
    diag_inds = tuple(map(np.ravel, (i for d in ds for i in np.indices(d))))
    out[diag_inds] = array.ravel()
    return out
