"""
Utility functions related to concat
"""

import numpy as np
import pandas._libs.tslib as tslib
from pandas import compat
from pandas.core.dtypes.common import (
    is_categorical_dtype,
    is_sparse,
    is_datetimetz,
    is_datetime64_dtype,
    is_timedelta64_dtype,
    is_period_dtype,
    is_object_dtype,
    is_bool_dtype,
    is_dtype_equal,
    _NS_DTYPE,
    _TD_DTYPE)
from pandas.core.dtypes.generic import (
    ABCDatetimeIndex, ABCTimedeltaIndex,
    ABCPeriodIndex)

def get_dtype_kinds(l):
    """
    Parameters
    ----------
    l : list of arrays

    Returns
    -------
    a set of kinds that exist in this list of arrays
    """

    typs = set()
    for arr in l:

        dtype = arr.dtype
        if is_categorical_dtype(dtype):
            typ = 'category'
        elif is_sparse(arr):
            typ = 'sparse'
        elif is_datetimetz(arr):
            # if to_concat contains different tz,
            # the result must be object dtype
            typ = str(arr.dtype)
        elif is_datetime64_dtype(dtype):
            typ = 'datetime'
        elif is_timedelta64_dtype(dtype):
            typ = 'timedelta'
        elif is_object_dtype(dtype):
            typ = 'object'
        elif is_bool_dtype(dtype):
            typ = 'bool'
        elif is_period_dtype(dtype):
            typ = str(arr.dtype)
        else:
            typ = dtype.kind
        typs.add(typ)
    return typs

def _get_series_result_type(result):
    """
    return appropriate class of Series concat
    input is either dict or array-like
    """
    if isinstance(result, dict):
        # concat Series with axis 1
        if all(is_sparse(c) for c in compat.itervalues(result)):
            from pandas.core.sparse.api import SparseDataFrame
            return SparseDataFrame
        else:
            from pandas.core.frame import DataFrame
            return DataFrame

    elif is_sparse(result):
        # concat Series with axis 1
        from pandas.core.sparse.api import SparseSeries
        return SparseSeries
    else:
        from pandas.core.series import Series
        return Series

def _get_frame_result_type(result, objs):
    """
    return appropriate class of DataFrame-like concat
    if any block is SparseBlock, return SparseDataFrame
    otherwise, return 1st obj
    """
    if any(b.is_sparse for b in result.blocks):
        from pandas.core.sparse.api import SparseDataFrame
        return SparseDataFrame
    else:
        return objs[0]

def _concat_compat(to_concat, axis=0):
    """
    provide concatenation of an array of arrays each of which is a single
    'normalized' dtypes (in that for example, if it's object, then it is a
    non-datetimelike and provide a combined dtype for the resulting array that
    preserves the overall dtype if possible)

    Parameters
    ----------
    to_concat : array of arrays
    axis : axis to provide concatenation

    Returns
    -------
    a single array, preserving the combined dtypes
    """

    # filter empty arrays
    # 1-d dtypes always are included here
    def is_nonempty(x):
        try:
            return x.shape[axis] > 0
        except Exception:
            return True

    nonempty = [x for x in to_concat if is_nonempty(x)]

    # If all arrays are empty, there's nothing to convert, just short-cut to
    # the concatenation, #3121.
    #
    # Creating an empty array directly is tempting, but the winnings would be
    # marginal given that it would still require shape & dtype calculation and
    # np.concatenate which has them both implemented is compiled.

    typs = get_dtype_kinds(to_concat)

    _contains_datetime = any(typ.startswith('datetime') for typ in typs)
    _contains_period = any(typ.startswith('period') for typ in typs)

    if 'category' in typs:
        # this must be priort to _concat_datetime,
        # to support Categorical + datetime-like
        return _concat_categorical(to_concat, axis=axis)

    elif _contains_datetime or 'timedelta' in typs or _contains_period:
        return _concat_datetime(to_concat, axis=axis, typs=typs)

    # these are mandated to handle empties as well
    elif 'sparse' in typs:
        return _concat_sparse(to_concat, axis=axis, typs=typs)

    if not nonempty:
        # we have all empties, but may need to coerce the result dtype to
        # object if we have non-numeric type operands (numpy would otherwise
        # cast this to float)
        typs = get_dtype_kinds(to_concat)
        if len(typs) != 1:

            if (not len(typs - set(['i', 'u', 'f'])) or
                    not len(typs - set(['bool', 'i', 'u']))):
                # let numpy coerce
                pass
            else:
                # coerce to object
                to_concat = [x.astype('object') for x in to_concat]

    return np.concatenate(to_concat, axis=axis)

def _concat_categorical(to_concat, axis=0):
    """Concatenate an object/categorical array of arrays, each of which is a
    single dtype

    Parameters
    ----------
    to_concat : array of arrays
    axis : int
        Axis to provide concatenation in the current implementation this is
        always 0, e.g. we only have 1D categoricals

    Returns
    -------
    Categorical
        A single array, preserving the combined dtypes
    """

    def _concat_asobject(to_concat):
        to_concat = [x.get_values() if is_categorical_dtype(x.dtype)
                     else x.ravel() for x in to_concat]
        res = _concat_compat(to_concat)
        if axis == 1:
            return res.reshape(1, len(res))
        else:
            return res

    # we could have object blocks and categoricals here
    # if we only have a single categoricals then combine everything
    # else its a non-compat categorical
    categoricals = [x for x in to_concat if is_categorical_dtype(x.dtype)]

    # validate the categories
    if len(categoricals) != len(to_concat):
        pass
    else:
        # when all categories are identical
        first = to_concat[0]
        if all(first.is_dtype_equal(other) for other in to_concat[1:]):
            return union_categoricals(categoricals)

    return _concat_asobject(to_concat)

def union_categoricals(to_union, sort_categories=False, ignore_order=False):
    """
    Combine list-like of Categorical-like, unioning categories. All
    categories must have the same dtype.

    .. versionadded:: 0.19.0

    Parameters
    ----------
    to_union : list-like of Categorical, CategoricalIndex,
               or Series with dtype='category'
    sort_categories : boolean, default False
        If true, resulting categories will be lexsorted, otherwise
        they will be ordered as they appear in the data.
    ignore_order: boolean, default False
        If true, the ordered attribute of the Categoricals will be ignored.
        Results in an unordered categorical.

        .. versionadded:: 0.20.0

    Returns
    -------
    result : Categorical

    Raises
    ------
    TypeError
        - all inputs do not have the same dtype
        - all inputs do not have the same ordered property
        - all inputs are ordered and their categories are not identical
        - sort_categories=True and Categoricals are ordered
    ValueError
        Empty list of categoricals passed
    """
    from pandas import Index, Categorical, CategoricalIndex, Series

    if len(to_union) == 0:
        raise ValueError('No Categoricals to union')

    def _maybe_unwrap(x):
        if isinstance(x, (CategoricalIndex, Series)):
            return x.values
        elif isinstance(x, Categorical):
            return x
        else:
            raise TypeError("all components to combine must be Categorical")

    to_union = [_maybe_unwrap(x) for x in to_union]
    first = to_union[0]

    if not all(is_dtype_equal(other.categories.dtype, first.categories.dtype)
               for other in to_union[1:]):
        raise TypeError("dtype of categories must be the same")

    ordered = False
    if all(first.is_dtype_equal(other) for other in to_union[1:]):
        # identical categories - fastpath
        categories = first.categories
        ordered = first.ordered
        new_codes = np.concatenate([c.codes for c in to_union])

        if sort_categories and not ignore_order and ordered:
            raise TypeError("Cannot use sort_categories=True with "
                            "ordered Categoricals")

        if sort_categories and not categories.is_monotonic_increasing:
            categories = categories.sort_values()
            indexer = categories.get_indexer(first.categories)

            from pandas.core.algorithms import take_1d
            new_codes = take_1d(indexer, new_codes, fill_value=-1)
    elif ignore_order or all(not c.ordered for c in to_union):
        # different categories - union and recode
        cats = first.categories.append([c.categories for c in to_union[1:]])
        categories = Index(cats.unique())
        if sort_categories:
            categories = categories.sort_values()

        new_codes = []
        for c in to_union:
            if len(c.categories) > 0:
                indexer = categories.get_indexer(c.categories)

                from pandas.core.algorithms import take_1d
                new_codes.append(take_1d(indexer, c.codes, fill_value=-1))
            else:
                # must be all NaN
                new_codes.append(c.codes)
        new_codes = np.concatenate(new_codes)
    else:
        # ordered - to show a proper error message
        if all(c.ordered for c in to_union):
            msg = ("to union ordered Categoricals, "
                   "all categories must be the same")
            raise TypeError(msg)
        else:
            raise TypeError('Categorical.ordered must be the same')

    if ignore_order:
        ordered = False

    return Categorical(new_codes, categories=categories, ordered=ordered,
                       fastpath=True)

def _concat_datetime(to_concat, axis=0, typs=None):
    """
    provide concatenation of an datetimelike array of arrays each of which is a
    single M8[ns], datetimet64[ns, tz] or m8[ns] dtype

    Parameters
    ----------
    to_concat : array of arrays
    axis : axis to provide concatenation
    typs : set of to_concat dtypes

    Returns
    -------
    a single array, preserving the combined dtypes
    """

    def convert_to_pydatetime(x, axis):
        # coerce to an object dtype

        # if dtype is of datetimetz or timezone
        if x.dtype.kind == _NS_DTYPE.kind:
            if getattr(x, 'tz', None) is not None:
                x = x.asobject.values
            else:
                shape = x.shape
                x = tslib.ints_to_pydatetime(x.view(np.int64).ravel(),
                                             box=True)
                x = x.reshape(shape)

        elif x.dtype == _TD_DTYPE:
            shape = x.shape
            x = tslib.ints_to_pytimedelta(x.view(np.int64).ravel(), box=True)
            x = x.reshape(shape)

        if axis == 1:
            x = np.atleast_2d(x)
        return x

    if typs is None:
        typs = get_dtype_kinds(to_concat)

    # must be single dtype
    if len(typs) == 1:
        _contains_datetime = any(typ.startswith('datetime') for typ in typs)
        _contains_period = any(typ.startswith('period') for typ in typs)

        if _contains_datetime:

            if 'datetime' in typs:
                new_values = np.concatenate([x.view(np.int64) for x in
                                             to_concat], axis=axis)
                return new_values.view(_NS_DTYPE)
            else:
                # when to_concat has different tz, len(typs) > 1.
                # thus no need to care
                return _concat_datetimetz(to_concat)

        elif 'timedelta' in typs:
            new_values = np.concatenate([x.view(np.int64) for x in to_concat],
                                        axis=axis)
            return new_values.view(_TD_DTYPE)

        elif _contains_period:
            # PeriodIndex must be handled by PeriodIndex,
            # Thus can't meet this condition ATM
            # Must be changed when we adding PeriodDtype
            raise NotImplementedError

    # need to coerce to object
    to_concat = [convert_to_pydatetime(x, axis) for x in to_concat]
    return np.concatenate(to_concat, axis=axis)

def _concat_datetimetz(to_concat, name=None):
    """
    concat DatetimeIndex with the same tz
    all inputs must be DatetimeIndex
    it is used in DatetimeIndex.append also
    """
    # do not pass tz to set because tzlocal cannot be hashed
    if len(set([str(x.dtype) for x in to_concat])) != 1:
        raise ValueError('to_concat must have the same tz')
    tz = to_concat[0].tz
    # no need to localize because internal repr will not be changed
    new_values = np.concatenate([x.asi8 for x in to_concat])
    return to_concat[0]._simple_new(new_values, tz=tz, name=name)

def _concat_index_asobject(to_concat, name=None):
    """
    concat all inputs as object. DatetimeIndex, TimedeltaIndex and
    PeriodIndex are converted to object dtype before concatenation
    """

    klasses = ABCDatetimeIndex, ABCTimedeltaIndex, ABCPeriodIndex
    to_concat = [x.asobject if isinstance(x, klasses) else x
                 for x in to_concat]

    from pandas import Index
    self = to_concat[0]
    attribs = self._get_attributes_dict()
    attribs['name'] = name

    to_concat = [x._values if isinstance(x, Index) else x
                 for x in to_concat]
    return self._shallow_copy_with_infer(np.concatenate(to_concat), **attribs)

def _concat_sparse(to_concat, axis=0, typs=None):
    """
    provide concatenation of an sparse/dense array of arrays each of which is a
    single dtype

    Parameters
    ----------
    to_concat : array of arrays
    axis : axis to provide concatenation
    typs : set of to_concat dtypes

    Returns
    -------
    a single array, preserving the combined dtypes
    """

    from pandas.core.sparse.array import SparseArray, _make_index

    def convert_sparse(x, axis):
        # coerce to native type
        if isinstance(x, SparseArray):
            x = x.get_values()
        x = x.ravel()
        if axis > 0:
            x = np.atleast_2d(x)
        return x

    if typs is None:
        typs = get_dtype_kinds(to_concat)

    if len(typs) == 1:
        # concat input as it is if all inputs are sparse
        # and have the same fill_value
        fill_values = set(c.fill_value for c in to_concat)
        if len(fill_values) == 1:
            sp_values = [c.sp_values for c in to_concat]
            indexes = [c.sp_index.to_int_index() for c in to_concat]

            indices = []
            loc = 0
            for idx in indexes:
                indices.append(idx.indices + loc)
                loc += idx.length
            sp_values = np.concatenate(sp_values)
            indices = np.concatenate(indices)
            sp_index = _make_index(loc, indices, kind=to_concat[0].sp_index)

            return SparseArray(sp_values, sparse_index=sp_index,
                               fill_value=to_concat[0].fill_value)

    # input may be sparse / dense mixed and may have different fill_value
    # input must contain sparse at least 1
    sparses = [c for c in to_concat if is_sparse(c)]
    fill_values = [c.fill_value for c in sparses]
    sp_indexes = [c.sp_index for c in sparses]

    # densify and regular concat
    to_concat = [convert_sparse(x, axis) for x in to_concat]
    result = np.concatenate(to_concat, axis=axis)

    if not len(typs - set(['sparse', 'f', 'i'])):
        # sparsify if inputs are sparse and dense numerics
        # first sparse input's fill_value and SparseIndex is used
        result = SparseArray(result.ravel(), fill_value=fill_values[0],
                             kind=sp_indexes[0])
    else:
        # coerce to object if needed
        result = result.astype('object')
    return result


"""
Tests for impulse responses of time series

Author: Chad Fulton
License: Simplified-BSD
"""
import warnings

import numpy as np
import pandas as pd
from scipy.stats import ortho_group
import pytest
from numpy.testing import assert_, assert_allclose

from statsmodels.tools.sm_exceptions import EstimationWarning
from statsmodels.tsa.statespace import (mlemodel, sarimax, structural, varmax,
                                        dynamic_factor)

def test_sarimax():
    # AR(1)
    mod = sarimax.SARIMAX([0], order=(1, 0, 0))
    phi = 0.5
    actual = mod.impulse_responses([phi, 1], steps=10)
    desired = np.r_[[phi**i for i in range(11)]]
    assert_allclose(actual, desired)

    # MA(1)
    mod = sarimax.SARIMAX([0], order=(0, 0, 1))
    theta = 0.5
    actual = mod.impulse_responses([theta, 1], steps=10)
    desired = np.r_[1, theta, [0]*9]
    assert_allclose(actual, desired)

    # ARMA(2, 2) + constant
    # Stata:
    # webuse lutkepohl2
    # arima dln_inc, arima(2, 0, 2)
    # irf create irf1, set(irf1) step(10)
    # irf table irf
    params = [.01928228, -.03656216, .7588994,
              .27070341, -.72928328, .01122177**0.5]
    mod = sarimax.SARIMAX([0], order=(2, 0, 2), trend='c')
    actual = mod.impulse_responses(params, steps=10)
    desired = [1, .234141, .021055, .17692, .00951, .133917, .002321, .101544,
               -.001951, .077133, -.004301]
    assert_allclose(actual, desired, atol=1e-6)

    # SARIMAX(1,1,1)x(1,0,1,4) + constant + exog
    # Stata:
    # webuse lutkepohl2
    # gen exog = _n^2
    # arima inc exog, arima(1,1,1) sarima(1,0,1,4)
    # irf create irf2, set(irf2) step(10)
    # irf table irf
    params = [.12853289, 12.207156, .86384742, -.71463236,
              .81878967, -.9533955, 14.043884**0.5]
    exog = np.arange(1, 92)**2
    mod = sarimax.SARIMAX(np.zeros(91), order=(1, 1, 1),
                          seasonal_order=(1, 0, 1, 4), trend='c', exog=exog,
                          simple_differencing=True)
    actual = mod.impulse_responses(params, steps=10)
    desired = [1, .149215, .128899, .111349, -.038417, .063007, .054429,
               .047018, -.069598, .018641, .016103]
    assert_allclose(actual, desired, atol=1e-6)

def test_structural():
    steps = 10

    # AR(1)
    mod = structural.UnobservedComponents([0], autoregressive=1)
    phi = 0.5
    actual = mod.impulse_responses([1, phi], steps)
    desired = np.r_[[phi**i for i in range(steps + 1)]]
    assert_allclose(actual, desired)

    # ARX(1)
    # This is adequately tested in test_simulate.py, since in the time-varying
    # case `impulse_responses` just calls `simulate`

    # Irregular
    mod = structural.UnobservedComponents([0], 'irregular')
    actual = mod.impulse_responses([1.], steps)
    assert_allclose(actual, 0)

    # Fixed intercept
    # (in practice this is a deterministic constant, because an irregular
    #  component must be added)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        mod = structural.UnobservedComponents([0], 'fixed intercept')
    actual = mod.impulse_responses([1.], steps)
    assert_allclose(actual, 0)

    # Deterministic constant
    mod = structural.UnobservedComponents([0], 'deterministic constant')
    actual = mod.impulse_responses([1.], steps)
    assert_allclose(actual, 0)

    # Local level
    mod = structural.UnobservedComponents([0], 'local level')
    actual = mod.impulse_responses([1., 1.], steps)
    assert_allclose(actual, 1)

    # Random walk
    mod = structural.UnobservedComponents([0], 'random walk')
    actual = mod.impulse_responses([1.], steps)
    assert_allclose(actual, 1)

    # Fixed slope
    # (in practice this is a deterministic trend, because an irregular
    #  component must be added)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        mod = structural.UnobservedComponents([0], 'fixed slope')
    actual = mod.impulse_responses([1.], steps)
    assert_allclose(actual, 0)

    # Deterministic trend
    mod = structural.UnobservedComponents([0], 'deterministic trend')
    actual = mod.impulse_responses([1.], steps)
    assert_allclose(actual, 0)

    # Local linear deterministic trend
    mod = structural.UnobservedComponents(
        [0], 'local linear deterministic trend')
    actual = mod.impulse_responses([1., 1.], steps)
    assert_allclose(actual, 1)

    # Random walk with drift
    mod = structural.UnobservedComponents([0], 'random walk with drift')
    actual = mod.impulse_responses([1.], steps)
    assert_allclose(actual, 1)

    # Local linear trend
    mod = structural.UnobservedComponents([0], 'local linear trend')
    # - shock the level
    actual = mod.impulse_responses([1., 1., 1.], steps)
    assert_allclose(actual, 1)
    # - shock the trend
    actual = mod.impulse_responses([1., 1., 1.], steps, impulse=1)
    assert_allclose(actual, np.arange(steps + 1))

    # Smooth trend
    mod = structural.UnobservedComponents([0], 'smooth trend')
    actual = mod.impulse_responses([1., 1.], steps)
    assert_allclose(actual, np.arange(steps + 1))

    # Random trend
    mod = structural.UnobservedComponents([0], 'random trend')
    actual = mod.impulse_responses([1., 1.], steps)
    assert_allclose(actual, np.arange(steps + 1))

    # Seasonal (deterministic)
    mod = structural.UnobservedComponents([0], 'irregular', seasonal=2,
                                          stochastic_seasonal=False)
    actual = mod.impulse_responses([1.], steps)
    assert_allclose(actual, 0)

    # Seasonal (stochastic)
    mod = structural.UnobservedComponents([0], 'irregular', seasonal=2)
    actual = mod.impulse_responses([1., 1.], steps)
    desired = np.r_[1, np.tile([-1, 1], steps // 2)]
    assert_allclose(actual, desired)

    # Cycle (deterministic)
    mod = structural.UnobservedComponents([0], 'irregular', cycle=True)
    actual = mod.impulse_responses([1., 1.2], steps)
    assert_allclose(actual, 0)

    # Cycle (stochastic)
    mod = structural.UnobservedComponents([0], 'irregular', cycle=True,
                                          stochastic_cycle=True)
    actual = mod.impulse_responses([1., 1., 1.2], steps=10)
    x1 = [np.cos(1.2), np.sin(1.2)]
    x2 = [-np.sin(1.2), np.cos(1.2)]
    T = np.array([x1, x2])
    desired = np.zeros(steps + 1)
    states = [1, 0]
    for i in range(steps + 1):
        desired[i] += states[0]
        states = np.dot(T, states)
    assert_allclose(actual, desired)

def test_varmax():
    steps = 10

    # Clear warnings
    varmax.__warningregistry__ = {}

    # VAR(2) - single series
    mod1 = varmax.VARMAX([[0]], order=(2, 0), trend='n')
    mod2 = sarimax.SARIMAX([0], order=(2, 0, 0))
    actual = mod1.impulse_responses([0.5, 0.2, 1], steps)
    desired = mod2.impulse_responses([0.5, 0.2, 1], steps)
    assert_allclose(actual, desired)

    # VMA(2) - single series
    mod1 = varmax.VARMAX([[0]], order=(0, 2), trend='n')
    mod2 = sarimax.SARIMAX([0], order=(0, 0, 2))
    actual = mod1.impulse_responses([0.5, 0.2, 1], steps)
    desired = mod2.impulse_responses([0.5, 0.2, 1], steps)
    assert_allclose(actual, desired)

    # VARMA(2, 2) - single series
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        mod1 = varmax.VARMAX([[0]], order=(2, 2), trend='n')
    mod2 = sarimax.SARIMAX([0], order=(2, 0, 2))
    actual = mod1.impulse_responses([0.5, 0.2, 0.1, -0.2, 1], steps)
    desired = mod2.impulse_responses([0.5, 0.2, 0.1, -0.2, 1], steps)
    assert_allclose(actual, desired)

    # VARMA(2, 2) + trend - single series
    warning = EstimationWarning
    match = r'VARMA\(p,q\) models is not'
    with pytest.warns(warning, match=match):
        mod1 = varmax.VARMAX([[0]], order=(2, 2), trend='c')
    mod2 = sarimax.SARIMAX([0], order=(2, 0, 2), trend='c')
    actual = mod1.impulse_responses([10, 0.5, 0.2, 0.1, -0.2, 1], steps)
    desired = mod2.impulse_responses([10, 0.5, 0.2, 0.1, -0.2, 1], steps)
    assert_allclose(actual, desired)

    # VAR(2) + constant
    # Stata:
    # webuse lutkepohl2
    # var dln_inv dln_inc, lags(1/2)
    # irf create irf3, set(irf3) step(10)
    # irf table irf
    # irf table oirf
    params = [-.00122728, .01503679,
              -.22741923, .71030531, -.11596357, .51494891,
              .05974659, .02094608, .05635125, .08332519,
              .04297918, .00159473, .01096298]
    irf_00 = [1, -.227419, -.021806, .093362, -.001875, -.00906, .009605,
              .001323, -.001041, .000769, .00032]
    irf_01 = [0, .059747, .044015, -.008218, .007845, .004629, .000104,
              .000451, .000638, .000063, .000042]
    irf_10 = [0, .710305, .36829, -.065697, .084398, .043038, .000533,
              .005755, .006051, .000548, .000526]
    irf_11 = [1, .020946, .126202, .066419, .028735, .007477, .009878,
              .003287, .001266, .000986, .0005]
    oirf_00 = [0.042979, -0.008642, -0.00035, 0.003908, 0.000054, -0.000321,
               0.000414, 0.000066, -0.000035, 0.000034, 0.000015]
    oirf_01 = [0.001595, 0.002601, 0.002093, -0.000247, 0.000383, 0.000211,
               0.00002, 0.000025, 0.000029, 4.30E-06, 2.60E-06]
    oirf_10 = [0, 0.007787, 0.004037, -0.00072, 0.000925, 0.000472, 5.80E-06,
               0.000063, 0.000066, 6.00E-06, 5.80E-06]
    oirf_11 = [0.010963, 0.00023, 0.001384, 0.000728, 0.000315, 0.000082,
               0.000108, 0.000036, 0.000014, 0.000011, 5.50E-06]

    mod = varmax.VARMAX([[0, 0]], order=(2, 0), trend='c')

    # IRFs
    actual = mod.impulse_responses(params, steps, impulse=0)
    assert_allclose(actual, np.c_[irf_00, irf_01], atol=1e-6)

    actual = mod.impulse_responses(params, steps, impulse=1)
    assert_allclose(actual, np.c_[irf_10, irf_11], atol=1e-6)

    # Orthogonalized IRFs
    actual = mod.impulse_responses(params, steps, impulse=0,
                                   orthogonalized=True)
    assert_allclose(actual, np.c_[oirf_00, oirf_01], atol=1e-6)

    actual = mod.impulse_responses(params, steps, impulse=1,
                                   orthogonalized=True)
    assert_allclose(actual, np.c_[oirf_10, oirf_11], atol=1e-6)

    # VARMA(2, 2) + trend + exog
    # TODO: This is just a smoke test
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        mod = varmax.VARMAX(
            np.random.normal(size=(steps, 2)), order=(2, 2), trend='c',
            exog=np.ones(steps), enforce_stationarity=False,
            enforce_invertibility=False)
    mod.impulse_responses(mod.start_params, steps)

def test_dynamic_factor():
    steps = 10
    exog = np.random.normal(size=steps)

    # DFM: 2 series, AR(2) factor
    mod1 = dynamic_factor.DynamicFactor([[0, 0]], k_factors=1, factor_order=2)
    mod2 = sarimax.SARIMAX([0], order=(2, 0, 0))
    actual = mod1.impulse_responses([-0.9, 0.8, 1., 1., 0.5, 0.2], steps)
    desired = mod2.impulse_responses([0.5, 0.2, 1], steps)
    assert_allclose(actual[:, 0], -0.9 * desired)
    assert_allclose(actual[:, 1], 0.8 * desired)

    # DFM: 2 series, AR(2) factor, exog
    mod1 = dynamic_factor.DynamicFactor(np.zeros((steps, 2)), k_factors=1,
                                        factor_order=2, exog=exog)
    mod2 = sarimax.SARIMAX([0], order=(2, 0, 0))
    actual = mod1.impulse_responses(
        [-0.9, 0.8, 5, -2, 1., 1., 0.5, 0.2], steps)
    desired = mod2.impulse_responses([0.5, 0.2, 1], steps)
    assert_allclose(actual[:, 0], -0.9 * desired)
    assert_allclose(actual[:, 1], 0.8 * desired)

    # DFM, 3 series, VAR(2) factor, exog, error VAR
    # TODO: This is just a smoke test
    mod = dynamic_factor.DynamicFactor(np.random.normal(size=(steps, 3)),
                                       k_factors=2, factor_order=2, exog=exog,
                                       error_order=2, error_var=True,
                                       enforce_stationarity=False)
    mod.impulse_responses(mod.start_params, steps)

def test_time_varying_ssm():
    # Create an ad-hoc time-varying model
    mod = sarimax.SARIMAX([0] * 11, order=(1, 0, 0))
    mod.update([0.5, 1.0])
    T = np.zeros((1, 1, 11))
    T[..., :5] = 0.5
    T[..., 5:] = 0.2
    mod['transition'] = T

    irfs = mod.ssm.impulse_responses()
    desired = np.cumprod(np.r_[1, [0.5] * 4, [0.2] * 5]).reshape(10, 1)
    assert_allclose(irfs, desired)

class TVSS(mlemodel.MLEModel):
    """
    Time-varying state space model for testing

    This creates a state space model with randomly generated time-varying
    system matrices. When used in a test, that test should use
    `reset_randomstate` to ensure consistent test runs.
    """
    def __init__(self, endog, _k_states=None):
        k_states = 2
        k_posdef = 2
        # Allow subcasses to add additional states
        if _k_states is None:
            _k_states = k_states
        super(TVSS, self).__init__(endog, k_states=_k_states,
                                   k_posdef=k_posdef, initialization='diffuse')

        self['obs_intercept'] = np.random.normal(
            size=(self.k_endog, self.nobs))
        self['design'] = np.zeros((self.k_endog, self.k_states, self.nobs))
        self['transition'] = np.zeros(
            (self.k_states, self.k_states, self.nobs))
        self['selection'] = np.zeros(
            (self.k_states, self.ssm.k_posdef, self.nobs))
        self['design', :, :k_states, :] = np.random.normal(
            size=(self.k_endog, k_states, self.nobs))
        # For the transition matrices, enforce eigenvalues not too far outside
        # unit circle. Otherwise, the random draws will often lead to large
        # eigenvalues that cause the covariance matrices to blow up to huge
        # values during long periods of missing data, which leads to numerical
        # problems and essentially spurious test failures
        D = [np.diag(d)
             for d in np.random.uniform(-1.1, 1.1, size=(self.nobs, k_states))]
        Q = ortho_group.rvs(k_states, size=self.nobs)
        self['transition', :k_states, :k_states, :] = (
            Q @ D @ Q.transpose(0, 2, 1)).transpose(1, 2, 0)
        self['selection', :k_states, :, :] = np.random.normal(
            size=(k_states, self.ssm.k_posdef, self.nobs))

        # Need to make sure the covariances are positive definite
        H05 = np.random.normal(size=(self.k_endog, self.k_endog, self.nobs))
        Q05 = np.random.normal(
            size=(self.ssm.k_posdef, self.ssm.k_posdef, self.nobs))
        H = np.zeros_like(H05)
        Q = np.zeros_like(Q05)
        for t in range(self.nobs):
            H[..., t] = np.dot(H05[..., t], H05[..., t].T)
            Q[..., t] = np.dot(Q05[..., t], Q05[..., t].T)
        self['obs_cov'] = H
        self['state_cov'] = Q

    def clone(self, endog, exog=None, **kwargs):
        mod = self.__class__(endog, **kwargs)

        for key in self.ssm.shapes.keys():
            if key in ['obs', 'state_intercept']:
                continue
            n = min(self.nobs, mod.nobs)
            mod[key, ..., :n] = self.ssm[key, ..., :n]

        return mod

def test_time_varying_in_sample(reset_randomstate):
    mod = TVSS(np.zeros((10, 2)))

    # Compute the max number of in-sample IRFs
    irfs = mod.impulse_responses([], steps=mod.nobs - 1)
    # Compute the same thing, but now with explicit anchor argument
    irfs_anchor = mod.impulse_responses([], steps=mod.nobs - 1, anchor=0)

    # Cumulative IRFs
    cirfs = mod.impulse_responses([], steps=mod.nobs - 1, cumulative=True)
    # Orthogonalized IRFs
    oirfs = mod.impulse_responses([], steps=mod.nobs - 1, orthogonalized=True)
    # Cumulative, orthogonalized IRFs
    coirfs = mod.impulse_responses([], steps=mod.nobs - 1, cumulative=True,
                                   orthogonalized=True)

    # Compute IRFs manually
    Z = mod['design']
    T = mod['transition']
    R = mod['selection']
    Q = mod['state_cov', ..., 0]
    L = np.linalg.cholesky(Q)

    desired_irfs = np.zeros((mod.nobs - 1, 2)) * np.nan
    desired_oirfs = np.zeros((mod.nobs - 1, 2)) * np.nan
    tmp = R[..., 0]
    for i in range(1, mod.nobs):
        desired_irfs[i - 1] = Z[:, :, i].dot(tmp)[:, 0]
        desired_oirfs[i - 1] = Z[:, :, i].dot(tmp).dot(L)[:, 0]
        tmp = T[:, :, i].dot(tmp)

    assert_allclose(irfs, desired_irfs)
    assert_allclose(irfs_anchor, desired_irfs)

    assert_allclose(cirfs, np.cumsum(desired_irfs, axis=0))
    assert_allclose(oirfs, desired_oirfs)
    assert_allclose(coirfs, np.cumsum(desired_oirfs, axis=0))

def test_time_varying_out_of_sample(reset_randomstate):
    mod = TVSS(np.zeros((10, 2)))

    # Compute all in-sample IRFs and also one out-of-sample IRF
    new_Z = np.random.normal(size=mod['design', :, :, -1].shape)
    new_T = np.random.normal(size=mod['transition', :, :, -1].shape)
    irfs = mod.impulse_responses(
        [], steps=mod.nobs, design=new_Z[:, :, None],
        transition=new_T[:, :, None])
    # Compute the same thing, but now with explicit anchor argument
    irfs_anchor = mod.impulse_responses(
        [], steps=mod.nobs, anchor=0, design=new_Z[:, :, None],
        transition=new_T[:, :, None])

    # Cumulative IRFs
    cirfs = mod.impulse_responses([], steps=mod.nobs, design=new_Z[:, :, None],
                                  transition=new_T[:, :, None],
                                  cumulative=True)
    # Orthogonalized IRFs
    oirfs = mod.impulse_responses([], steps=mod.nobs, design=new_Z[:, :, None],
                                  transition=new_T[:, :, None],
                                  orthogonalized=True)
    # Cumulative, orthogonalized IRFs
    coirfs = mod.impulse_responses(
        [], steps=mod.nobs, design=new_Z[:, :, None],
        transition=new_T[:, :, None], cumulative=True, orthogonalized=True)

    # Compute IRFs manually
    Z = mod['design']
    T = mod['transition']
    R = mod['selection']
    Q = mod['state_cov', ..., 0]
    L = np.linalg.cholesky(Q)

    desired_irfs = np.zeros((mod.nobs, 2)) * np.nan
    desired_oirfs = np.zeros((mod.nobs, 2)) * np.nan
    tmp = R[..., 0]
    for i in range(1, mod.nobs):
        desired_irfs[i - 1] = Z[:, :, i].dot(tmp)[:, 0]
        desired_oirfs[i - 1] = Z[:, :, i].dot(tmp).dot(L)[:, 0]
        tmp = T[:, :, i].dot(tmp)
    desired_irfs[mod.nobs - 1] = new_Z.dot(tmp)[:, 0]
    desired_oirfs[mod.nobs - 1] = new_Z.dot(tmp).dot(L)[:, 0]

    assert_allclose(irfs, desired_irfs)
    assert_allclose(irfs_anchor, desired_irfs)

    assert_allclose(cirfs, np.cumsum(desired_irfs, axis=0))
    assert_allclose(oirfs, desired_oirfs)
    assert_allclose(coirfs, np.cumsum(desired_oirfs, axis=0))

def test_time_varying_in_sample_anchored(reset_randomstate):
    mod = TVSS(np.zeros((10, 2)))

    # Compute the max number of in-sample IRFs
    anchor = 2
    irfs = mod.impulse_responses(
        [], steps=mod.nobs - 1 - anchor, anchor=anchor)

    # Cumulative IRFs
    cirfs = mod.impulse_responses(
        [], steps=mod.nobs - 1 - anchor, anchor=anchor,
        cumulative=True)
    # Orthogonalized IRFs
    oirfs = mod.impulse_responses(
        [], steps=mod.nobs - 1 - anchor, anchor=anchor,
        orthogonalized=True)
    # Cumulative, orthogonalized IRFs
    coirfs = mod.impulse_responses(
        [], steps=mod.nobs - 1 - anchor, anchor=anchor,
        cumulative=True, orthogonalized=True)

    # Compute IRFs manually
    Z = mod['design']
    T = mod['transition']
    R = mod['selection']
    Q = mod['state_cov', ..., anchor]
    L = np.linalg.cholesky(Q)

    desired_irfs = np.zeros((mod.nobs - anchor - 1, 2)) * np.nan
    desired_oirfs = np.zeros((mod.nobs - anchor - 1, 2)) * np.nan
    tmp = R[..., anchor]
    for i in range(1, mod.nobs - anchor):
        desired_irfs[i - 1] = Z[:, :, i + anchor].dot(tmp)[:, 0]
        desired_oirfs[i - 1] = Z[:, :, i + anchor].dot(tmp).dot(L)[:, 0]
        tmp = T[:, :, i + anchor].dot(tmp)

    assert_allclose(irfs, desired_irfs)

    assert_allclose(cirfs, np.cumsum(desired_irfs, axis=0))
    assert_allclose(oirfs, desired_oirfs)
    assert_allclose(coirfs, np.cumsum(desired_oirfs, axis=0))

def test_time_varying_out_of_sample_anchored(reset_randomstate):
    mod = TVSS(np.zeros((10, 2)))

    # Compute all in-sample IRFs and also one out-of-sample IRF
    anchor = 2

    new_Z = mod['design', :, :, -1]
    new_T = mod['transition', :, :, -1]
    irfs = mod.impulse_responses(
        [], steps=mod.nobs - anchor, anchor=anchor, design=new_Z[:, :, None],
        transition=new_T[:, :, None])

    # Cumulative IRFs
    cirfs = mod.impulse_responses(
        [], steps=mod.nobs - anchor, anchor=anchor,
        design=new_Z[:, :, None], transition=new_T[:, :, None],
        cumulative=True)
    # Orthogonalized IRFs
    oirfs = mod.impulse_responses(
        [], steps=mod.nobs - anchor, anchor=anchor,
        design=new_Z[:, :, None], transition=new_T[:, :, None],
        orthogonalized=True)
    # Cumulative, orthogonalized IRFs
    coirfs = mod.impulse_responses(
        [], steps=mod.nobs - anchor, anchor=anchor,
        design=new_Z[:, :, None], transition=new_T[:, :, None],
        cumulative=True, orthogonalized=True)

    # Compute IRFs manually
    Z = mod['design']
    T = mod['transition']
    R = mod['selection']
    Q = mod['state_cov', ..., anchor]
    L = np.linalg.cholesky(Q)

    desired_irfs = np.zeros((mod.nobs - anchor, 2)) * np.nan
    desired_oirfs = np.zeros((mod.nobs - anchor, 2)) * np.nan
    tmp = R[..., anchor]
    for i in range(1, mod.nobs - anchor):
        desired_irfs[i - 1] = Z[:, :, i + anchor].dot(tmp)[:, 0]
        desired_oirfs[i - 1] = Z[:, :, i + anchor].dot(tmp).dot(L)[:, 0]
        tmp = T[:, :, i + anchor].dot(tmp)
    desired_irfs[mod.nobs - anchor - 1] = new_Z.dot(tmp)[:, 0]
    desired_oirfs[mod.nobs - anchor - 1] = new_Z.dot(tmp).dot(L)[:, 0]

    assert_allclose(irfs, desired_irfs)

    assert_allclose(cirfs, np.cumsum(desired_irfs, axis=0))
    assert_allclose(oirfs, desired_oirfs)
    assert_allclose(coirfs, np.cumsum(desired_oirfs, axis=0))

def test_time_varying_out_of_sample_anchored_end(reset_randomstate):
    mod = TVSS(np.zeros((10, 2)))

    # Cannot compute the any in-sample IRFs when anchoring at the end
    with pytest.raises(ValueError, match='Model has time-varying'):
        mod.impulse_responses([], steps=2, anchor='end')

    # Compute two out-of-sample IRFs
    new_Z = np.random.normal(size=mod['design', :, :, -2:].shape)
    new_T = np.random.normal(size=mod['transition', :, :, -2:].shape)
    irfs = mod.impulse_responses([], steps=2, anchor='end',
                                 design=new_Z, transition=new_T)

    # Cumulative IRFs
    cirfs = mod.impulse_responses(
        [], steps=2, anchor='end', design=new_Z, transition=new_T,
        cumulative=True)
    # Orthogonalized IRFs
    oirfs = mod.impulse_responses(
        [], steps=2, anchor='end', design=new_Z, transition=new_T,
        orthogonalized=True)
    # Cumulative, orthogonalized IRFs
    coirfs = mod.impulse_responses(
        [], steps=2, anchor='end', design=new_Z, transition=new_T,
        cumulative=True, orthogonalized=True)

    # Compute IRFs manually
    R = mod['selection']
    Q = mod['state_cov', ..., -1]
    L = np.linalg.cholesky(Q)

    desired_irfs = np.zeros((2, 2)) * np.nan
    desired_oirfs = np.zeros((2, 2)) * np.nan
    # desired[0] = 0
    # Z_{T+1} R_T
    tmp = R[..., -1]
    desired_irfs[0] = new_Z[:, :, 0].dot(tmp)[:, 0]
    desired_oirfs[0] = new_Z[:, :, 0].dot(tmp).dot(L)[:, 0]
    # T_{T+1} R_T
    tmp = new_T[..., 0].dot(tmp)
    # Z_{T+2} T_{T+1} R_T
    desired_irfs[1] = new_Z[:, :, 1].dot(tmp)[:, 0]
    desired_oirfs[1] = new_Z[:, :, 1].dot(tmp).dot(L)[:, 0]

    assert_allclose(irfs, desired_irfs)

    assert_allclose(cirfs, np.cumsum(desired_irfs, axis=0))
    assert_allclose(oirfs, desired_oirfs)
    assert_allclose(coirfs, np.cumsum(desired_oirfs, axis=0))

def test_pandas_univariate_rangeindex():
    # Impulse responses have RangeIndex
    endog = pd.Series(np.zeros(1))
    mod = sarimax.SARIMAX(endog)
    res = mod.filter([0.5, 1.])

    actual = res.impulse_responses(2)
    desired = pd.Series([1., 0.5, 0.25])
    assert_allclose(res.impulse_responses(2), desired)
    assert_(actual.index.equals(desired.index))

def test_pandas_univariate_dateindex():
    # Impulse responses still have RangeIndex (i.e. aren't wrapped with dates)
    ix = pd.date_range(start='2000', periods=1, freq='M')
    endog = pd.Series(np.zeros(1), index=ix)
    mod = sarimax.SARIMAX(endog)
    res = mod.filter([0.5, 1.])

    actual = res.impulse_responses(2)
    desired = pd.Series([1., 0.5, 0.25])
    assert_allclose(res.impulse_responses(2), desired)
    assert_(actual.index.equals(desired.index))

def test_pandas_multivariate_rangeindex():
    # Impulse responses have RangeIndex
    endog = pd.DataFrame(np.zeros((1, 2)))
    mod = varmax.VARMAX(endog, trend='n')
    res = mod.filter([0.5, 0., 0., 0.2, 1., 0., 1.])

    actual = res.impulse_responses(2)
    desired = pd.DataFrame([[1., 0.5, 0.25], [0., 0., 0.]]).T
    assert_allclose(actual, desired)
    assert_(actual.index.equals(desired.index))

def test_pandas_multivariate_dateindex():
    # Impulse responses still have RangeIndex (i.e. aren't wrapped with dates)
    ix = pd.date_range(start='2000', periods=1, freq='M')
    endog = pd.DataFrame(np.zeros((1, 2)), index=ix)
    mod = varmax.VARMAX(endog, trend='n')
    res = mod.filter([0.5, 0., 0., 0.2, 1., 0., 1.])

    actual = res.impulse_responses(2)
    desired = pd.DataFrame([[1., 0.5, 0.25], [0., 0., 0.]]).T
    assert_allclose(actual, desired)
    assert_(actual.index.equals(desired.index))

def test_pandas_anchor():
    # Test that anchor with dates works
    ix = pd.date_range(start='2000', periods=10, freq='M')
    endog = pd.DataFrame(np.zeros((10, 2)), index=ix)
    mod = TVSS(endog)
    res = mod.filter([])

    desired = res.impulse_responses(2, anchor=1)

    # Anchor to date
    actual = res.impulse_responses(2, anchor=ix[1])
    assert_allclose(actual, desired)
    assert_(actual.index.equals(desired.index))

    # Anchor to negative index
    actual = res.impulse_responses(2, anchor=-9)
    assert_allclose(actual, desired)
    assert_(actual.index.equals(desired.index))


from __future__ import print_function
import pandas as pd
import pickle as pk
import cv2
import os
import re
import progressbar
import imutils
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from mahotas.features import haralick
import json
from sklearn.decomposition import PCA
plt.style.use('seaborn-dark')

def get_fg_bg_rects(fg):
    b, g, r, a = cv2.split(fg)
    h, w = fg.shape[:2]
    h -= 1
    w -= 1 # to avoid indexing problems
    rectDims = [10, 10] # h, w of rectangles
    hRects = h / rectDims[0]
    wRects = w / rectDims[1]
    fgRects = []
    bgRects = []
    for i in range(wRects):
        for j in range(hRects):
            pt1 = (i * rectDims[0], j * rectDims[1])
            pt2 = ((i + 1) * rectDims[0], (j + 1) * rectDims[1])
            # alpha is 255 over the part of the dog
            if a[pt1[1], pt1[0]] == 255 and a[pt2[1], pt2[0]] == 255:
                fgRects.append([pt1, pt2])
                #cv2.rectangle(fgcp, pt1, pt2, [0, 0, 255], 2) # for debugging
            elif a[pt1[1], pt1[0]] == 0 and a[pt2[1], pt2[0]] == 0:
                bgRects.append([pt1, pt2])
                #cv2.rectangle(bgcp, pt1, pt2, [0, 0, 255], 2)

    return fgRects, bgRects

def get_avg_hara(im, rects):
    # returns the haralick texture averaged over all rectangles in an image
    if len(rects)==0:
        return None
    im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    hara = 0
    for r in rects:
        # slice images as: img[y0:y1, x0:x1]
        hara += haralick(im[r[0][1]:r[1][1], r[0][0]:r[1][0]]).mean(0)
    hara /= (len(rects))
    return hara

def make_hara_map(im, rects):
    # draws heatmap of haralick texture PCA dim1 variance
    if len(rects)==0:
        return None
    gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    hara = []
    for r in rects:
        # slice images as: img[y0:y1, x0:x1]
        hara.append(pcaFG.transform(haralick(im[r[0][1]:r[1][1], r[0][0]:r[1][0]]).mean(0).reshape(1, -1)))

    hara = np.array(hara)
    haraMean = np.mean(hara, axis=0)
    haraStd = np.std(hara, axis=0)
    haraMins = np.min(hara, axis=0)
    haraMaxs = np.max(hara, axis=0)
    norm = (haraMaxs-haraMins)
    copy = im.copy()
    copy = cv2.cvtColor(copy, cv2.COLOR_BGRA2RGBA)
    im = cv2.cvtColor(im, cv2.COLOR_BGRA2RGBA)

    for i in range(hara.shape[0]):
        brightScale = 255*(hara[i] - haraMins)/norm
        bright = brightScale[0][0]
        r = rects[i]
        cv2.rectangle(copy, r[0], r[1], [0, bright, 0, 255], -1)

    f, axarr = plt.subplots(2, 1)
    axarr[0].imshow(copy)
    axarr[1].imshow(im)
    plt.show()

# load configuration
with open('../../config.json', 'rb') as f:
    config = json.load(f)

mainImPath = config['image_dir']
pDir = config['pickle_dir']

pcaFG = pk.load(open(pDir + 'pcaFG.pk', 'rb'))

bb = pk.load(open(pDir + 'pDogs-bounding-boxes-clean.pd.pk', 'rb'))
bb.dropna(inplace=True)

# do something like sorted(bb.breed.unique().tolist())[50:] to check another breed
for breed in sorted(bb.breed.unique().tolist())[50:]:
    print('breed:', breed)
    cropDir = mainImPath + breed + '/grabcut/'
    fgDir = cropDir + 'fg/'
    fgFiles = os.listdir(fgDir)

    for fi in fgFiles:
        try:
            fg = cv2.imread(fgDir + fi, -1) # -1 tells it to load alpha channel
        except Exception as err:
            print('exception:', err)
            continue
        fgRects, bgRects = get_fg_bg_rects(fg)
        make_hara_map(fg, fgRects)


"""
PDBBind binding pocket dataset loader.
"""

from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

import os
import numpy as np
import pandas as pd
import shutil
import time
import re
from rdkit import Chem
import deepchem as dc

def compute_binding_pocket_features(pocket_featurizer, ligand_featurizer,
                                    pdb_subdir, pdb_code, threshold=.3):
  """Compute features for a given complex"""
  protein_file = os.path.join(pdb_subdir, "%s_protein.pdb" % pdb_code)
  ligand_file = os.path.join(pdb_subdir, "%s_ligand.sdf" % pdb_code)
  ligand_mol2 = os.path.join(pdb_subdir, "%s_ligand.mol2" % pdb_code)

  # Extract active site
  active_site_box, active_site_atoms, active_site_coords = (
      dc.dock.binding_pocket.extract_active_site(
          protein_file, ligand_file))

  # Featurize ligand
  mol = Chem.MolFromMol2File(str(ligand_mol2), removeHs=False)
  if mol is None:
    return None, None
  # Default for CircularFingerprint
  n_ligand_features = 1024
  ligand_features = ligand_featurizer.featurize([mol])

  # Featurize pocket
  finder = dc.dock.ConvexHullPocketFinder()
  pockets, pocket_atoms, pocket_coords = finder.find_pockets(protein_file, ligand_file)
  n_pockets = len(pockets)
  n_pocket_features = dc.feat.BindingPocketFeaturizer.n_features

  features = np.zeros((n_pockets, n_pocket_features+n_ligand_features))
  pocket_features = pocket_featurizer.featurize(
      protein_file, pockets, pocket_atoms, pocket_coords)
  # Note broadcast operation
  features[:, :n_pocket_features] = pocket_features
  features[:, n_pocket_features:] = ligand_features

  # Compute labels for pockets
  labels = np.zeros(n_pockets)
  pocket_atoms[active_site_box] = active_site_atoms
  for ind, pocket in enumerate(pockets):
    overlap = dc.dock.binding_pocket.compute_overlap(
        pocket_atoms, active_site_box, pocket)
    if overlap > threshold:
      labels[ind] = 1
    else:
      labels[ind] = 0
  return features, labels

def load_pdbbind_labels(labels_file):
  """Loads pdbbind labels as dataframe"""
  # Some complexes have labels but no PDB files. Filter these manually
  missing_pdbs = ["1d2v", "1jou", "1s8j", "1cam", "4mlt", "4o7d"]
  contents = []
  with open(labels_file) as f:
    for line in f:
      if line.startswith("#"):
        continue
      else:
        # Some of the ligand-names are of form (FMN ox). Use regex
        # to merge into form (FMN-ox)
        p = re.compile('\(([^\)\s]*) ([^\)\s]*)\)')
        line = p.sub('(\\1-\\2)', line)
        elts = line.split()
        # Filter if missing PDB files
        if elts[0] in missing_pdbs:
          continue
        contents.append(elts)
  contents_df = pd.DataFrame(
      contents,
      columns=("PDB code", "resolution", "release year", "-logKd/Ki", "Kd/Ki",
               "ignore-this-field", "reference", "ligand name"))
  return contents_df

def featurize_pdbbind_pockets(data_dir=None, subset="core"):
  """Featurizes pdbbind according to provided featurization"""
  tasks = ["active-site"]
  current_dir = os.path.dirname(os.path.realpath(__file__))
  data_dir = os.path.join(current_dir, "%s_pockets" % (subset))
  if os.path.exists(data_dir):
    return dc.data.DiskDataset(data_dir), tasks
  pdbbind_dir = os.path.join(current_dir, "../pdbbind/v2015")

  # Load PDBBind dataset
  if subset == "core":
    labels_file = os.path.join(pdbbind_dir, "INDEX_core_data.2013")
  elif subset == "refined":
    labels_file = os.path.join(pdbbind_dir, "INDEX_refined_data.2015")
  elif subset == "full":
    labels_file = os.path.join(pdbbind_dir, "INDEX_general_PL_data.2015")
  else:
    raise ValueError("Only core, refined, and full subsets supported.")
  print("About to load contents.")
  if not os.path.exists(labels_file):
    raise ValueError("Run ../pdbbind/get_pdbbind.sh to download dataset.")
  contents_df = load_pdbbind_labels(labels_file)
  ids = contents_df["PDB code"].values
  y = np.array([float(val) for val in contents_df["-logKd/Ki"].values])

  # Define featurizers
  pocket_featurizer = dc.feat.BindingPocketFeaturizer()
  ligand_featurizer = dc.feat.CircularFingerprint(size=1024)

  # Featurize Dataset
  all_features = []
  all_labels = []
  missing_pdbs = []
  all_ids = []
  time1 = time.time()
  for ind, pdb_code in enumerate(ids):
    print("Processing complex %d, %s" % (ind, str(pdb_code)))
    pdb_subdir = os.path.join(pdbbind_dir, pdb_code)
    if not os.path.exists(pdb_subdir):
      print("%s is missing!" % pdb_subdir)
      missing_pdbs.append(pdb_subdir)
      continue
    features, labels = compute_binding_pocket_features(
        pocket_featurizer, ligand_featurizer, pdb_subdir, pdb_code)
    if features is None:
      print("Featurization failed!")
      continue
    all_features.append(features)
    all_labels.append(labels)
    ids = np.array(["%s%d" % (pdb_code, i) for i in range(len(labels))])
    all_ids.append(ids)
  time2 = time.time()
  print("TIMING: PDBBind Pocket Featurization took %0.3f s" % (time2-time1))
  X = np.vstack(all_features)
  y = np.concatenate(all_labels)
  w = np.ones_like(y)
  ids = np.concatenate(all_ids)

  dataset = dc.data.DiskDataset.from_numpy(X, y, w, ids, data_dir=data_dir)
  return dataset, tasks

def load_pdbbind_pockets(split="index", subset="core"):
  """Load PDBBind datasets. Does not do train/test split"""
  dataset, tasks = featurize_pdbbind_pockets(subset=subset)

  splitters = {'index': dc.splits.IndexSplitter(),
               'random': dc.splits.RandomSplitter()}
  splitter = splitters[split]
  ########################################################### DEBUG
  print("dataset.X.shape")
  print(dataset.X.shape)
  print("dataset.y.shape")
  print(dataset.y.shape)
  print("dataset.w.shape")
  print(dataset.w.shape)
  print("dataset.ids.shape")
  print(dataset.ids.shape)
  ########################################################### DEBUG
  train, valid, test = splitter.train_valid_test_split(dataset)

  transformers = []
  for transformer in transformers:
    train = transformer.transform(train)
  for transformer in transformers:
    valid = transformer.transform(valid)
  for transformer in transformers:
    test = transformer.transform(test)

  return tasks, (train, valid, test), transformers


""" The data is available from Geonet, the official source of New
Zealand earthquake hazard data:

http://wfs.geonet.org.nz/geonet/ows?service=WFS&version=1.0.0&request=GetFeature&typeName=geonet:quake_search_v1&outputFormat=csv

Geonet Data policy
==================

All data and images are made available free of charge through the
GeoNet project to facilitate research into hazards and assessment of
risk. GeoNet is sponsored by the New Zealand Government through its
agencies: Earthquake Commission (EQC), GNS Science and Land
Information New Zealand (LINZ). The use of data or images is subject
to the following conditions:

Users are requested to acknowledge the GeoNet project sponsors as the
source of the data or images. (Suggested text: We acknowledge the New
Zealand GeoNet project and its sponsors EQC, GNS Science and LINZ, for
providing data/images used in this study.)

The GeoNet project sponsors accept no liability for any loss or
damage, direct or indirect, resulting from the use of the data or
images provided.  The GeoNet project sponsors do not make any
representation in respect of the information's accuracy, completeness
or fitness for any particular purpose.

"""

URL = "http://wfs.geonet.org.nz/geonet/ows?service=WFS&version=1.0.0&request=GetFeature&typeName=geonet:quake_search_v1&outputFormat=csv"

from pathlib import Path
import requests
import karmapi
import pandas

def get(path):

    path = Path(path)

    r = requests.get(URL)

    path.write_bytes(r.content)

def datefix(df):

    tt = df.origintime.apply(lambda x: x[:19])

    df.index = pandas.to_datetime(tt)

    return df


#!/usr/bin/env python

import matplotlib as mpl
from matplotlib import pyplot as plt
import argparse

def add_adiabatic_map_to_axis(axis, style, energies, color):
    """ add single set of energies to plot """

    # Energy horizontal decks
    x = style['START']
    for energy in energies:
        axis.plot([x, x+style['WIDTH']], [energy, energy],
            '-%s' % color, linewidth=2)
        x += style['SPACING']

    # Connect steps
    x = style['START']
    for i in range(1, len(energies)):
        x1 = x + style['WIDTH']
        x2 = x + style['SPACING']
        y1 = energies[i-1]
        y2 = energies[i]
        axis.plot([x1, x2], [y1, y2], '-%s' % color)
        x += style['SPACING']

def getargs():
    parser = argparse.ArgumentParser(description="""
        Make plot from user provided energies.
        Can read multiple sets of energies.""",
    formatter_class=argparse.RawTextHelpFormatter
    )

    parser.add_argument('-o', '--output',
        default='PES.svg',
        help='File name of output figure')
    parser.add_argument('--dpi',
        default=300, type=int,
        help='Resolution for bitmaps')
    parser.add_argument('-e', '--energies',
        nargs='+', type=float, action='append',
        help='Energies for any number of stationary points')
    parser.add_argument('-l', '--labels', nargs='+',
        help='Name of stationary points')
    parser.add_argument('-c', '--colors', nargs='+',
        help='Color codes')

    args = parser.parse_args()

    # less colors than PES ? add 'k'
    if args.colors:
        missing_colors = len(args.energies) - len(args.colors)
        missing_colors = (missing_colors > 0) * missing_colors
        args.colors += 'k' * missing_colors

    return args

def makelabels(N):
    """ Make automatic labels: TS1, INT1, TS2, etc.."""

    labels = ['R']
    n_ts = N / 2
    n_i = (N - 2) / 2
    n_i = n_i * (n_i > 0) # becomes zero if negative
    for i in range(n_ts + n_i):
        if i % 2:
            labels.append('INT%d' % (i/2+1))
        else:
            labels.append('TS%d' % (i/2+1))
    if N % 2 and N >= 3:
        labels.append('P')

    return labels

def configure_axis_limits(axis, style, energies):

    # Appearance
    ymin, ymax = float('+inf'), float('-inf')
    maxlen = 0
    for energy_set in energies:
        ymin = min(ymin, min(energy_set))
        ymax = max(ymax, max(energy_set))
        maxlen = max(len(energy_set), maxlen)
    yrange = ymax-ymin
    axis.set_ylim(ymin-0.1*yrange, ymax+0.1*yrange)
    xmax = style['START']*2 + style['WIDTH'] + (maxlen-1)*style['SPACING']
    axis.set_xlim(0, xmax)
    axis.set_xticks([
        style['START']+i*style['SPACING']+style['WIDTH']/2.0 for i in range(maxlen)])

    return maxlen

def main():

    # get user input
    args = getargs()

    # important style features
    style = {
        'WIDTH'     : 4,    # width of horizontal bars
        'SPACING'   : 10,   # spacing between center of horizontal bars
        'START'     : 3     # x-offset from y-axis
    }

    # Configure Figure
    fig = plt.gcf()
    fig.set_size_inches(3.3, 2.5)
    mpl.rcParams.update({'font.size': 7, 'axes.linewidth':0.5})
    plt.subplots_adjust(bottom=.15)
    plt.subplots_adjust(left=.15)
    plt.ylabel('Energy (kcal/mol)')
    plt.xlabel('Reaction coordinate')
    ax = fig.gca()
    ax.grid(True)
    maxlen = configure_axis_limits(ax, style, args.energies)

    if not args.labels:
        args.labels = makelabels(maxlen)

    ax.set_xticklabels(args.labels)

    # plot stuff
    color = 'k'
    for j,energies in enumerate(args.energies):
        if args.colors:
            color = args.colors[j]
        add_adiabatic_map_to_axis(ax, style, energies, color)
    plt.savefig(args.output, dpi=args.dpi)

if __name__ == '__main__':
    main()


# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2013, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program 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 Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------
"""
Provides two classes with the same signature for writing data out of NuPIC
models.
(This is a component of the One Hot Gym Prediction Tutorial.)
"""
import csv
from collections import deque
from abc import ABCMeta, abstractmethod
# Try to import matplotlib, but we don't have to.
try:
  import matplotlib
  matplotlib.use('TKAgg')
  import matplotlib.pyplot as plt
  import matplotlib.gridspec as gridspec
  from matplotlib.dates import date2num
except ImportError:
  pass

WINDOW = 100

class NuPICOutput(object):

  __metaclass__ = ABCMeta

  def __init__(self, names, showAnomalyScore=False):
    self.names = names
    self.showAnomalyScore = showAnomalyScore

  @abstractmethod
  def write(self, timestamps, actualValues, predictedValues,
            predictionStep=1):
    pass

  @abstractmethod
  def close(self):
    pass

class NuPICFileOutput(NuPICOutput):

  def __init__(self, *args, **kwargs):
    super(NuPICFileOutput, self).__init__(*args, **kwargs)
    self.outputFiles = []
    self.outputWriters = []
    self.lineCounts = []
    headerRow = ['timestamp', 'kw_energy_consumption', 'prediction']
    for name in self.names:
      self.lineCounts.append(0)
      outputFileName = "%s_out.csv" % name
      print "Preparing to output %s data to %s" % (name, outputFileName)
      outputFile = open(outputFileName, "w")
      self.outputFiles.append(outputFile)
      outputWriter = csv.writer(outputFile)
      self.outputWriters.append(outputWriter)
      outputWriter.writerow(headerRow)

  def write(self, timestamps, actualValues, predictedValues,
            predictionStep=1):

    assert len(timestamps) == len(actualValues) == len(predictedValues)

    for index in range(len(self.names)):
      timestamp = timestamps[index]
      actual = actualValues[index]
      prediction = predictedValues[index]
      writer = self.outputWriters[index]

      if timestamp is not None:
        outputRow = [timestamp, actual, prediction]
        writer.writerow(outputRow)
        self.lineCounts[index] += 1

  def close(self):
    for index, name in enumerate(self.names):
      self.outputFiles[index].close()
      print "Done. Wrote %i data lines to %s." % (self.lineCounts[index], name)

class NuPICPlotOutput(NuPICOutput):

  def __init__(self, *args, **kwargs):
    super(NuPICPlotOutput, self).__init__(*args, **kwargs)
    # Turn matplotlib interactive mode on.
    plt.ion()
    self.dates = []
    self.convertedDates = []
    self.actualValues = []
    self.predictedValues = []
    self.actualLines = []
    self.predictedLines = []
    self.linesInitialized = False
    self.graphs = []
    plotCount = len(self.names)
    plotHeight = max(plotCount * 3, 6)
    fig = plt.figure(figsize=(14, plotHeight))
    gs = gridspec.GridSpec(plotCount, 1)
    for index in range(len(self.names)):
      self.graphs.append(fig.add_subplot(gs[index, 0]))
      plt.title(self.names[index])
      plt.ylabel('KW Energy Consumption')
      plt.xlabel('Date')
    plt.tight_layout()

  def initializeLines(self, timestamps):
    for index in range(len(self.names)):
      print "initializing %s" % self.names[index]
      # graph = self.graphs[index]
      self.dates.append(deque([timestamps[index]] * WINDOW, maxlen=WINDOW))
      self.convertedDates.append(deque(
        [date2num(date) for date in self.dates[index]], maxlen=WINDOW
      ))
      self.actualValues.append(deque([0.0] * WINDOW, maxlen=WINDOW))
      self.predictedValues.append(deque([0.0] * WINDOW, maxlen=WINDOW))

      actualPlot, = self.graphs[index].plot(
        self.dates[index], self.actualValues[index]
      )
      self.actualLines.append(actualPlot)
      predictedPlot, = self.graphs[index].plot(
        self.dates[index], self.predictedValues[index]
      )
      self.predictedLines.append(predictedPlot)
    self.linesInitialized = True

  def write(self, timestamps, actualValues, predictedValues,
            predictionStep=1):

    assert len(timestamps) == len(actualValues) == len(predictedValues)

    # We need the first timestamp to initialize the lines at the right X value,
    # so do that check first.
    if not self.linesInitialized:
      self.initializeLines(timestamps)

    for index in range(len(self.names)):
      self.dates[index].append(timestamps[index])
      self.convertedDates[index].append(date2num(timestamps[index]))
      self.actualValues[index].append(actualValues[index])
      self.predictedValues[index].append(predictedValues[index])

      # Update data
      self.actualLines[index].set_xdata(self.convertedDates[index])
      self.actualLines[index].set_ydata(self.actualValues[index])
      self.predictedLines[index].set_xdata(self.convertedDates[index])
      self.predictedLines[index].set_ydata(self.predictedValues[index])

      self.graphs[index].relim()
      self.graphs[index].autoscale_view(True, True, True)

    plt.draw()
    plt.legend(('actual','predicted'), loc=3)

  def close(self):
    plt.ioff()
    plt.show()

NuPICOutput.register(NuPICFileOutput)
NuPICOutput.register(NuPICPlotOutput)


#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
Plotter for 2D slices of GEOS-Chem output NetCDFs files.

NOTES
---
 - This is setup for Cly, but many other options (plot/species) are availible
    by just updating passed variables/plotting function called.
"""

import AC_tools as AC
import numpy as np
import matplotlib.pyplot as plt

def main():
    """
    Basic plotter of NetCDF files using AC_tools
    """
    # --- Local settings hardwired here...
    fam = 'Cly'  # Family to plot
    # print species in family for reference...
    print((AC.GC_var(fam)))

    # --- Get working directory etc from command line (as a dictionary object)
    # (1st argument is fil directory with folder, 2nd is filename)
    Var_rc = AC.get_default_variable_dict()
    # Get details on extracted data (inc. resolution)
    Data_rc = AC.get_shared_data_as_dict(Var_rc=Var_rc)

    # --- extract data and units of data for family/species...
    arr, units = AC.fam_data_extractor(wd=Var_rc['wd'], fam=fam,
                                       res=Data_rc['res'], rtn_units=True, annual_mean=False)

    # --- Process data (add and extra processing of data here... )
    # take average over time
    print((arr.shape))
    arr = arr.mean(axis=-1)
    # Select surface values
    print((arr.shape))
    arr = arr[..., 0]
    # convert to pptv
    arr = arr*1E12
    units = 'pptv'

    # --- Plot up data...
    print((arr.shape))
    #  - Plot a (very) simple plot ...
#    AC.map_plot( arr.T, res=Data_rc['res'] )
    #  - plot a slightly better plot...
    # (loads of options here - just type help(AC.plot_spatial_figure) in ipython)
    # set range for data...
    fixcb = np.array([0., 100.])
    # number of ticks on colorbar (make sure the fixcb range divides by this)
    nticks = 6
    interval = (1/3.)  # number of lat/lon labels... (x*15 degrees... )
    # set limits of plot
    lat_min = 5.
    lat_max = 75.
    lon_min = -30.
    lon_max = 60.
    left_cb_pos = 0.85  # set X (fractional) position
    axis_titles = True  # add labels for lat and lon
    # title for plot
    title = "Plot of annual average {}".format(fam)
    #  save as pdf (just set to True) or show?
#    figsize = (7,5) # figsize to use? (e.g. square or rectangular plot)

    # call plotter...
    AC.plot_spatial_figure(arr, res=Data_rc['res'], units=units, fixcb=fixcb,
                           lat_min=lat_min, lat_max=lat_max, lon_min=lon_min, lon_max=lon_max,
                           axis_titles=axis_titles, left_cb_pos=left_cb_pos,
                           nticks=nticks, interval=interval, title=title, show=False)

    # are the spacings right? - if not just up
    bottom = 0.1
    top = 0.9
    left = 0.1
    right = 0.9
    fig = plt.gcf()
    fig.subplots_adjust(bottom=bottom, top=top, left=left, right=right)

    # show and save as PDF?
    plt.savefig('pete_plot.png')
    AC.show_plot()

if __name__ == "__main__":
    main()


import itertools

from common import ImplementsReduce
from ops import inject_reduce_methods
import variable
import dataset
import numpy as np

def unique_value_groups(ar):
    """Group an array by its unique values.

    Parameters
    ----------
    ar : array-like
        Input array. This will be flattened if it is not already 1-D.

    Returns
    -------
    values : np.ndarray
        Sorted, unique values as returned by `np.unique`.
    indices : list of lists of int
        Each element provides the integer indices in `ar` with values given by
        the corresponding value in `unique_values`.
    """
    values, inverse = np.unique(ar, return_inverse=True)
    groups = [[] for _ in range(len(values))]
    for n, g in enumerate(inverse):
        groups[g].append(n)
    return values, groups

def peek_at(iterable):
    """Returns the first value from iterable, as well as a new iterable with
    the same content as the original iterable
    """
    gen = iter(iterable)
    peek = gen.next()
    return peek, itertools.chain([peek], gen)

class GroupBy(object):
    """A object that implements the split-apply-combine pattern.

    Modeled after `pandas.GroupBy`. The `GroupBy` object can be iterated over
    (unique_value, grouped_array) pairs, but the main way to interact with a
    groupby object are with the `apply` or `reduce` methods. You can also
    directly call numpy methods like `mean` or `std`.

    You should create a GroupBy object by using the `DataArray.groupby` or
    `Dataset.groupby` methods.

    See Also
    --------
    XArray.groupby
    DataArray.groupby
    """
    def __init__(self, obj, group_coord, squeeze=True):
        """Create a GroupBy object

        Parameters
        ----------
        obj : Dataset or DataArray
            Object to group.
        group_coord : DataArray
            1-dimensional array with the group values.
        squeeze : boolean, optional
            If "group" is a coordinate of object, `squeeze` controls whether
            the subarrays have a dimension of length 1 along that coordinate or
            if the dimension is squeezed out.
        """
        if group_coord.ndim != 1:
            # TODO: remove this limitation?
            raise ValueError('`group_coord` must be 1 dimensional')

        self.obj = obj
        self.group_coord = group_coord
        self.group_dim, = group_coord.dimensions

        expected_size = dataset.as_dataset(obj).dimensions[self.group_dim]
        if group_coord.size != expected_size:
            raise ValueError('the group variable\'s length does not '
                             'match the length of this variable along its '
                             'dimension')

        if group_coord.name in obj.dimensions:
            # assume that group_coord already has sorted, unique values
            if group_coord.dimensions != (group_coord.name,):
                raise ValueError('`group_coord` is required to be a '
                                 'coordinate variable if `group_coord.name` '
                                 'is a dimension in `obj`')
            group_indices = np.arange(group_coord.size)
            if not squeeze:
                # group_indices = group_indices.reshape(-1, 1)
                # use slices to do views instead of fancy indexing
                group_indices = [slice(i, i + 1) for i in group_indices]
            unique_coord = group_coord
        else:
            # look through group_coord to find the unique values
            unique_values, group_indices = unique_value_groups(group_coord)
            # TODO: switch this to using the new DataArray constructor when we
            # get around to writing it:
            # unique_coord = xary.DataArray(unique_values, name=group_coord.name)
            variables = {group_coord.name: (group_coord.name, unique_values)}
            unique_coord = dataset.Dataset(variables)[group_coord.name]

        self.group_indices = group_indices
        self.unique_coord = unique_coord
        self._groups = None

    @property
    def groups(self):
        # provided to mimic pandas.groupby
        if self._groups is None:
            self._groups = dict(zip(self.unique_coord.values,
                                    self.group_indices))
        return self._groups

    def __len__(self):
        return self.unique_coord.size

    def __iter__(self):
        return itertools.izip(self.unique_coord.values, self._iter_grouped())

    def _iter_grouped(self):
        """Iterate over each element in this group"""
        for indices in self.group_indices:
            yield self.obj.indexed(**{self.group_dim: indices})

    def _infer_concat_args(self, applied_example):
        if self.group_dim in applied_example.dimensions:
            concat_dim = self.group_coord
            indexers = self.group_indices
        else:
            concat_dim = self.unique_coord
            indexers = np.arange(self.unique_coord.size)
        return concat_dim, indexers

    @property
    def _combine(self):
        return type(self.obj).concat

class ArrayGroupBy(GroupBy, ImplementsReduce):
    """GroupBy object specialized to grouping DataArray objects
    """
    def _iter_grouped_shortcut(self):
        """Fast version of `_iter_grouped` that yields XArrays without metadata
        """
        array = variable.as_variable(self.obj)

        # build the new dimensions
        if isinstance(self.group_indices[0], int):
            # group_dim is squeezed out
            dims = tuple(d for d in array.dimensions if d != self.group_dim)
        else:
            dims = array.dimensions

        # slice the data and build the new Arrays directly
        indexer = [slice(None)] * array.ndim
        group_axis = array.get_axis_num(self.group_dim)
        for indices in self.group_indices:
            indexer[group_axis] = indices
            data = array.values[tuple(indexer)]
            yield variable.Variable(dims, data)

    def _combine_shortcut(self, applied, concat_dim, indexers):
        stacked = variable.Variable.concat(
            applied, concat_dim, indexers, shortcut=True)
        stacked.attrs.update(self.obj.attrs)

        name = self.obj.name
        ds = self.obj.dataset.unselect(name)
        ds[concat_dim.name] = concat_dim
        # remove extraneous dimensions
        for dim in self.obj.dimensions:
            if dim not in stacked.dimensions:
                del ds[dim]
        ds[name] = stacked
        return ds[name]

    def _restore_dim_order(self, stacked, concat_dim):
        def lookup_order(dimension):
            if dimension == self.group_coord.name:
                dimension, = concat_dim.dimensions
            if dimension in self.obj.dimensions:
                axis = self.obj.get_axis_num(dimension)
            else:
                axis = 1e6  # some arbitrarily high value
            return axis

        new_order = sorted(stacked.dimensions, key=lookup_order)
        return stacked.transpose(*new_order)

    def apply(self, func, shortcut=False, **kwargs):
        """Apply a function over each array in the group and concatenate them
        together into a new array.

        `func` is called like `func(ar, *args, **kwargs)` for each array `ar`
        in this group.

        Apply uses heuristics (like `pandas.GroupBy.apply`) to figure out how
        to stack together the array. The rule is:
        1. If the dimension along which the group coordinate is defined is
           still in the first grouped array after applying `func`, then stack
           over this dimension.
        2. Otherwise, stack over the new dimension given by name of this
           grouping (the argument to the `groupby` function).

        Parameters
        ----------
        func : function
            Callable to apply to each array.
        shortcut : bool, optional
            Whether or not to shortcut evaluation under the assumptions that:
            (1) The action of `func` does not depend on any of the array
                metadata (attributes, indices or other contained arrays) but
                only on the data and dimensions.
            (2) The action of `func` creates arrays with homogeneous metadata,
                that is, with the same dimensions and attributes.
            If these conditions are satisfied `shortcut` provides significant
            speedup. This should be the case for many common groupby operations
            (e.g., applying numpy ufuncs).
        **kwargs
            Used to call `func(ar, **kwargs)` for each array `ar.

        Returns
        -------
        applied : DataArray
            The result of splitting, applying and combining this array.
        """
        if shortcut:
            grouped = self._iter_grouped_shortcut()
        else:
            grouped = self._iter_grouped()
        applied = (func(arr, **kwargs) for arr in grouped)

        # peek at applied to determine which coordinate to stack over
        applied_example, applied = peek_at(applied)
        concat_dim, indexers = self._infer_concat_args(applied_example)
        if shortcut:
            combined = self._combine_shortcut(applied, concat_dim, indexers)
        else:
            combined = self._combine(applied, concat_dim, indexers)

        reordered = self._restore_dim_order(combined, concat_dim)
        return reordered

    def reduce(self, func, dimension=None, axis=None, shortcut=True,
               **kwargs):
        """Reduce the items in this group by applying `func` along some
        dimension(s).

        Parameters
        ----------
        func : function
            Function which can be called in the form
            `func(x, axis=axis, **kwargs)` to return the result of collapsing an
            np.ndarray over an integer valued axis.
        dimension : str or sequence of str, optional
            Dimension(s) over which to apply `func`.
        axis : int or sequence of int, optional
            Axis(es) over which to apply `func`. Only one of the 'dimension'
            and 'axis' arguments can be supplied. If neither are supplied, then
            `func` is calculated over all dimension for each group item.
        **kwargs : dict
            Additional keyword arguments passed on to `func`.

        Returns
        -------
        reduced : Array
            Array with summarized data and the indicated dimension(s)
            removed.
        """
        def reduce_array(ar):
            return ar.reduce(func, dimension, axis, **kwargs)
        return self.apply(reduce_array, shortcut=shortcut)

    _reduce_method_docstring = \
        """Reduce the items in this group by applying `{name}` along some
        dimension(s).

        Parameters
        ----------
        dimension : str or sequence of str, optional
            Dimension(s) over which to apply `{name}`.
        axis : int or sequence of int, optional
            Axis(es) over which to apply `{name}`. Only one of the 'dimension'
            and 'axis' arguments can be supplied. If neither are supplied, then
            `{name}` is calculated over all dimension for each group item.
        **kwargs : dict
            Additional keyword arguments passed on to `{name}`.

        Returns
        -------
        reduced : {cls}
            New {cls} object with `{name}` applied to its data and the
            indicated dimension(s) removed.
        """

inject_reduce_methods(ArrayGroupBy)

class DatasetGroupBy(GroupBy):
    def apply(self, func, **kwargs):
        """Apply a function over each Dataset in the group and concatenate them
        together into a new Dataset.

        `func` is called like `func(ds, *args, **kwargs)` for each dataset `ds`
        in this group.

        Apply uses heuristics (like `pandas.GroupBy.apply`) to figure out how
        to stack together the datasets. The rule is:
        1. If the dimension along which the group coordinate is defined is
           still in the first grouped item after applying `func`, then stack
           over this dimension.
        2. Otherwise, stack over the new dimension given by name of this
           grouping (the argument to the `groupby` function).

        Parameters
        ----------
        func : function
            Callable to apply to each sub-dataset.
        **kwargs
            Used to call `func(ds, **kwargs)` for each sub-dataset `ar`.

        Returns
        -------
        applied : Dataset
            The result of splitting, applying and combining this dataset.
        """
        applied = [func(ds, **kwargs) for ds in self._iter_grouped()]
        concat_dim, indexers = self._infer_concat_args(applied[0])
        combined = self._combine(applied, concat_dim, indexers)
        return combined


# -*- coding: utf-8 -*-
# vispy: testskip
# -----------------------------------------------------------------------------
# Copyright (c) 2015, Vispy Development Team.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
# -----------------------------------------------------------------------------
"""
Example demonstrating how to use vispy.pyplot, which uses mplexporter
to convert matplotlib commands to vispy draw commands.

Requires matplotlib.
"""

import numpy as np

# You can use either matplotlib or vispy to render this example:
# import matplotlib.pyplot as plt
import vispy.mpl_plot as plt

from vispy.io import read_png, load_data_file

n = 200
freq = 10
fs = 100.
t = np.arange(n) / fs
tone = np.sin(2*np.pi*freq*t)
noise = np.random.RandomState(0).randn(n)
signal = tone + noise
magnitude = np.abs(np.fft.fft(signal))
freqs = np.fft.fftfreq(n, 1. / fs)
flim = n // 2

# Signal
fig = plt.figure()
ax = plt.subplot(311)
ax.imshow(read_png(load_data_file('pyplot/logo.png')))

ax = plt.subplot(312)
ax.plot(t, signal, 'k-')

# Frequency content
ax = plt.subplot(313)
idx = np.argmax(magnitude[:flim])
ax.text(freqs[idx], magnitude[idx], 'Max: %s Hz' % freqs[idx],
        verticalalignment='top')
ax.plot(freqs[:flim], magnitude[:flim], 'k-o')

plt.draw()

# NOTE: show() has currently been overwritten to convert to vispy format, so:
# 1. It must be called to show the results, and
# 2. Any plotting commands executed after this will not take effect.
# We are working to remove this limitation.

if __name__ == '__main__':
    plt.show(True)


"""
Number crunching code for multinomial logit.
``mnl_estimate`` and ``mnl_simulate`` especially are used by
``urbansim.models.lcm``.

"""
from __future__ import print_function

import logging

import numpy as np
import pandas as pd
import scipy.optimize

import pmat
from pmat import PMAT

from ..utils.logutil import log_start_finish

logger = logging.getLogger(__name__)

# right now MNL can only estimate location choice models, where every equation
# is the same
# it might be better to use stats models for a non-location choice problem

# data should be column matrix of dimensions NUMVARS x (NUMALTS*NUMOBVS)
# beta is a row vector of dimensions 1 X NUMVARS

def mnl_probs(data, beta, numalts):
    logging.debug('start: calculate MNL probabilities')
    clamp = data.typ == 'numpy'
    utilities = beta.multiply(data)
    if numalts == 0:
        raise Exception("Number of alternatives is zero")
    utilities.reshape(numalts, utilities.size() / numalts)

    exponentiated_utility = utilities.exp(inplace=True)
    if clamp:
        exponentiated_utility.inftoval(1e20)
    if clamp:
        exponentiated_utility.clamptomin(1e-300)
    sum_exponentiated_utility = exponentiated_utility.sum(axis=0)
    probs = exponentiated_utility.divide_by_row(
        sum_exponentiated_utility, inplace=True)
    if clamp:
        probs.nantoval(1e-300)
    if clamp:
        probs.clamptomin(1e-300)

    logging.debug('finish: calculate MNL probabilities')
    return probs

def get_hessian(derivative):
    return np.linalg.inv(np.dot(derivative, np.transpose(derivative)))

def get_standard_error(hessian):
    return np.sqrt(np.diagonal(hessian))

# data should be column matrix of dimensions NUMVARS x (NUMALTS*NUMOBVS)
# beta is a row vector of dimensions 1 X NUMVARS

def mnl_loglik(beta, data, chosen, numalts, weights=None, lcgrad=False,
               stderr=0):
    logger.debug('start: calculate MNL log-likelihood')
    numvars = beta.size
    numobs = data.size() / numvars / numalts

    beta = np.reshape(beta, (1, beta.size))
    beta = PMAT(beta, data.typ)

    probs = mnl_probs(data, beta, numalts)

    # lcgrad is the special gradient for the latent class membership model
    if lcgrad:
        assert weights
        gradmat = weights.subtract(probs).reshape(probs.size(), 1)
        gradarr = data.multiply(gradmat)
    else:
        if not weights:
            gradmat = chosen.subtract(probs).reshape(probs.size(), 1)
        else:
            gradmat = chosen.subtract(probs).multiply_by_row(
                weights).reshape(probs.size(), 1)
        gradarr = data.multiply(gradmat)

    if stderr:
        gradmat = data.multiply_by_row(gradmat.reshape(1, gradmat.size()))
        gradmat.reshape(numvars, numalts * numobs)
        return get_standard_error(get_hessian(gradmat.get_mat()))

    chosen.reshape(numalts, numobs)
    if weights is not None:
        if probs.shape() == weights.shape():
            loglik = ((probs.log(inplace=True)
                       .element_multiply(weights, inplace=True)
                       .element_multiply(chosen, inplace=True))
                      .sum(axis=1).sum(axis=0))
        else:
            loglik = ((probs.log(inplace=True)
                       .multiply_by_row(weights, inplace=True)
                       .element_multiply(chosen, inplace=True))
                      .sum(axis=1).sum(axis=0))
    else:
        loglik = (probs.log(inplace=True).element_multiply(
            chosen, inplace=True)).sum(axis=1).sum(axis=0)

    if loglik.typ == 'numpy':
        loglik, gradarr = loglik.get_mat(), gradarr.get_mat().flatten()
    else:
        loglik = loglik.get_mat()[0, 0]
        gradarr = np.reshape(gradarr.get_mat(), (1, gradarr.size()))[0]

    logger.debug('finish: calculate MNL log-likelihood')
    return -1 * loglik, -1 * gradarr

def mnl_simulate(data, coeff, numalts, GPU=False, returnprobs=True):
    """
    Get the probabilities for each chooser choosing between `numalts`
    alternatives.

    Parameters
    ----------
    data : 2D array
        The data are expected to be in "long" form where each row is for
        one alternative. Alternatives are in groups of `numalts` rows per
        choosers. Alternatives must be in the same order for each chooser.
    coeff : 1D array
        The model coefficients corresponding to each column in `data`.
    numalts : int
        The number of alternatives available to each chooser.
    GPU : bool, optional
    returnprobs : bool, optional
        If True, return the probabilities for each chooser/alternative instead
        of actual choices.

    Returns
    -------
    probs or choices: 2D array
        If `returnprobs` is True the probabilities are a 2D array with a
        row for each chooser and columns for each alternative.

    """
    logger.debug(
        'start: MNL simulation with len(data)={} and numalts={}'.format(
            len(data), numalts))
    atype = 'numpy' if not GPU else 'cuda'

    data = np.transpose(data)
    coeff = np.reshape(np.array(coeff), (1, len(coeff)))

    data, coeff = PMAT(data, atype), PMAT(coeff, atype)

    probs = mnl_probs(data, coeff, numalts)

    if returnprobs:
        return np.transpose(probs.get_mat())

    # convert to cpu from here on - gpu doesn't currently support these ops
    if probs.typ == 'cuda':
        probs = PMAT(probs.get_mat())

    probs = probs.cumsum(axis=0)
    r = pmat.random(probs.size() / numalts)
    choices = probs.subtract(r, inplace=True).firstpositive(axis=0)

    logger.debug('finish: MNL simulation')
    return choices.get_mat()

def mnl_estimate(data, chosen, numalts, GPU=False, coeffrange=(-3, 3),
                 weights=None, lcgrad=False, beta=None):
    """
    Calculate coefficients of the MNL model.

    Parameters
    ----------
    data : 2D array
        The data are expected to be in "long" form where each row is for
        one alternative. Alternatives are in groups of `numalts` rows per
        choosers. Alternatives must be in the same order for each chooser.
    chosen : 2D array
        This boolean array has a row for each chooser and a column for each
        alternative. The column ordering for alternatives is expected to be
        the same as their row ordering in the `data` array.
        A one (True) indicates which alternative each chooser has chosen.
    numalts : int
        The number of alternatives.
    GPU : bool, optional
    coeffrange : tuple of floats, optional
        Limits of (min, max) to which coefficients are clipped.
    weights : ndarray, optional
    lcgrad : bool, optional
    beta : 1D array, optional
        Any initial guess for the coefficients.

    Returns
    -------
    log_likelihood : dict
        Dictionary of log-likelihood values describing the quality of
        the model fit.
    fit_parameters : pandas.DataFrame
        Table of fit parameters with columns 'Coefficient', 'Std. Error',
        'T-Score'. Each row corresponds to a column in `data` and are given
        in the same order as in `data`.

    See Also
    --------
    scipy.optimize.fmin_l_bfgs_b : The optimization routine used.

    """
    logger.debug(
        'start: MNL fit with len(data)={} and numalts={}'.format(
            len(data), numalts))
    atype = 'numpy' if not GPU else 'cuda'

    numvars = data.shape[1]
    numobs = data.shape[0] / numalts

    if chosen is None:
        chosen = np.ones((numobs, numalts))  # used for latent classes

    data = np.transpose(data)
    chosen = np.transpose(chosen)

    data, chosen = PMAT(data, atype), PMAT(chosen, atype)
    if weights is not None:
        weights = PMAT(np.transpose(weights), atype)

    if beta is None:
        beta = np.zeros(numvars)
    bounds = [coeffrange] * numvars

    with log_start_finish('scipy optimization for MNL fit', logger):
        args = (data, chosen, numalts, weights, lcgrad)
        bfgs_result = scipy.optimize.fmin_l_bfgs_b(mnl_loglik,
                                                   beta,
                                                   args=args,
                                                   fprime=None,
                                                   factr=10,
                                                   approx_grad=False,
                                                   bounds=bounds
                                                   )
    beta = bfgs_result[0]
    stderr = mnl_loglik(
        beta, data, chosen, numalts, weights, stderr=1, lcgrad=lcgrad)

    l0beta = np.zeros(numvars)
    l0 = -1 * mnl_loglik(l0beta, *args)[0]
    l1 = -1 * mnl_loglik(beta, *args)[0]

    log_likelihood = {
        'null': float(l0[0][0]),
        'convergence': float(l1[0][0]),
        'ratio': float((1 - (l1 / l0))[0][0])
    }

    fit_parameters = pd.DataFrame({
        'Coefficient': beta,
        'Std. Error': stderr,
        'T-Score': beta / stderr})

    logger.debug('finish: MNL fit')
    return log_likelihood, fit_parameters


# -*- coding: utf-8 -*-
"""
==================================
Regularized OT with generic solver
==================================

Illustrates the use of the generic solver for regularized OT with
user-designed regularization term. It uses Conditional gradient as in [6] and
generalized Conditional Gradient as proposed in [5][7].

[5] N. Courty; R. Flamary; D. Tuia; A. Rakotomamonjy, Optimal Transport for
Domain Adaptation, in IEEE Transactions on Pattern Analysis and Machine
Intelligence , vol.PP, no.99, pp.1-1.

[6] Ferradans, S., Papadakis, N., Peyré, G., & Aujol, J. F. (2014).
Regularized discrete optimal transport. SIAM Journal on Imaging Sciences,
7(3), 1853-1882.

[7] Rakotomamonjy, A., Flamary, R., & Courty, N. (2015). Generalized
conditional gradient: analysis of convergence and applications.
arXiv preprint arXiv:1510.06567.

"""

import numpy as np
import matplotlib.pylab as pl
import ot
import ot.plot

##############################################################################
# Generate data
# -------------

#%% parameters

n = 100  # nb bins

# bin positions
x = np.arange(n, dtype=np.float64)

# Gaussian distributions
a = ot.datasets.get_1D_gauss(n, m=20, s=5)  # m= mean, s= std
b = ot.datasets.get_1D_gauss(n, m=60, s=10)

# loss matrix
M = ot.dist(x.reshape((n, 1)), x.reshape((n, 1)))
M /= M.max()

##############################################################################
# Solve EMD
# ---------

#%% EMD

G0 = ot.emd(a, b, M)

pl.figure(3, figsize=(5, 5))
ot.plot.plot1D_mat(a, b, G0, 'OT matrix G0')

##############################################################################
# Solve EMD with Frobenius norm regularization
# --------------------------------------------

#%% Example with Frobenius norm regularization

def f(G):
    return 0.5 * np.sum(G**2)

def df(G):
    return G

reg = 1e-1

Gl2 = ot.optim.cg(a, b, M, reg, f, df, verbose=True)

pl.figure(3)
ot.plot.plot1D_mat(a, b, Gl2, 'OT matrix Frob. reg')

##############################################################################
# Solve EMD with entropic regularization
# --------------------------------------

#%% Example with entropic regularization

def f(G):
    return np.sum(G * np.log(G))

def df(G):
    return np.log(G) + 1.

reg = 1e-3

Ge = ot.optim.cg(a, b, M, reg, f, df, verbose=True)

pl.figure(4, figsize=(5, 5))
ot.plot.plot1D_mat(a, b, Ge, 'OT matrix Entrop. reg')

##############################################################################
# Solve EMD with Frobenius norm + entropic regularization
# -------------------------------------------------------

#%% Example with Frobenius norm + entropic regularization with gcg

def f(G):
    return 0.5 * np.sum(G**2)

def df(G):
    return G

reg1 = 1e-3
reg2 = 1e-1

Gel2 = ot.optim.gcg(a, b, M, reg1, reg2, f, df, verbose=True)

pl.figure(5, figsize=(5, 5))
ot.plot.plot1D_mat(a, b, Gel2, 'OT entropic + matrix Frob. reg')
pl.show()


def bivariate_normal(X, Y, sigmax=1.0, sigmay=1.0, mux=0.0, muy=0.0, sigmaxy=0.0):
    Xmu = X-mux
    Ymu = Y-muy
    rho = sigmaxy/(sigmax*sigmay)
    z = Xmu**2/sigmax**2 + Ymu**2/sigmay**2 - 2*rho*Xmu*Ymu/(sigmax*sigmay)
    denom = 2*np.pi*sigmax*sigmay*np.sqrt(1-rho**2)
    return np.exp(-z/(2*(1-rho**2))) / denom

def MolToMPL(mol,size=(300,300),kekulize=True, wedgeBonds=True, imageType=None, fitImage=False, options=None, **kwargs):
  if not mol:
    raise ValueError('Null molecule provided')
  from rdkit.Chem.Draw.mplCanvas import Canvas
  canvas = Canvas(size)
  if options is None:
    options = DrawingOptions()
    options.bgColor=None
  if fitImage:
      drawingOptions.dotsPerAngstrom = int(min(size) / 10)
  options.wedgeDashedBonds=wedgeBonds
  drawer = MolDrawing(canvas=canvas, drawingOptions=options)
  omol=mol
  if kekulize:
    from rdkit import Chem
    mol = Chem.Mol(mol.ToBinary())
    Chem.Kekulize(mol)

  if not mol.GetNumConformers():
    from rdkit.Chem import AllChem
    AllChem.Compute2DCoords(mol)

  drawer.AddMol(mol,**kwargs)
  omol._atomPs=drawer.atomPs[mol]
  for k,v in iteritems(omol._atomPs):
    omol._atomPs[k]=canvas.rescalePt(v)
  canvas._figure.set_size_inches(float(size[0])/100,float(size[1])/100)
  return canvas._figure

def calcAtomGaussians(mol,a=0.03,step=0.02,weights=None):
  import numpy
  from matplotlib import mlab
  x = numpy.arange(0,1,step)
  y = numpy.arange(0,1,step)
  X,Y = numpy.meshgrid(x,y)
  if weights is None:
    weights=[1.]*mol.GetNumAtoms()
  Z = mlab.bivariate_normal(X,Y,a,a,mol._atomPs[0][0], mol._atomPs[0][1])*weights[0] # this is not bivariate case ... only univariate no mixtures #matplotlib.mlab.bivariate_normal(X, Y, sigmax=1.0, sigmay=1.0, mux=0.0, muy=0.0, sigmaxy=0.0)
  for i in range(1,mol.GetNumAtoms()):
    Zp = mlab.bivariate_normal(X,Y,a,a,mol._atomPs[i][0], mol._atomPs[i][1])
    Z += Zp*weights[i]
  return X,Y,Z

def GetSimilarityMapFromWeights(mol, weights, colorMap=cm.PiYG, scale=-1, size=(250, 250), sigma=None,  #@UndefinedVariable  #pylint: disable=E1101
                                coordScale=1.5, step=0.01, colors='k', contourLines=10, alpha=0.5, **kwargs):
  if mol.GetNumAtoms() < 2: raise ValueError("too few atoms")
  fig = Draw.MolToMPL(mol, coordScale=coordScale, size=size, **kwargs)
  if sigma is None:
    if mol.GetNumBonds() > 0:
      bond = mol.GetBondWithIdx(0)
      idx1 = bond.GetBeginAtomIdx()
      idx2 = bond.GetEndAtomIdx()
      sigma = 0.3 * math.sqrt(sum([(mol._atomPs[idx1][i]-mol._atomPs[idx2][i])**2 for i in range(2)]))
    else:
      sigma = 0.3 * math.sqrt(sum([(mol._atomPs[0][i]-mol._atomPs[1][i])**2 for i in range(2)]))
    sigma = round(sigma, 2)
  x, y, z = Draw.calcAtomGaussians(mol, sigma, weights=weights, step=step)
  # scaling
  if scale <= 0.0: maxScale = max(math.fabs(numpy.min(z)), math.fabs(numpy.max(z)))
  else: maxScale = scale
  # coloring
  fig.axes[0].imshow(z, cmap=colorMap, interpolation='bilinear', origin='lower', extent=(0,1,0,1), vmin=-maxScale, vmax=maxScale)
  # contour lines
  # only draw them when at least one weight is not zero
  if len([w for w in weights if w != 0.0]):
      fig.axes[0].contour(x, y, z, contourLines, colors=colors, alpha=alpha, **kwargs)
  return fig


# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#         Gael Varoquaux <gael.varoquaux@normalesup.org>
#         Virgile Fritsch <virgile.fritsch@inria.fr>
#
# License: BSD 3 clause

import numpy as np

from sklearn.utils.testing import assert_almost_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_raises
from sklearn.utils.validation import NotFittedError

from sklearn import datasets
from sklearn.covariance import empirical_covariance, MinCovDet, \
    EllipticEnvelope

X = datasets.load_iris().data
X_1d = X[:, 0]
n_samples, n_features = X.shape

def test_mcd():
    # Tests the FastMCD algorithm implementation
    # Small data set
    # test without outliers (random independent normal data)
    launch_mcd_on_dataset(100, 5, 0, 0.01, 0.1, 80)
    # test with a contaminated data set (medium contamination)
    launch_mcd_on_dataset(100, 5, 20, 0.01, 0.01, 70)
    # test with a contaminated data set (strong contamination)
    launch_mcd_on_dataset(100, 5, 40, 0.1, 0.1, 50)

    # Medium data set
    launch_mcd_on_dataset(1000, 5, 450, 0.1, 0.1, 540)

    # Large data set
    launch_mcd_on_dataset(1700, 5, 800, 0.1, 0.1, 870)

    # 1D data set
    launch_mcd_on_dataset(500, 1, 100, 0.001, 0.001, 350)

def launch_mcd_on_dataset(n_samples, n_features, n_outliers, tol_loc, tol_cov,
                          tol_support):

    rand_gen = np.random.RandomState(0)
    data = rand_gen.randn(n_samples, n_features)
    # add some outliers
    outliers_index = rand_gen.permutation(n_samples)[:n_outliers]
    outliers_offset = 10. * \
        (rand_gen.randint(2, size=(n_outliers, n_features)) - 0.5)
    data[outliers_index] += outliers_offset
    inliers_mask = np.ones(n_samples).astype(bool)
    inliers_mask[outliers_index] = False

    pure_data = data[inliers_mask]
    # compute MCD by fitting an object
    mcd_fit = MinCovDet(random_state=rand_gen).fit(data)
    T = mcd_fit.location_
    S = mcd_fit.covariance_
    H = mcd_fit.support_
    # compare with the estimates learnt from the inliers
    error_location = np.mean((pure_data.mean(0) - T) ** 2)
    assert(error_location < tol_loc)
    error_cov = np.mean((empirical_covariance(pure_data) - S) ** 2)
    assert(error_cov < tol_cov)
    assert(np.sum(H) >= tol_support)
    assert_array_almost_equal(mcd_fit.mahalanobis(data), mcd_fit.dist_)

def test_mcd_issue1127():
    # Check that the code does not break with X.shape = (3, 1)
    # (i.e. n_support = n_samples)
    rnd = np.random.RandomState(0)
    X = rnd.normal(size=(3, 1))
    mcd = MinCovDet()
    mcd.fit(X)

def test_outlier_detection():
    rnd = np.random.RandomState(0)
    X = rnd.randn(100, 10)
    clf = EllipticEnvelope(contamination=0.1)
    assert_raises(NotFittedError, clf.predict, X)
    assert_raises(NotFittedError, clf.decision_function, X)
    clf.fit(X)
    y_pred = clf.predict(X)
    decision = clf.decision_function(X, raw_values=True)
    decision_transformed = clf.decision_function(X, raw_values=False)

    assert_array_almost_equal(
        decision, clf.mahalanobis(X))
    assert_array_almost_equal(clf.mahalanobis(X), clf.dist_)
    assert_almost_equal(clf.score(X, np.ones(100)),
                        (100 - y_pred[y_pred == -1].size) / 100.)
    assert(sum(y_pred == -1) == sum(decision_transformed < 0))