from __future__ import division, print_function, absolute_import

import math
import numpy as np

from abc import ABCMeta, abstractmethod

from sklearn.neighbors import NearestNeighbors
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.neighbors.dist_metrics import METRIC_MAPPING

class Stat(object):
    """The similarity statistics container.

    The similarity statistics is a container to pass the
    calculated measure of similarity between the case
    identified by the case id and the query case between
    functions.

    Parameters
    ----------
    case_id : int
        The case's id.
    similarity : float
        The similarity measure.

    """
    __slots__ = ('_case_id', '_similarity')

    @property
    def case_id(self):
        """The case's id.

        Returns
        -------
        int :
            The case's id

        """
        return self._case_id

    @property
    def similarity(self):
        """The similarity measure.

        Returns
        -------
        float :
            The similarity measure.

        """
        return self._similarity

    def __init__(self, case_id, similarity=None):
        self._case_id = case_id
        self._similarity = similarity

class SimilarityFactory(object):
    """The similarity factory.

    An instance of a similarity model can be created by passing
    the similarity model type.

    Examples
    --------
    >>> from mlpy.knowledgerep.cbr.similarity import SimilarityFactory
    >>> SimilarityFactory.create('float', **{})

    """
    @staticmethod
    def create(_type, **kwargs):
        """
        Create a feature of the given type.

        Parameters
        ----------
        _type : str
            The feature type. Valid feature types are:

                knn
                    A k-nearest-neighbor algorithm is used to determine similarity
                    between cases (:class:`NeighborSimilarity`). The value
                    ``n_neighbors`` must be specified.

                radius-n
                    Similarity between cases is determined by the nearest neighbors
                    within a radius (:class:`NeighborSimilarity`). The value ``radius``
                    must be specified.

                kmeans
                    Similarity is determined by a KMeans clustering algorithm
                    (:class:`KMeansSimilarity`). The value ``n_clusters`` must be specified.

                exact-match
                    Only exact matches are considered similar (:class:`ExactMatchSimilarity`).

                cosine
                    A cosine similarity measure is used to determine similarity between
                    cases (:class:`CosineSimilarity`).

        kwargs : dict, optional
            Non-positional arguments to pass to the class of the given type
            for initialization.

        Returns
        -------
        ISimilarity :
            A similarity instance of the given type.

        """
        try:
            if _type == "knn":
                kwargs["n_neighbors"] = kwargs["method_params"]
            elif _type == "radius-n":
                kwargs["radius"] = kwargs["method_params"]
            elif _type == "kmeans":
                kwargs["n_cluster"] = kwargs["method_params"]
            elif _type == "cosine":
                kwargs["threshold"] = kwargs["method_params"]
            del kwargs["method_params"]

            return {
                "knn": NeighborSimilarity,
                "radius-n": NeighborSimilarity,
                "kmeans": KMeansSimilarity,
                "exact-match": ExactMatchSimilarity,
                "cosine": CosineSimilarity,
            }[_type](**kwargs)

        except KeyError:
            return None

class ISimilarity(object):
    """The similarity model interface.

    The similarity model keeps an internal indexing structure of
    the relevant case data to efficiently computing the similarity
    measure between data points.

    Notes
    -----
    All similarity models must inherit from this class.

    """
    __metaclass__ = ABCMeta

    def __init__(self):
        #: The indexing structure
        self._indexing_structure = None
        #: The mapping of the data points to their case ids
        self._id_map = None
        """:ivar: dict"""

    @abstractmethod
    def build_indexing_structure(self, data, id_map):
        """Build the indexing structure.

        Parameters
        ----------
        data : ndarray[ndarray[float]]
            The raw data points to be indexed.
        id_map : dict[int, int]
            The mapping from the data points to their case ids.

        Raises
        ------
        NotImplementedError
            If the child class does not implement this function.

        """
        raise NotImplementedError

    @abstractmethod
    def compute_similarity(self, data_point):
        """Computes the similarity.

        Computes the similarity between the data point and the data in
        the indexing structure returning the results in a collection of
        similarity statistics (:class:`Stat`).

        Parameters
        ----------
        data_point : list[float]
            The raw data point to compare against the data points stored in the
            indexing structure.

        Returns
        -------
        list[Stat] :
            A collection of similarity statistics.

        Raises
        ------
        NotImplementedError
            If the child class does not implement this function.

        """
        raise NotImplementedError

class NeighborSimilarity(ISimilarity):
    """The neighborhood similarity model.

    The neighbor similarity model determines similarity between the data
    in the indexing structure and the query data by using the nearest
    neighbor algorithm :class:`sklearn.neighbors.NearestNeighbors`.

    Both a k-neighbors classifier and a radius-neighbor-classifier are implemented.
    To choose between the classifiers either `n_neighbors` or `radius` must be
    specified.

    Parameters
    ----------
    n_neighbors : int
        The number of data points considered to be closest neighbors.
    radius : int
        The radius around the query data point, within which the data points
        are considered closest neighbors.
    algorithm : str
        The internal indexing structure of the training data. Defaults to
        `kd-tree`.
    metric : str
        The metric used to compute the distances between pairs of points.
        Refer to :class:`sklearn.neighbors.DistanceMetric` for valid
        identifiers. Default is `euclidean`.
    metric_params : dict
        Parameters relevant to the specified metric.

    Raises
    ------
    UserWarning :
        If the either both or none of `n_neighbors` and `radius` are given.

    See Also
    --------
    :class:`sklearn.neighbors.KNeighborsClassifier`, :class:`sklearn.neighbors.RadiusNeighborsClassifier`

    """
    def __init__(self, n_neighbors=None, radius=None, algorithm=None, metric=None, metric_params=None):
        super(NeighborSimilarity, self).__init__()

        if (n_neighbors is not None and radius is not None) or not (n_neighbors is None or radius is None):
            raise UserWarning("Exactly one of n_neighbors or radius must be initialized.")

        self._n_neighbors = n_neighbors
        self._radius = radius

        if algorithm is not None:
            if algorithm not in ["ball_tree", "kd_tree", "brute", "auto"]:
                raise ValueError("%s is not a valid retrieval algorithm" % algorithm)
            self._algorithm = algorithm
        else:
            self._algorithm = "kd_tree"

        if metric is not None:
            if metric not in METRIC_MAPPING:
                raise ValueError("%s is not a valid retrieval metric" % metric)
            self._metric = metric
        else:
            self._metric = "euclidean"

        self._metric_params = metric_params if metric_params is not None else 2

    def build_indexing_structure(self, data, id_map):
        """Build the indexing structure.

        Build the indexing structure by fitting the data according to the
        specified algorithm.

        Parameters
        ----------
        data : ndarray[ndarray[float]]
            The raw data points to be indexed.
        id_map : dict[int, int]
            The mapping from the data points to their case ids.
        """
        self._id_map = id_map

        if self._n_neighbors is not None:
            self._indexing_structure = NearestNeighbors(n_neighbors=self._n_neighbors, algorithm=self._algorithm,
                                                        metric=self._metric, p=self._metric_params).fit(data)
        else:
            self._indexing_structure = NearestNeighbors(radius=self._radius, algorithm=self._algorithm,
                                                        metric=self._metric, p=self._metric_params).fit(data)

    def compute_similarity(self, data_point):
        """Computes the similarity.

        Computes the similarity between the data point and the data in
        the indexing structure using the :class:`sklearn.neighbors.NearestNeighbors`
        algorithm. The results are returned in a collection of similarity statistics
        (:class:`Stat`).

        Parameters
        ----------
        data_point : list[float]
            The raw data point to compare against the data points stored in the
            indexing structure.

        Returns
        -------
        list[Stat] :
            A collection of similarity statistics.

        """
        if self._n_neighbors is not None:
            # noinspection PyProtectedMember
            raw_data = self._indexing_structure._fit_X
            if len(raw_data) < self._n_neighbors:
                result = []
                for i, feat in enumerate(raw_data):
                    dist = np.linalg.norm(np.asarray(data_point) - np.asarray(feat))
                    result.append(Stat(self._id_map[i], dist))

                # noinspection PyShadowingNames
                result = sorted(result, key=lambda x: x.similarity)
            else:
                d, key_lists = self._indexing_structure.kneighbors(data_point)
                result = [Stat(self._id_map[x], d[0][i]) for i, x in enumerate(key_lists[0])]

        else:
            d, key_lists = self._indexing_structure.radius_neighbors(data_point)
            result = [Stat(self._id_map[x], d[0][i]) for i, x in enumerate(key_lists[0])]
        return result

class KMeansSimilarity(ISimilarity):
    """The KMeans similarity model.

    The KMeans similarity model determines similarity between the data in the
    indexing structure and the query data by using the :class:`sklearn.cluster.KMeans`
    algorithm.

    Parameters
    ----------
    n_cluster : int
        The number of clusters to fit the raw data in.

    """
    def __init__(self, n_cluster=None):
        super(KMeansSimilarity, self).__init__()

        self._n_cluster = n_cluster if n_cluster is None else 8

    def build_indexing_structure(self, data, id_map):
        """Build the indexing structure.

        Build the indexing structure by fitting the data into `n_cluster`
        clusters.

        Parameters
        ----------
        data : ndarray[ndarray[float]]
            The raw data points to be indexed.
        id_map : dict[int, int]
            The mapping from the data points to their case ids.
        """
        self._id_map = id_map
        self._indexing_structure = KMeans(init='k-means++', n_clusters=self._n_cluster, n_init=10).fit(data)

    def compute_similarity(self, data_point):
        """Computes the similarity.

        Computes the similarity between the data point and the data in
        the indexing structure using the :class:`sklearn.cluster.KMeans`
        clustering algorithm. The results are returned in a collection
        of similarity statistics (:class:`Stat`).

        Parameters
        ----------
        data_point : list[float]
            The raw data point to compare against the data points stored in the
            indexing structure.

        Returns
        -------
        list[Stat] :
            A collection of similarity statistics.

        """
        label = self._indexing_structure.predict(data_point)

        result = []
        try:
            # noinspection PyTypeChecker,PyUnresolvedReferences
            key_lists = np.nonzero(self._indexing_structure.labels_ == label[0])[0]
            result = [Stat(self._id_map[x]) for x in key_lists]
        except IndexError:
            pass

        return result

class ExactMatchSimilarity(ISimilarity):
    """The exact match similarity model.

    The exact match similarity model considered only exact matches between
    the data in the indexing structure and the query data as similar.

    """
    # noinspection PyUnusedLocal
    def __init__(self, **kwargs):
        super(ExactMatchSimilarity, self).__init__()

    def build_indexing_structure(self, data, id_map):
        """Build the indexing structure.

        To determine exact matches a brute-force algorithm is used thus
        the data remains as is and no special indexing structure is
        implemented.

        Parameters
        ----------
        data : ndarray[ndarray[float]]
            The raw data points to be indexed.
        id_map : dict[int, int]
            The mapping from the data points to their case ids.

        .. todo::
            It might be worth looking into a more efficient way of determining
            exact matches.

        """
        self._id_map = id_map
        self._indexing_structure = data

    def compute_similarity(self, data_point):
        """Computes the similarity.

        Computes the similarity between the data point and the data in
        the indexing structure identifying exact matches. The results are
        returned in a collection of similarity statistics (:class:`Stat`).

        Parameters
        ----------
        data_point : list[float]
            The raw data point to compare against the data points stored in the
            indexing structure.

        Returns
        -------
        list[Stat] :
            A collection of similarity statistics.

        """
        result = []

        for i, feat in enumerate(self._indexing_structure):
            total = 0
            for j, val in enumerate(data_point):
                total += math.pow(val - feat[j], 2)

            if total == 0.0:
                result.append(Stat(self._id_map[i]))

        return result

class CosineSimilarity(ISimilarity):
    """The cosine similarity model.

    Cosine similarity is a measure of similarity between two vectors of an inner
    product space that measures the cosine of the angle between them. The cosine
    of 0 degree is 1, and it is less than 1 for any other angle. It is thus a
    judgement of orientation and not magnitude: tow vectors with the same
    orientation have a cosine similarity of 1, two vectors at 90 degrees have a
    similarity of 0, and two vectors diametrically opposed have a similarity of -1,
    independent of their magnitude [1]_.

    The cosine model employs the
    `cosine_similarity <http://scikit-learn.org/stable/modules/metrics.html#cosine-similarity>`_
    function from the :mod:`sklearn.metrics.pairwise` module to determine similarity.

    .. seealso::
        `Machine Learning::Cosine Similarity for Vector Space Models (Part III)
        <http://blog.christianperone.com/?p=2497>`_

    References
    ----------
    .. [1] `Wikipidia::cosine_similarity <https://en.wikipedia.org/wiki/Cosine_similarity>`_

    """
    # noinspection PyUnusedLocal
    def __init__(self, **kwargs):
        super(CosineSimilarity, self).__init__()

    def build_indexing_structure(self, data, id_map):
        """Build the indexing structure.

        The cosine_similarity function from :mod:`sklearn.metrics.pairwise` takes
        the raw data as input. Thus the data remains as is and no special indexing
        structure is implemented.

        Parameters
        ----------
        data : ndarray[ndarray[float]]
            The raw data points to be indexed.
        id_map : dict[int, int]
            The mapping from the data points to their case ids.

        """
        self._id_map = id_map
        self._indexing_structure = data

    def compute_similarity(self, data_point):
        """Computes the similarity.

        Computes the similarity between the data point and the data in
        the indexing structure using the function :func:`cosine_similarity` from
        :mod:`sklearn.metrics.pairwise`.

        The resulting similarity ranges from -1 meaning exactly opposite, to 1
        meaning exactly the same, with 0 indicating orthogonality (decorrelation),
        and in-between values indicating intermediate similarity or dissimilarity.
        The results are returned in a collection of similarity statistics (:class:`Stat`).

        Parameters
        ----------
        data_point : list[float]
            The raw data point to compare against the data points stored in the
            indexing structure.

        Returns
        -------
        list[Stat] :
            A collection of similarity statistics.

        """
        similarity = cosine_similarity(data_point, self._indexing_structure)
        if not np.any(data_point):
            similarity = np.array([[float(np.array_equal(data_point, m)) for m in np.array(self._indexing_structure)]])

        return [Stat(self._id_map[i], x) for i, x in enumerate(similarity[0])]


"""
Common code for all metrics

"""
# Authors: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#          Mathieu Blondel <mathieu@mblondel.org>
#          Olivier Grisel <olivier.grisel@ensta.org>
#          Arnaud Joly <a.joly@ulg.ac.be>
#          Jochen Wersdorfer <jochen@wersdoerfer.de>
#          Lars Buitinck
#          Joel Nothman <joel.nothman@gmail.com>
#          Noel Dawe <noel@dawe.me>
# License: BSD 3 clause

from __future__ import division

import numpy as np

from ..utils import check_array, check_consistent_length
from ..utils.multiclass import type_of_target

from ..exceptions import UndefinedMetricWarning as _UndefinedMetricWarning
from ..utils import deprecated

@deprecated("UndefinedMetricWarning has been moved into the sklearn.exceptions"
            " module. It will not be available here from version 0.19")
class UndefinedMetricWarning(_UndefinedMetricWarning):
    pass

def _average_binary_score(binary_metric, y_true, y_score, average,
                          sample_weight=None):
    """Average a binary metric for multilabel classification

    Parameters
    ----------
    y_true : array, shape = [n_samples] or [n_samples, n_classes]
        True binary labels in binary label indicators.

    y_score : array, shape = [n_samples] or [n_samples, n_classes]
        Target scores, can either be probability estimates of the positive
        class, confidence values, or binary decisions.

    average : string, [None, 'micro', 'macro' (default), 'samples', 'weighted']
        If ``None``, the scores for each class are returned. Otherwise,
        this determines the type of averaging performed on the data:

        ``'micro'``:
            Calculate metrics globally by considering each element of the label
            indicator matrix as a label.
        ``'macro'``:
            Calculate metrics for each label, and find their unweighted
            mean.  This does not take label imbalance into account.
        ``'weighted'``:
            Calculate metrics for each label, and find their average, weighted
            by support (the number of true instances for each label).
        ``'samples'``:
            Calculate metrics for each instance, and find their average.

    sample_weight : array-like of shape = [n_samples], optional
        Sample weights.

    binary_metric : callable, returns shape [n_classes]
        The binary metric function to use.

    Returns
    -------
    score : float or array of shape [n_classes]
        If not ``None``, average the score, else return the score for each
        classes.

    """
    average_options = (None, 'micro', 'macro', 'weighted', 'samples')
    if average not in average_options:
        raise ValueError('average has to be one of {0}'
                         ''.format(average_options))

    y_type = type_of_target(y_true)
    if y_type not in ("binary", "multilabel-indicator"):
        raise ValueError("{0} format is not supported".format(y_type))

    if y_type == "binary":
        return binary_metric(y_true, y_score, sample_weight=sample_weight)

    check_consistent_length(y_true, y_score, sample_weight)
    y_true = check_array(y_true)
    y_score = check_array(y_score)

    not_average_axis = 1
    score_weight = sample_weight
    average_weight = None

    if average == "micro":
        if score_weight is not None:
            score_weight = np.repeat(score_weight, y_true.shape[1])
        y_true = y_true.ravel()
        y_score = y_score.ravel()

    elif average == 'weighted':
        if score_weight is not None:
            average_weight = np.sum(np.multiply(
                y_true, np.reshape(score_weight, (-1, 1))), axis=0)
        else:
            average_weight = np.sum(y_true, axis=0)
        if average_weight.sum() == 0:
            return 0

    elif average == 'samples':
        # swap average_weight <-> score_weight
        average_weight = score_weight
        score_weight = None
        not_average_axis = 0

    if y_true.ndim == 1:
        y_true = y_true.reshape((-1, 1))

    if y_score.ndim == 1:
        y_score = y_score.reshape((-1, 1))

    n_classes = y_score.shape[not_average_axis]
    score = np.zeros((n_classes,))
    for c in range(n_classes):
        y_true_c = y_true.take([c], axis=not_average_axis).ravel()
        y_score_c = y_score.take([c], axis=not_average_axis).ravel()
        score[c] = binary_metric(y_true_c, y_score_c,
                                 sample_weight=score_weight)

    # Average the results
    if average is not None:
        return np.average(score, weights=average_weight)
    else:
        return score


import unittest

import pandas as pd
import numpy as np

from pgmpy.estimators import HillClimbSearch, K2Score
from pgmpy.models import BayesianNetwork

class TestHillClimbEstimator(unittest.TestCase):
    def setUp(self):
        self.rand_data = pd.DataFrame(
            np.random.randint(0, 5, size=(int(1e4), 2)), columns=list("AB")
        )
        self.rand_data["C"] = self.rand_data["B"]
        self.est_rand = HillClimbSearch(self.rand_data)
        k2score = K2Score(self.rand_data)
        self.score_rand = k2score.local_score
        self.score_structure_prior = k2score.structure_prior_ratio

        self.model1 = BayesianNetwork()
        self.model1.add_nodes_from(["A", "B", "C"])
        self.model1_possible_edges = set(
            [(u, v) for u in self.model1.nodes() for v in self.model1.nodes()]
        )

        self.model2 = self.model1.copy()
        self.model2.add_edge("A", "B")
        self.model2_possible_edges = set(
            [(u, v) for u in self.model2.nodes() for v in self.model2.nodes()]
        )

        # link to dataset: "https://www.kaggle.com/c/titanic/download/train.csv"
        self.titanic_data = pd.read_csv(
            "pgmpy/tests/test_estimators/testdata/titanic_train.csv"
        )
        self.titanic_data1 = self.titanic_data[
            ["Survived", "Sex", "Pclass", "Age", "Embarked"]
        ]
        self.est_titanic1 = HillClimbSearch(self.titanic_data1)
        self.score_titanic1 = K2Score(self.titanic_data1).local_score

        self.titanic_data2 = self.titanic_data[["Survived", "Sex", "Pclass"]]
        self.est_titanic2 = HillClimbSearch(self.titanic_data2)
        self.score_titanic2 = K2Score(self.titanic_data2).local_score

    def test_legal_operations(self):
        model2_legal_ops = list(
            self.est_rand._legal_operations(
                model=self.model2,
                score=self.score_rand,
                structure_score=self.score_structure_prior,
                tabu_list=set(),
                max_indegree=float("inf"),
                black_list=set(),
                white_list=self.model2_possible_edges,
                fixed_edges=set(),
            )
        )
        model2_legal_ops_ref = [
            (("+", ("C", "A")), -28.15602208305154),
            (("+", ("A", "C")), -28.155467430966382),
            (("+", ("C", "B")), 7636.947544933631),
            (("+", ("B", "C")), 7937.805375579936),
            (("-", ("A", "B")), 28.155467430966382),
            (("flip", ("A", "B")), -0.0005546520851567038),
        ]
        self.assertSetEqual(
            set([op for op, score in model2_legal_ops]),
            set([op for op, score in model2_legal_ops_ref]),
        )

    def test_legal_operations_blacklist_whitelist(self):
        model2_legal_ops_bl = list(
            self.est_rand._legal_operations(
                model=self.model2,
                score=self.score_rand,
                structure_score=self.score_structure_prior,
                tabu_list=set(),
                max_indegree=float("inf"),
                black_list=set([("A", "B"), ("A", "C"), ("C", "A"), ("C", "B")]),
                white_list=self.model2_possible_edges,
                fixed_edges=set(),
            )
        )
        model2_legal_ops_bl_ref = [
            ("+", ("B", "C")),
            ("-", ("A", "B")),
            ("flip", ("A", "B")),
        ]
        self.assertSetEqual(
            set([op for op, score in model2_legal_ops_bl]), set(model2_legal_ops_bl_ref)
        )

        model2_legal_ops_wl = list(
            self.est_rand._legal_operations(
                model=self.model2,
                score=self.score_rand,
                structure_score=self.score_structure_prior,
                tabu_list=set(),
                max_indegree=float("inf"),
                black_list=set(),
                white_list=set([("A", "B"), ("A", "C"), ("C", "A"), ("A", "B")]),
                fixed_edges=set(),
            )
        )
        model2_legal_ops_wl_ref = [
            ("+", ("A", "C")),
            ("+", ("C", "A")),
            ("-", ("A", "B")),
        ]
        self.assertSetEqual(
            set([op for op, score in model2_legal_ops_wl]), set(model2_legal_ops_wl_ref)
        )

    def test_legal_operations_titanic(self):
        start_model = BayesianNetwork(
            [("Survived", "Sex"), ("Pclass", "Age"), ("Pclass", "Embarked")]
        )
        all_possible_edges = set(
            [(u, v) for u in start_model.nodes() for v in start_model.nodes()]
        )
        legal_ops = self.est_titanic1._legal_operations(
            model=start_model,
            score=self.score_titanic1,
            structure_score=self.score_structure_prior,
            tabu_list=[],
            max_indegree=float("inf"),
            black_list=set(),
            white_list=all_possible_edges,
            fixed_edges=set(),
        )
        self.assertEqual(len(list(legal_ops)), 20)

        tabu_list = [
            ("-", ("Survived", "Sex")),
            ("-", ("Survived", "Pclass")),
            ("flip", ("Age", "Pclass")),
        ]
        legal_ops_tabu = self.est_titanic1._legal_operations(
            model=start_model,
            score=self.score_titanic1,
            structure_score=self.score_structure_prior,
            tabu_list=tabu_list,
            max_indegree=float("inf"),
            black_list=set(),
            white_list=all_possible_edges,
            fixed_edges=set(),
        )
        self.assertEqual(len(list(legal_ops_tabu)), 18)

        legal_ops_indegree = self.est_titanic1._legal_operations(
            model=start_model,
            score=self.score_titanic1,
            structure_score=self.score_structure_prior,
            tabu_list=[],
            max_indegree=1,
            black_list=set(),
            white_list=all_possible_edges,
            fixed_edges=set(),
        )
        self.assertEqual(len(list(legal_ops_indegree)), 11)

        legal_ops_both = self.est_titanic1._legal_operations(
            model=start_model,
            score=self.score_titanic1,
            structure_score=self.score_structure_prior,
            tabu_list=tabu_list,
            max_indegree=1,
            black_list=set(),
            white_list=all_possible_edges,
            fixed_edges=set(),
        )

        legal_ops_both_ref = {
            ("+", ("Embarked", "Survived")): 10.050632580087495,
            ("+", ("Survived", "Pclass")): 41.8886804654893,
            ("+", ("Age", "Survived")): -23.635716036430722,
            ("+", ("Pclass", "Survived")): 41.81314459373152,
            ("+", ("Sex", "Pclass")): 4.772261678791324,
            ("-", ("Pclass", "Age")): 11.546515590730905,
            ("-", ("Pclass", "Embarked")): -32.17148283253266,
            ("flip", ("Pclass", "Embarked")): 3.3563814191275583,
            ("flip", ("Survived", "Sex")): 0.0397370279797542,
        }
        self.assertSetEqual(
            set([op for op, score in legal_ops_both]), set(legal_ops_both_ref)
        )
        for op, score in legal_ops_both:
            self.assertAlmostEqual(score, legal_ops_both_ref[op])

    def test_estimate_rand(self):
        est1 = self.est_rand.estimate()
        self.assertSetEqual(set(est1.nodes()), set(["A", "B", "C"]))
        self.assertTrue(
            list(est1.edges()) == [("B", "C")] or list(est1.edges()) == [("C", "B")]
        )

        est2 = self.est_rand.estimate(
            start_dag=BayesianNetwork([("A", "B"), ("A", "C")])
        )
        self.assertTrue(
            list(est2.edges()) == [("B", "C")] or list(est2.edges()) == [("C", "B")]
        )

        est3 = self.est_rand.estimate(fixed_edges=[("B", "C")])
        self.assertTrue([("B", "C")] == list(est3.edges()))

    def test_estimate_titanic(self):
        self.assertSetEqual(
            set(self.est_titanic2.estimate().edges()),
            set([("Survived", "Pclass"), ("Sex", "Pclass"), ("Sex", "Survived")]),
        )

        self.assertTrue(
            ("Pclass", "Survived")
            in self.est_titanic2.estimate(fixed_edges=[("Pclass", "Survived")]).edges()
        )

    def test_no_legal_operation(self):
        data = pd.DataFrame(
            [
                [1, 0, 0, 1, 0, 0, 1, 1, 0],
                [1, 0, 1, 0, 0, 1, 0, 1, 0],
                [1, 0, 0, 0, 0, 1, 0, 1, 1],
                [1, 1, 0, 1, 0, 1, 1, 0, 0],
                [0, 0, 1, 0, 0, 1, 1, 0, 0],
            ],
            columns=list("ABCDEFGHI"),
        )
        est = HillClimbSearch(data)
        best_model = est.estimate(
            fixed_edges=[("A", "B"), ("B", "C")], white_list=[("F", "C")]
        )
        self.assertEqual(
            set(best_model.edges()), set([("A", "B"), ("B", "C"), ("F", "C")])
        )

    def test_estimate(self):
        for score in ["k2score", "bdeuscore", "bdsscore", "bicscore"]:
            dag = self.est_rand.estimate(scoring_method=score)
            dag = self.est_titanic1.estimate(scoring_method=score)

    def tearDown(self):
        del self.rand_data
        del self.est_rand
        del self.model1
        del self.titanic_data
        del self.titanic_data1
        del self.titanic_data2
        del self.est_titanic1
        del self.est_titanic2


import warnings
import numpy as np
from pandas import compat
from pandas._libs import reduction
from pandas.core.dtypes.generic import ABCSeries
from pandas.core.dtypes.common import (
    is_extension_type,
    is_dict_like,
    is_list_like,
    is_sequence)
from pandas.util._decorators import cache_readonly

from pandas.io.formats.printing import pprint_thing

def frame_apply(obj, func, axis=0, broadcast=None,
                raw=False, reduce=None, result_type=None,
                ignore_failures=False,
                args=None, kwds=None):
    """ construct and return a row or column based frame apply object """

    axis = obj._get_axis_number(axis)
    if axis == 0:
        klass = FrameRowApply
    elif axis == 1:
        klass = FrameColumnApply

    return klass(obj, func, broadcast=broadcast,
                 raw=raw, reduce=reduce, result_type=result_type,
                 ignore_failures=ignore_failures,
                 args=args, kwds=kwds)

class FrameApply(object):

    def __init__(self, obj, func, broadcast, raw, reduce, result_type,
                 ignore_failures, args, kwds):
        self.obj = obj
        self.raw = raw
        self.ignore_failures = ignore_failures
        self.args = args or ()
        self.kwds = kwds or {}

        if result_type not in [None, 'reduce', 'broadcast', 'expand']:
            raise ValueError("invalid value for result_type, must be one "
                             "of {None, 'reduce', 'broadcast', 'expand'}")

        if broadcast is not None:
            warnings.warn("The broadcast argument is deprecated and will "
                          "be removed in a future version. You can specify "
                          "result_type='broadcast' to broadcast the result "
                          "to the original dimensions",
                          FutureWarning, stacklevel=4)
            if broadcast:
                result_type = 'broadcast'

        if reduce is not None:
            warnings.warn("The reduce argument is deprecated and will "
                          "be removed in a future version. You can specify "
                          "result_type='reduce' to try to reduce the result "
                          "to the original dimensions",
                          FutureWarning, stacklevel=4)
            if reduce:

                if result_type is not None:
                    raise ValueError(
                        "cannot pass both reduce=True and result_type")

                result_type = 'reduce'

        self.result_type = result_type

        # curry if needed
        if ((kwds or args) and
                not isinstance(func, (np.ufunc, compat.string_types))):

            def f(x):
                return func(x, *args, **kwds)
        else:
            f = func

        self.f = f

        # results
        self.result = None
        self.res_index = None
        self.res_columns = None

    @property
    def columns(self):
        return self.obj.columns

    @property
    def index(self):
        return self.obj.index

    @cache_readonly
    def values(self):
        return self.obj.values

    @cache_readonly
    def dtypes(self):
        return self.obj.dtypes

    @property
    def agg_axis(self):
        return self.obj._get_agg_axis(self.axis)

    def get_result(self):
        """ compute the results """

        # dispatch to agg
        if is_list_like(self.f) or is_dict_like(self.f):
            return self.obj.aggregate(self.f, axis=self.axis,
                                      *self.args, **self.kwds)

        # all empty
        if len(self.columns) == 0 and len(self.index) == 0:
            return self.apply_empty_result()

        # string dispatch
        if isinstance(self.f, compat.string_types):
            # Support for `frame.transform('method')`
            # Some methods (shift, etc.) require the axis argument, others
            # don't, so inspect and insert if necessary.
            func = getattr(self.obj, self.f)
            sig = compat.signature(func)
            if 'axis' in sig.args:
                self.kwds['axis'] = self.axis
            return func(*self.args, **self.kwds)

        # ufunc
        elif isinstance(self.f, np.ufunc):
            with np.errstate(all='ignore'):
                results = self.f(self.values)
            return self.obj._constructor(data=results, index=self.index,
                                         columns=self.columns, copy=False)

        # broadcasting
        if self.result_type == 'broadcast':
            return self.apply_broadcast()

        # one axis empty
        elif not all(self.obj.shape):
            return self.apply_empty_result()

        # raw
        elif self.raw and not self.obj._is_mixed_type:
            return self.apply_raw()

        return self.apply_standard()

    def apply_empty_result(self):
        """
        we have an empty result; at least 1 axis is 0

        we will try to apply the function to an empty
        series in order to see if this is a reduction function
        """

        # we are not asked to reduce or infer reduction
        # so just return a copy of the existing object
        if self.result_type not in ['reduce', None]:
            return self.obj.copy()

        # we may need to infer
        reduce = self.result_type == 'reduce'

        from pandas import Series
        if not reduce:

            EMPTY_SERIES = Series([])
            try:
                r = self.f(EMPTY_SERIES, *self.args, **self.kwds)
                reduce = not isinstance(r, Series)
            except Exception:
                pass

        if reduce:
            return self.obj._constructor_sliced(np.nan, index=self.agg_axis)
        else:
            return self.obj.copy()

    def apply_raw(self):
        """ apply to the values as a numpy array """

        try:
            result = reduction.reduce(self.values, self.f, axis=self.axis)
        except Exception:
            result = np.apply_along_axis(self.f, self.axis, self.values)

        # TODO: mixed type case
        if result.ndim == 2:
            return self.obj._constructor(result,
                                         index=self.index,
                                         columns=self.columns)
        else:
            return self.obj._constructor_sliced(result,
                                                index=self.agg_axis)

    def apply_broadcast(self, target):
        result_values = np.empty_like(target.values)

        # axis which we want to compare compliance
        result_compare = target.shape[0]

        for i, col in enumerate(target.columns):
            res = self.f(target[col])
            ares = np.asarray(res).ndim

            # must be a scalar or 1d
            if ares > 1:
                raise ValueError("too many dims to broadcast")
            elif ares == 1:

                # must match return dim
                if result_compare != len(res):
                    raise ValueError("cannot broadcast result")

            result_values[:, i] = res

        # we *always* preserve the original index / columns
        result = self.obj._constructor(result_values,
                                       index=target.index,
                                       columns=target.columns)
        return result

    def apply_standard(self):

        # try to reduce first (by default)
        # this only matters if the reduction in values is of different dtype
        # e.g. if we want to apply to a SparseFrame, then can't directly reduce

        # we cannot reduce using non-numpy dtypes,
        # as demonstrated in gh-12244
        if (self.result_type in ['reduce', None] and
                not self.dtypes.apply(is_extension_type).any()):

            # Create a dummy Series from an empty array
            from pandas import Series
            values = self.values
            index = self.obj._get_axis(self.axis)
            labels = self.agg_axis
            empty_arr = np.empty(len(index), dtype=values.dtype)
            dummy = Series(empty_arr, index=index, dtype=values.dtype)

            try:
                result = reduction.reduce(values, self.f,
                                          axis=self.axis,
                                          dummy=dummy,
                                          labels=labels)
                return self.obj._constructor_sliced(result, index=labels)
            except Exception:
                pass

        # compute the result using the series generator
        self.apply_series_generator()

        # wrap results
        return self.wrap_results()

    def apply_series_generator(self):
        series_gen = self.series_generator
        res_index = self.result_index

        i = None
        keys = []
        results = {}
        if self.ignore_failures:
            successes = []
            for i, v in enumerate(series_gen):
                try:
                    results[i] = self.f(v)
                    keys.append(v.name)
                    successes.append(i)
                except Exception:
                    pass

            # so will work with MultiIndex
            if len(successes) < len(res_index):
                res_index = res_index.take(successes)

        else:
            try:
                for i, v in enumerate(series_gen):
                    results[i] = self.f(v)
                    keys.append(v.name)
            except Exception as e:
                if hasattr(e, 'args'):

                    # make sure i is defined
                    if i is not None:
                        k = res_index[i]
                        e.args = e.args + ('occurred at index %s' %
                                           pprint_thing(k), )
                raise

        self.results = results
        self.res_index = res_index
        self.res_columns = self.result_columns

    def wrap_results(self):
        results = self.results

        # see if we can infer the results
        if len(results) > 0 and is_sequence(results[0]):

            return self.wrap_results_for_axis()

        # dict of scalars
        result = self.obj._constructor_sliced(results)
        result.index = self.res_index

        return result

class FrameRowApply(FrameApply):
    axis = 0

    def apply_broadcast(self):
        return super(FrameRowApply, self).apply_broadcast(self.obj)

    @property
    def series_generator(self):
        return (self.obj._ixs(i, axis=1)
                for i in range(len(self.columns)))

    @property
    def result_index(self):
        return self.columns

    @property
    def result_columns(self):
        return self.index

    def wrap_results_for_axis(self):
        """ return the results for the rows """

        results = self.results
        result = self.obj._constructor(data=results)

        if not isinstance(results[0], ABCSeries):
            try:
                result.index = self.res_columns
            except ValueError:
                pass

        try:
            result.columns = self.res_index
        except ValueError:
            pass

        return result

class FrameColumnApply(FrameApply):
    axis = 1

    def apply_broadcast(self):
        result = super(FrameColumnApply, self).apply_broadcast(self.obj.T)
        return result.T

    @property
    def series_generator(self):
        constructor = self.obj._constructor_sliced
        return (constructor(arr, index=self.columns, name=name)
                for i, (arr, name) in enumerate(zip(self.values,
                                                    self.index)))

    @property
    def result_index(self):
        return self.index

    @property
    def result_columns(self):
        return self.columns

    def wrap_results_for_axis(self):
        """ return the results for the columns """
        results = self.results

        # we have requested to expand
        if self.result_type == 'expand':
            result = self.infer_to_same_shape()

        # we have a non-series and don't want inference
        elif not isinstance(results[0], ABCSeries):
            from pandas import Series
            result = Series(results)
            result.index = self.res_index

        # we may want to infer results
        else:
            result = self.infer_to_same_shape()

        return result

    def infer_to_same_shape(self):
        """ infer the results to the same shape as the input object """
        results = self.results

        result = self.obj._constructor(data=results)
        result = result.T

        # set the index
        result.index = self.res_index

        # infer dtypes
        result = result.infer_objects()

        return result


import os
from setuptools import find_packages, setup

def extract_version():
    """
    Extracts version values from the main matplotlib __init__.py and
    returns them as a dictionary.
    """
    with open('ggplot/__init__.py') as fd:
        for line in fd.readlines():
            if (line.startswith('__version__')):
                exec(line.strip())
    return locals()["__version__"]

def get_package_data():
    baseline_images = [
        'tests/baseline_images/%s/*' % x
        for x in os.listdir('ggplot/tests/baseline_images')]

    return {
        'ggplot':
        baseline_images +
        [
            "exampledata/*.csv",
            "geoms/*.png"
        ]}

setup(
    name="ggplot",
    # Increase the version in ggplot/__init__.py
    version=extract_version(),
    author="Greg Lamp",
    author_email="greg@yhathq.com",
    url="https://github.com/yhat/ggplot/",
    license="BSD",
    packages=find_packages(),
    package_dir={"ggplot": "ggplot"},
    package_data=get_package_data(),
    description="ggplot for python",
    # run pandoc --from=markdown --to=rst --output=README.rst README.md
    long_description=open("README.rst").read(),
    # numpy is here to make installing easier... Needs to be at the last position,
    # as that's the first installed with "python setup.py install"
    install_requires=["six", "statsmodels", "brewer2mpl", "matplotlib", "scipy",
                      "patsy", "pandas", "numpy"],
    classifiers=['Intended Audience :: Science/Research',
                 'Intended Audience :: Developers',
                 'Programming Language :: Python',
                 'Topic :: Software Development',
                 'Topic :: Scientific/Engineering',
                 'Operating System :: Microsoft :: Windows',
                 'Operating System :: POSIX',
                 'Operating System :: Unix',
                 'Operating System :: MacOS',
                 'Programming Language :: Python :: 2',
                 'Programming Language :: Python :: 2.7',
                 'Programming Language :: Python :: 3',
                 'Programming Language :: Python :: 3.3'],
    zip_safe=False)


# Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
# License: BSD 3 clause

import numpy as np
from scipy import sparse

from sklearn.utils.testing import (assert_array_almost_equal, assert_equal,
                                   assert_greater, assert_almost_equal,
                                   assert_greater_equal,
                                   assert_array_equal,
                                   assert_raises,
                                   assert_warns_message)
from sklearn.datasets import make_classification, make_blobs
from sklearn.naive_bayes import MultinomialNB
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.svm import LinearSVC
from sklearn.linear_model import Ridge
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import Imputer
from sklearn.metrics import brier_score_loss, log_loss
from sklearn.calibration import CalibratedClassifierCV
from sklearn.calibration import _sigmoid_calibration, _SigmoidCalibration
from sklearn.calibration import calibration_curve

def test_calibration():
    """Test calibration objects with isotonic and sigmoid"""
    n_samples = 100
    X, y = make_classification(n_samples=2 * n_samples, n_features=6,
                               random_state=42)
    sample_weight = np.random.RandomState(seed=42).uniform(size=y.size)

    X -= X.min()  # MultinomialNB only allows positive X

    # split train and test
    X_train, y_train, sw_train = \
        X[:n_samples], y[:n_samples], sample_weight[:n_samples]
    X_test, y_test = X[n_samples:], y[n_samples:]

    # Naive-Bayes
    clf = MultinomialNB().fit(X_train, y_train, sample_weight=sw_train)
    prob_pos_clf = clf.predict_proba(X_test)[:, 1]

    pc_clf = CalibratedClassifierCV(clf, cv=y.size + 1)
    assert_raises(ValueError, pc_clf.fit, X, y)

    # Naive Bayes with calibration
    for this_X_train, this_X_test in [(X_train, X_test),
                                      (sparse.csr_matrix(X_train),
                                       sparse.csr_matrix(X_test))]:
        for method in ['isotonic', 'sigmoid']:
            pc_clf = CalibratedClassifierCV(clf, method=method, cv=2)
            # Note that this fit overwrites the fit on the entire training
            # set
            pc_clf.fit(this_X_train, y_train, sample_weight=sw_train)
            prob_pos_pc_clf = pc_clf.predict_proba(this_X_test)[:, 1]

            # Check that brier score has improved after calibration
            assert_greater(brier_score_loss(y_test, prob_pos_clf),
                           brier_score_loss(y_test, prob_pos_pc_clf))

            # Check invariance against relabeling [0, 1] -> [1, 2]
            pc_clf.fit(this_X_train, y_train + 1, sample_weight=sw_train)
            prob_pos_pc_clf_relabeled = pc_clf.predict_proba(this_X_test)[:, 1]
            assert_array_almost_equal(prob_pos_pc_clf,
                                      prob_pos_pc_clf_relabeled)

            # Check invariance against relabeling [0, 1] -> [-1, 1]
            pc_clf.fit(this_X_train, 2 * y_train - 1, sample_weight=sw_train)
            prob_pos_pc_clf_relabeled = pc_clf.predict_proba(this_X_test)[:, 1]
            assert_array_almost_equal(prob_pos_pc_clf,
                                      prob_pos_pc_clf_relabeled)

            # Check invariance against relabeling [0, 1] -> [1, 0]
            pc_clf.fit(this_X_train, (y_train + 1) % 2,
                       sample_weight=sw_train)
            prob_pos_pc_clf_relabeled = \
                pc_clf.predict_proba(this_X_test)[:, 1]
            if method == "sigmoid":
                assert_array_almost_equal(prob_pos_pc_clf,
                                          1 - prob_pos_pc_clf_relabeled)
            else:
                # Isotonic calibration is not invariant against relabeling
                # but should improve in both cases
                assert_greater(brier_score_loss(y_test, prob_pos_clf),
                               brier_score_loss((y_test + 1) % 2,
                                                prob_pos_pc_clf_relabeled))

        # check that calibration can also deal with regressors that have
        # a decision_function
        clf_base_regressor = CalibratedClassifierCV(Ridge())
        clf_base_regressor.fit(X_train, y_train)
        clf_base_regressor.predict(X_test)

        # Check failure cases:
        # only "isotonic" and "sigmoid" should be accepted as methods
        clf_invalid_method = CalibratedClassifierCV(clf, method="foo")
        assert_raises(ValueError, clf_invalid_method.fit, X_train, y_train)

        # base-estimators should provide either decision_function or
        # predict_proba (most regressors, for instance, should fail)
        clf_base_regressor = \
            CalibratedClassifierCV(RandomForestRegressor(), method="sigmoid")
        assert_raises(RuntimeError, clf_base_regressor.fit, X_train, y_train)

def test_sample_weight_warning():
    n_samples = 100
    X, y = make_classification(n_samples=2 * n_samples, n_features=6,
                               random_state=42)

    sample_weight = np.random.RandomState(seed=42).uniform(size=len(y))
    X_train, y_train, sw_train = \
        X[:n_samples], y[:n_samples], sample_weight[:n_samples]
    X_test = X[n_samples:]

    for method in ['sigmoid', 'isotonic']:
        base_estimator = LinearSVC(random_state=42)
        calibrated_clf = CalibratedClassifierCV(base_estimator, method=method)
        # LinearSVC does not currently support sample weights but they
        # can still be used for the calibration step (with a warning)
        msg = "LinearSVC does not support sample_weight."
        assert_warns_message(
            UserWarning, msg,
            calibrated_clf.fit, X_train, y_train, sample_weight=sw_train)
        probs_with_sw = calibrated_clf.predict_proba(X_test)

        # As the weights are used for the calibration, they should still yield
        # a different predictions
        calibrated_clf.fit(X_train, y_train)
        probs_without_sw = calibrated_clf.predict_proba(X_test)

        diff = np.linalg.norm(probs_with_sw - probs_without_sw)
        assert_greater(diff, 0.1)

def test_calibration_multiclass():
    """Test calibration for multiclass """
    # test multi-class setting with classifier that implements
    # only decision function
    clf = LinearSVC()
    X, y_idx = make_blobs(n_samples=100, n_features=2, random_state=42,
                          centers=3, cluster_std=3.0)

    # Use categorical labels to check that CalibratedClassifierCV supports
    # them correctly
    target_names = np.array(['a', 'b', 'c'])
    y = target_names[y_idx]

    X_train, y_train = X[::2], y[::2]
    X_test, y_test = X[1::2], y[1::2]

    clf.fit(X_train, y_train)
    for method in ['isotonic', 'sigmoid']:
        cal_clf = CalibratedClassifierCV(clf, method=method, cv=2)
        cal_clf.fit(X_train, y_train)
        probas = cal_clf.predict_proba(X_test)
        assert_array_almost_equal(np.sum(probas, axis=1), np.ones(len(X_test)))

        # Check that log-loss of calibrated classifier is smaller than
        # log-loss of naively turned OvR decision function to probabilities
        # via softmax
        def softmax(y_pred):
            e = np.exp(-y_pred)
            return e / e.sum(axis=1).reshape(-1, 1)
        uncalibrated_log_loss = \
            log_loss(y_test, softmax(clf.decision_function(X_test)))
        calibrated_log_loss = log_loss(y_test, probas)
        assert_greater_equal(uncalibrated_log_loss, calibrated_log_loss)

    # Test that calibration of a multiclass classifier decreases log-loss
    # for RandomForestClassifier
    X, y = make_blobs(n_samples=100, n_features=2, random_state=42,
                      cluster_std=3.0)
    X_train, y_train = X[::2], y[::2]
    X_test, y_test = X[1::2], y[1::2]

    clf = RandomForestClassifier(n_estimators=10, random_state=42)
    clf.fit(X_train, y_train)
    clf_probs = clf.predict_proba(X_test)
    loss = log_loss(y_test, clf_probs)

    for method in ['isotonic', 'sigmoid']:
        cal_clf = CalibratedClassifierCV(clf, method=method, cv=3)
        cal_clf.fit(X_train, y_train)
        cal_clf_probs = cal_clf.predict_proba(X_test)
        cal_loss = log_loss(y_test, cal_clf_probs)
        assert_greater(loss, cal_loss)

def test_calibration_prefit():
    """Test calibration for prefitted classifiers"""
    n_samples = 50
    X, y = make_classification(n_samples=3 * n_samples, n_features=6,
                               random_state=42)
    sample_weight = np.random.RandomState(seed=42).uniform(size=y.size)

    X -= X.min()  # MultinomialNB only allows positive X

    # split train and test
    X_train, y_train, sw_train = \
        X[:n_samples], y[:n_samples], sample_weight[:n_samples]
    X_calib, y_calib, sw_calib = \
        X[n_samples:2 * n_samples], y[n_samples:2 * n_samples], \
        sample_weight[n_samples:2 * n_samples]
    X_test, y_test = X[2 * n_samples:], y[2 * n_samples:]

    # Naive-Bayes
    clf = MultinomialNB()
    clf.fit(X_train, y_train, sw_train)
    prob_pos_clf = clf.predict_proba(X_test)[:, 1]

    # Naive Bayes with calibration
    for this_X_calib, this_X_test in [(X_calib, X_test),
                                      (sparse.csr_matrix(X_calib),
                                       sparse.csr_matrix(X_test))]:
        for method in ['isotonic', 'sigmoid']:
            pc_clf = CalibratedClassifierCV(clf, method=method, cv="prefit")

            for sw in [sw_calib, None]:
                pc_clf.fit(this_X_calib, y_calib, sample_weight=sw)
                y_prob = pc_clf.predict_proba(this_X_test)
                y_pred = pc_clf.predict(this_X_test)
                prob_pos_pc_clf = y_prob[:, 1]
                assert_array_equal(y_pred,
                                   np.array([0, 1])[np.argmax(y_prob, axis=1)])

                assert_greater(brier_score_loss(y_test, prob_pos_clf),
                               brier_score_loss(y_test, prob_pos_pc_clf))

def test_sigmoid_calibration():
    """Test calibration values with Platt sigmoid model"""
    exF = np.array([5, -4, 1.0])
    exY = np.array([1, -1, -1])
    # computed from my python port of the C++ code in LibSVM
    AB_lin_libsvm = np.array([-0.20261354391187855, 0.65236314980010512])
    assert_array_almost_equal(AB_lin_libsvm,
                              _sigmoid_calibration(exF, exY), 3)
    lin_prob = 1. / (1. + np.exp(AB_lin_libsvm[0] * exF + AB_lin_libsvm[1]))
    sk_prob = _SigmoidCalibration().fit(exF, exY).predict(exF)
    assert_array_almost_equal(lin_prob, sk_prob, 6)

    # check that _SigmoidCalibration().fit only accepts 1d array or 2d column
    # arrays
    assert_raises(ValueError, _SigmoidCalibration().fit,
                  np.vstack((exF, exF)), exY)

def test_calibration_curve():
    """Check calibration_curve function"""
    y_true = np.array([0, 0, 0, 1, 1, 1])
    y_pred = np.array([0., 0.1, 0.2, 0.8, 0.9, 1.])
    prob_true, prob_pred = calibration_curve(y_true, y_pred, n_bins=2)
    prob_true_unnormalized, prob_pred_unnormalized = \
        calibration_curve(y_true, y_pred * 2, n_bins=2, normalize=True)
    assert_equal(len(prob_true), len(prob_pred))
    assert_equal(len(prob_true), 2)
    assert_almost_equal(prob_true, [0, 1])
    assert_almost_equal(prob_pred, [0.1, 0.9])
    assert_almost_equal(prob_true, prob_true_unnormalized)
    assert_almost_equal(prob_pred, prob_pred_unnormalized)

    # probabilities outside [0, 1] should not be accepted when normalize
    # is set to False
    assert_raises(ValueError, calibration_curve, [1.1], [-0.1],
                  normalize=False)

def test_calibration_nan_imputer():
    """Test that calibration can accept nan"""
    X, y = make_classification(n_samples=10, n_features=2,
                               n_informative=2, n_redundant=0,
                               random_state=42)
    X[0, 0] = np.nan
    clf = Pipeline(
        [('imputer', Imputer()),
         ('rf', RandomForestClassifier(n_estimators=1))])
    clf_c = CalibratedClassifierCV(clf, cv=2, method='isotonic')
    clf_c.fit(X, y)
    clf_c.predict(X)


# from collections import OrderedDict

import numpy as np
from sklearn.metrics import precision_score

from sklearn_evaluation.preprocessing import binarize
from sklearn_evaluation import util
from sklearn_evaluation import validate

def compute_at_thresholds(fn, y_true, y_score, n_thresholds=10, start=0.0):
    """
    Given scores, binarize them at different thresholds, then compute
    metrics

    Examples
    --------
    >>> from sklearn_evaluation.metrics import compute_at_thresholds
    >>> from sklearn.metrics import accuracy_score
    >>> from sklearn.metrics import precision_score, recall_score, f1_score
    >>> import numpy as np
    >>> y_true = np.array([1, 1, 1, 1, 1, 0, 0, 0, 0, 0])
    >>> y_score = np.array([1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1])
    >>> binarized = compute_at_thresholds([accuracy_score, precision_score,
    ...                                    recall_score, f1_score],
    ...                                    y_true, y_score)
    """
    if util.isiter(fn):
        (thresholds,
         Y_pred) = binarize.scores_at_thresholds(y_score,
                                                 n_thresholds=n_thresholds)
        metrics = [np.array([fn_(y_true, y_pred) for y_pred in Y_pred])
                   for fn_ in fn]
        return thresholds, metrics
    else:
        (thresholds,
         Y_pred) = binarize.scores_at_thresholds(y_score,
                                                 n_thresholds=n_thresholds)
        metrics = np.array([fn(y_true, y_pred) for y_pred in Y_pred])
        return thresholds, metrics

def confusion_matrix(y_true, y_pred, normalize):
    pass

@validate.argument_is_proportion('top_proportion')
def precision_at(y_true, y_score, top_proportion, ignore_nas=False):
    '''
    Calculates precision at a given proportion.
    Only supports binary classification.
    '''
    # Sort scores in descending order
    scores_sorted = np.sort(y_score)[::-1]

    # Based on the proportion, get the index to split the data
    # if value is negative, return 0
    cutoff_index = max(int(len(y_true) * top_proportion) - 1, 0)
    # Get the cutoff value
    cutoff_value = scores_sorted[cutoff_index]

    # Convert scores to binary, by comparing them with the cutoff value
    scores_binary = np.array([int(y >= cutoff_value) for y in y_score])
    # Calculate precision using sklearn function
    if ignore_nas:
        precision = __precision(y_true, scores_binary)
    else:
        precision = precision_score(y_true, scores_binary)

    return precision, cutoff_value

def __precision(y_true, y_pred):
    '''
        Precision metric tolerant to unlabeled data in y_true,
        NA values are ignored for the precision calculation
    '''
    # make copies of the arrays to avoid modifying the original ones
    y_true = np.copy(y_true)
    y_pred = np.copy(y_pred)

    # precision = tp/(tp+fp)
    # True nehatives do not affect precision value, so for every missing
    # value in y_true, replace it with 0 and also replace the value
    # in y_pred with 0
    is_nan = np.isnan(y_true)
    y_true[is_nan] = 0
    y_pred[is_nan] = 0
    precision = precision_score(y_true, y_pred)
    return precision

@validate.argument_is_proportion('top_proportion')
def tp_at(y_true, y_score, top_proportion):
    y_pred = binarize.scores_at_top_proportion(y_score, top_proportion)
    tp = (y_pred == 1) & (y_true == 1)
    return tp.sum()

@validate.argument_is_proportion('top_proportion')
def fp_at(y_true, y_score, top_proportion):
    y_pred = binarize.scores_at_top_proportion(y_score, top_proportion)
    fp = (y_pred == 1) & (y_true == 0)
    return fp.sum()

@validate.argument_is_proportion('top_proportion')
def tn_at(y_true, y_score, top_proportion):
    y_pred = binarize.scores_at_top_proportion(y_score, top_proportion)
    tn = (y_pred == 0) & (y_true == 0)
    return tn.sum()

@validate.argument_is_proportion('top_proportion')
def fn_at(y_true, y_score, top_proportion):
    y_pred = binarize.scores_at_top_proportion(y_score, top_proportion)
    fn = (y_pred == 0) & (y_true == 1)
    return fn.sum()

@validate.argument_is_proportion('top_proportion')
def labels_at(y_true, y_score, top_proportion, normalize=False):
    '''
        Return the number of labels encountered in the top  X proportion
    '''
    # Get indexes of scores sorted in descending order
    indexes = np.argsort(y_score)[::-1]

    # Sort true values in the same order
    y_true_sorted = y_true[indexes]

    # Grab top x proportion of true values
    cutoff_index = max(int(len(y_true_sorted) * top_proportion) - 1, 0)
    # add one to index to grab values including that index
    y_true_top = y_true_sorted[:cutoff_index + 1]

    # Count the number of non-nas in the top x proportion
    # we are returning a count so it should be an int
    values = int((~np.isnan(y_true_top)).sum())

    if normalize:
        values = float(values) / (~np.isnan(y_true)).sum()

    return values


import numpy as np
from sklearn.ensemble import GradientBoostingRegressor
from random import sample, seed
from sklearn.decomposition import TruncatedSVD
from math import floor
from sklearn import cross_validation

import numpy as np
from numpy.linalg import norm, svd

def inexact_augmented_lagrange_multiplier(X, lmbda=.01, tol=1e-3,
                                          maxiter=100, verbose=True):
    """
    Inexact Augmented Lagrange Multiplier
    """
    Y = X
    norm_two = norm(Y.ravel(), 2)
    norm_inf = norm(Y.ravel(), np.inf) / lmbda
    dual_norm = np.max([norm_two, norm_inf])
    Y = Y / dual_norm
    A = np.zeros(Y.shape)
    E = np.zeros(Y.shape)
    dnorm = norm(X, 'fro')
    mu = 1.25 / norm_two
    rho = 1.5
    sv = 10.
    n = Y.shape[0]
    itr = 0
    while True:
        Eraw = X - A + (1/mu) * Y
        Eupdate = np.maximum(Eraw - lmbda / mu, 0) + np.minimum(Eraw + lmbda / mu, 0)
        U, S, V = svd(X - Eupdate + (1 / mu) * Y, full_matrices=False)
        svp = (S > 1 / mu).shape[0]
        if svp < sv:
            sv = np.min([svp + 1, n])
        else:
            sv = np.min([svp + round(.05 * n), n])
        Aupdate = np.dot(np.dot(U[:, :svp], np.diag(S[:svp] - 1 / mu)), V[:svp, :])
        A = Aupdate
        E = Eupdate
        Z = X - A - E
        Y = Y + mu * Z
        mu = np.min([mu * rho, mu * 1e7])
        itr += 1
        if ((norm(Z, 'fro') / dnorm) < tol) or (itr >= maxiter):
            break
    if verbose:
        print "Finished at iteration %d" % (itr)
    return A, E

class RegressionDriver(object):
    """Class for Regression-based analysis of Driver traces"""

    def __init__(self, driver, datadict, numberofrows=40): #, numfeatures = 200):
        """Initialize by providing a (positive) driver example and a dictionary of (negative) driver references."""
        seed(42)
        self.driver = driver
        self.numfeatures = self.driver.num_features
        featurelist = []
        self.__clf = GradientBoostingRegressor(n_estimators=300, max_depth=4, min_samples_leaf=2)
        # gbr = GradientBoostingRegressor(n_estimators=500, max_depth=10, max_features=numfeatures, random_state=42)
        # pca = PCA(whiten=True, n_components=numfeatures)
        # estimators = [('polyf', PolynomialFeatures()), ('scale', MinMaxScaler()), ('pca', PCA()), ('gbr', gbr)]
        # self.__clf = Pipeline(estimators)
        self.__indexlist = []
        for trace in self.driver.traces:
            self.__indexlist.append(trace.identifier)
            featurelist.append(trace.features)
        # Initialize train and test np arrays
        self.__traindata = np.asarray(featurelist)
        self.__testdata = np.asarray(featurelist)
        self.__trainlabels = np.ones((self.__traindata.shape[0],))
        data = np.empty((0, self.numfeatures), float)
        setkeys = datadict.keys()
        if driver.identifier in setkeys:
            setkeys.remove(driver.identifier)
        else:
            setkeys = sample(setkeys, len(setkeys) - 1)
        for key in setkeys:
            if key != driver.identifier:
                rand_smpl = [datadict[key][i] for i in sorted(sample(xrange(len(datadict[key])), numberofrows)) ]
                data = np.append(data, np.asarray(rand_smpl), axis=0)
        self.__traindata = np.append(self.__traindata, data, axis=0)
        self.__trainlabels = np.append(self.__trainlabels, np.zeros((data.shape[0],)), axis=0)
        self.__y = np.zeros((self.__testdata.shape[0],))

    def classify(self, nfolds=4):
        """Perform classification"""
        components = self.__traindata.shape[1]

        _, train_rpca_X_np = inexact_augmented_lagrange_multiplier(np.nan_to_num(self.__traindata))
        _, test_rpca_X_np = inexact_augmented_lagrange_multiplier(np.nan_to_num(self.__testdata))
        skf = cross_validation.StratifiedKFold(self.__trainlabels, n_folds=nfolds)

        for train_index, _ in skf:
            X_train = train_rpca_X_np[train_index]
            y_train = self.__trainlabels[train_index]
            self.__clf.fit(X_train, y_train)
            self.__y += self.__clf.predict(test_rpca_X_np)
        self.__y /= float(nfolds)
        # feature_importance = self.__clf.feature_importances_
        # feature_importance = 100.0 * (feature_importance / feature_importance.max())
        # print feature_importance

    def toKaggle(self):
        """Return string in Kaggle submission format"""
        returnstring = ""
        for i in xrange(len(self.__indexlist) - 1):
            returnstring += "%d_%d,%.6f\n" % (self.driver.identifier, self.__indexlist[i], self.__y[i])
        returnstring += "%d_%d,%.6f" % (self.driver.identifier, self.__indexlist[len(self.__indexlist)-1], self.__y[len(self.__indexlist)-1])
        return returnstring

    def validate(self, datadict):
        from sklearn.metrics import roc_auc_score
        testdata = np.empty((0, self.numfeatures), float)
        y_true = np.empty((0,), float)
        for key in datadict.keys():
            currenttestdata = np.asarray(datadict[key])
            testdata = np.append(testdata, currenttestdata, axis=0)
            if key != self.driver.identifier:
                y_true = np.append(y_true, np.zeros((currenttestdata.shape[0],)), axis=0)
            else:
                y_true = np.append(y_true, np.ones((currenttestdata.shape[0],)), axis=0)
        y_score = self.__clf.predict(testdata)
        result = roc_auc_score(y_true, y_score)
        return result


#  Copyright (C) 2010 Antoine Drouin
#
# This file is part of Paparazzi.
#
# Paparazzi is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2, or (at your option)
# any later version.
#
# Paparazzi 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Paparazzi; see the file COPYING.  If not, write to
# the Free Software Foundation, 59 Temple Place - Suite 330,
# Boston, MA 02111-1307, USA.
#

from __future__ import print_function, division

import re
import numpy as np
from numpy import sin, cos
from scipy import linalg, stats

import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def get_ids_in_log(filename):
    """Returns available ac_id from a log."""
    f = open(filename, 'r')
    ids = []
    pattern = re.compile("\S+ (\S+)")
    while True:
        line = f.readline().strip()
        if line == '':
            break
        m = re.match(pattern, line)
        if m:
            ac_id = m.group(1)
            if not ac_id in ids:
                ids.append(ac_id)
    return ids

def read_log(ac_id, filename, sensor):
    """Extracts raw sensor measurements from a log."""
    f = open(filename, 'r')
    pattern = re.compile("(\S+) "+ac_id+" IMU_"+sensor+"_RAW (\S+) (\S+) (\S+)")
    list_meas = []
    while True:
        line = f.readline().strip()
        if line == '':
            break
        m = re.match(pattern, line)
        if m:
            list_meas.append([float(m.group(2)), float(m.group(3)), float(m.group(4))])
    return np.array(list_meas)

def read_log_scaled(ac_id, filename, sensor, t_start, t_end):
    """Extracts scaled sensor measurements from a log."""
    f = open(filename, 'r')
    pattern = re.compile("(\S+) "+ac_id+" IMU_"+sensor+"_SCALED (\S+) (\S+) (\S+)")
    list_meas = []
    while True:
        line = f.readline().strip()
        if line == '':
            break
        m = re.match(pattern, line)
        if m:
            if (float(m.group(1)) >= float(t_start)) and (float(m.group(1)) < (float(t_end)+1.0)):
                list_meas.append([float(m.group(1)), float(m.group(2)), float(m.group(3)), float(m.group(4))])
    return np.array(list_meas)

def read_log_mag_current(ac_id, filename):
    """Extracts raw magnetometer and current measurements from a log."""
    f = open(filename, 'r')
    pattern = re.compile("(\S+) "+ac_id+" IMU_MAG_CURRENT_CALIBRATION (\S+) (\S+) (\S+) (\S+)")
    list_meas = []
    while True:
        line = f.readline().strip()
        if line == '':
            break
        m = re.match(pattern, line)
        if m:
            list_meas.append([float(m.group(2)), float(m.group(3)), float(m.group(4)), float(m.group(5))])
    return np.array(list_meas)

def filter_meas(meas, window_size, noise_threshold):
    """Select only non-noisy data."""
    filtered_meas = []
    filtered_idx = []
    for i in range(window_size, len(meas)-window_size):
        noise = meas[i-window_size:i+window_size, :].std(axis=0)
        if linalg.norm(noise) < noise_threshold:
            filtered_meas.append(meas[i, :])
            filtered_idx.append(i)
    return np.array(filtered_meas), filtered_idx

def get_min_max_guess(meas, scale):
    """Initial boundary based calibration."""
    max_meas = meas[:, :].max(axis=0)
    min_meas = meas[:, :].min(axis=0)
    range = max_meas - min_meas
    # check if we would get division by zero
    if range.all():
        n = (max_meas + min_meas) / 2
        sf = 2*scale/range
        return np.array([n[0], n[1], n[2], sf[0], sf[1], sf[2]])
    else:
        return np.array([0, 0, 0, 0])

def scale_measurements(meas, p):
    """Scale the set of measurements."""
    l_comp = []
    l_norm = []
    for m in meas[:, ]:
        sm = (m - p[0:3])*p[3:6]
        l_comp.append(sm)
        l_norm.append(linalg.norm(sm))
    return np.array(l_comp), np.array(l_norm)

def estimate_mag_current_relation(meas):
    """Calculate linear coefficient of magnetometer-current relation."""
    coefficient = []
    for i in range(0, 3):
        gradient, intercept, r_value, p_value, std_err = stats.linregress(meas[:, 3], meas[:, i])
        coefficient.append(gradient)
    return coefficient

def print_xml(p, sensor, res):
    """Print xml for airframe file."""
    print("")
    print("<define name=\""+sensor+"_X_NEUTRAL\" value=\""+str(int(round(p[0])))+"\"/>")
    print("<define name=\""+sensor+"_Y_NEUTRAL\" value=\""+str(int(round(p[1])))+"\"/>")
    print("<define name=\""+sensor+"_Z_NEUTRAL\" value=\""+str(int(round(p[2])))+"\"/>")
    print("<define name=\""+sensor+"_X_SENS\" value=\""+str(p[3]*2**res)+"\" integer=\"16\"/>")
    print("<define name=\""+sensor+"_Y_SENS\" value=\""+str(p[4]*2**res)+"\" integer=\"16\"/>")
    print("<define name=\""+sensor+"_Z_SENS\" value=\""+str(p[5]*2**res)+"\" integer=\"16\"/>")
    print("")

def print_imu_scaled(sensor, measurements, attrs):
    print("")
    print(sensor+" : Time Range("+str(measurements[:,0].min(axis=0))+" : "+str(measurements[:,0].max(axis=0))+")")
    np.set_printoptions(formatter={'float': '{:-7.3f}'.format})
    print("         " + attrs[2] + "      " + attrs[3] + "      " + attrs[4])
    print("Min   " + str(measurements[:,1:].min(axis=0)*attrs[0])  + " " + attrs[1])
    print("Max   " + str(measurements[:,1:].max(axis=0)*attrs[0])  + " " + attrs[1])
    print("Mean  " + str(measurements[:,1:].mean(axis=0)*attrs[0]) + " " + attrs[1])
    print("StDev " + str(measurements[:,1:].std(axis=0)*attrs[0])  + " " + attrs[1])

def plot_measurements(sensor, measurements):
    plt.plot(measurements[:, 0])
    plt.plot(measurements[:, 1])
    plt.plot(measurements[:, 2])
    plt.ylabel('ADC')
    plt.title("Raw %s measurements" % sensor)
    plt.show()

def plot_results(sensor, measurements, flt_idx, flt_meas, cp0, np0, cp1, np1, sensor_ref, blocking=True):
    """Plot calibration results."""
    # plot raw measurements with filtered ones marked as red circles
    plt.subplot(3, 1, 1)
    plt.plot(measurements[:, 0])
    plt.plot(measurements[:, 1])
    plt.plot(measurements[:, 2])
    plt.plot(flt_idx, flt_meas[:, 0], 'ro')
    plt.plot(flt_idx, flt_meas[:, 1], 'ro')
    plt.plot(flt_idx, flt_meas[:, 2], 'ro')
    plt.ylabel('ADC')
    plt.title('Raw '+sensor+', red dots are actually used measurements')

    plt.tight_layout()
    # show scaled measurements with initial guess
    plt.subplot(3, 2, 3)
    plt.plot(cp0[:, 0])
    plt.plot(cp0[:, 1])
    plt.plot(cp0[:, 2])
    plt.plot(-sensor_ref*np.ones(len(flt_meas)))
    plt.plot(sensor_ref*np.ones(len(flt_meas)))
    plt.title('scaled '+sensor+' (initial guess)')
    plt.xticks([])

    plt.subplot(3, 2, 4)
    plt.plot(np0)
    plt.plot(sensor_ref*np.ones(len(flt_meas)))
    plt.title('norm of '+sensor+' (initial guess)')
    plt.xticks([])

    # show scaled measurements after optimization
    plt.subplot(3, 2, 5)
    plt.plot(cp1[:, 0])
    plt.plot(cp1[:, 1])
    plt.plot(cp1[:, 2])
    plt.plot(-sensor_ref*np.ones(len(flt_meas)))
    plt.plot(sensor_ref*np.ones(len(flt_meas)))
    plt.title('scaled '+sensor+' (optimized)')
    plt.xticks([])

    plt.subplot(3, 2, 6)
    plt.plot(np1)
    plt.plot(sensor_ref*np.ones(len(flt_meas)))
    plt.title('norm of '+sensor+' (optimized)')
    plt.xticks([])

    # if we want to have another plot we only draw the figure (non-blocking)
    # also in matplotlib before 1.0.0 there is only one call to show possible
    if blocking:
        plt.show()
    else:
        plt.draw()

def plot_imu_scaled(sensor, measurements, attrs):
    """Plot imu scaled results."""
    plt.figure("Sensor Scaled")

    plt.subplot(4, 1, 1)
    plt.plot(measurements[:, 0], measurements[:, 1]*attrs[0])
    plt.plot(measurements[:, 0], measurements[:, 2]*attrs[0])
    plt.plot(measurements[:, 0], measurements[:, 3]*attrs[0])
    #plt.xlabel('Time (s)')
    plt.ylabel(attrs[1])
    plt.title(sensor)

    plt.subplot(4, 1, 2)
    plt.plot(measurements[:, 0], measurements[:, 1]*attrs[0], 'b')
    #plt.xlabel('Time (s)')
    plt.ylabel(attrs[2])

    plt.subplot(4, 1, 3)
    plt.plot(measurements[:, 0], measurements[:, 2]*attrs[0], 'g')
    #plt.xlabel('Time (s)')
    plt.ylabel(attrs[3])

    plt.subplot(4, 1, 4)
    plt.plot(measurements[:, 0], measurements[:, 3]*attrs[0], 'r')
    plt.xlabel('Time (s)')
    plt.ylabel(attrs[4])

    plt.show()

def plot_imu_scaled_fft(sensor, measurements, attrs):
    """Plot imu scaled fft results."""
    #dt = 0.0769
    #Fs = 1/dt
    Fs = 26.0

    plt.figure("Sensor Scaled - FFT")

    plt.subplot(3, 1, 1)
    plt.magnitude_spectrum(measurements[:, 1]*attrs[0], Fs=Fs, scale='linear')
    plt.ylabel(attrs[2])
    plt.title(sensor)

    plt.subplot(3, 1, 2)
    plt.magnitude_spectrum(measurements[:, 2]*attrs[0], Fs=Fs, scale='linear')
    plt.ylabel(attrs[3])

    plt.subplot(3, 1, 3)
    plt.magnitude_spectrum(measurements[:, 3]*attrs[0], Fs=Fs, scale='linear')
    plt.xlabel('Frequency')
    plt.ylabel(attrs[4])

    plt.show()

def plot_mag_3d(measured, calibrated, p):
    """Plot magnetometer measurements on 3D sphere."""
    # set up points for sphere and ellipsoid wireframes
    u = np.r_[0:2 * np.pi:20j]
    v = np.r_[0:np.pi:20j]
    wx = np.outer(cos(u), sin(v))
    wy = np.outer(sin(u), sin(v))
    wz = np.outer(np.ones(np.size(u)), cos(v))
    ex = p[0] * np.ones(np.size(u)) + np.outer(cos(u), sin(v)) / p[3]
    ey = p[1] * np.ones(np.size(u)) + np.outer(sin(u), sin(v)) / p[4]
    ez = p[2] * np.ones(np.size(u)) + np.outer(np.ones(np.size(u)), cos(v)) / p[5]

    # measurements
    mx = measured[:, 0]
    my = measured[:, 1]
    mz = measured[:, 2]

    # calibrated values
    cx = calibrated[:, 0]
    cy = calibrated[:, 1]
    cz = calibrated[:, 2]

    # axes size
    left = 0.02
    bottom = 0.05
    width = 0.46
    height = 0.9
    rect_l = [left, bottom, width, height]
    rect_r = [left/2+0.5, bottom, width, height]

    fig = plt.figure(figsize=plt.figaspect(0.5))
    if matplotlib.__version__.startswith('0'):
        ax = Axes3D(fig, rect=rect_l)
    else:
        ax = fig.add_subplot(1, 2, 1, position=rect_l, projection='3d')
    # plot measurements
    ax.scatter(mx, my, mz)
    plt.hold(True)
    # plot line from center to ellipsoid center
    ax.plot([0.0, p[0]], [0.0, p[1]], [0.0, p[2]], color='black', marker='+', markersize=10)
    # plot ellipsoid
    ax.plot_wireframe(ex, ey, ez, color='grey', alpha=0.5)

    # Create cubic bounding box to simulate equal aspect ratio
    max_range = np.array([mx.max() - mx.min(), my.max() - my.min(), mz.max() - mz.min()]).max()
    Xb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2, -1:2:2][0].flatten() + 0.5 * (mx.max() + mx.min())
    Yb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2, -1:2:2][1].flatten() + 0.5 * (my.max() + my.min())
    Zb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2, -1:2:2][2].flatten() + 0.5 * (mz.max() + mz.min())
    # add the fake bounding box:
    for xb, yb, zb in zip(Xb, Yb, Zb):
        ax.plot([xb], [yb], [zb], 'w')

    ax.set_title('MAG raw with fitted ellipsoid and center offset')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    if matplotlib.__version__.startswith('0'):
        ax = Axes3D(fig, rect=rect_r)
    else:
        ax = fig.add_subplot(1, 2, 2, position=rect_r, projection='3d')
    ax.plot_wireframe(wx, wy, wz, color='grey', alpha=0.5)
    plt.hold(True)
    ax.scatter(cx, cy, cz)

    ax.set_title('MAG calibrated on unit sphere')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    ax.set_xlim3d(-1, 1)
    ax.set_ylim3d(-1, 1)
    ax.set_zlim3d(-1, 1)
    plt.show()

def read_turntable_log(ac_id, tt_id, filename, _min, _max):
    """ Read a turntable log.
    return an array which first column is turnatble and next 3 are gyro
    """
    f = open(filename, 'r')
    pattern_g = re.compile("(\S+) "+str(ac_id)+" IMU_GYRO_RAW (\S+) (\S+) (\S+)")
    pattern_t = re.compile("(\S+) "+str(tt_id)+" IMU_TURNTABLE (\S+)")
    last_tt = None
    list_tt = []
    while True:
        line = f.readline().strip()
        if line == '':
            break
        m = re.match(pattern_t, line)
        if m:
            last_tt = float(m.group(2))
        m = re.match(pattern_g, line)
        if m and last_tt and _min < last_tt < _max:
            list_tt.append([last_tt, float(m.group(2)), float(m.group(3)), float(m.group(4))])
    return np.array(list_tt)


print("######################################################################")
print("# Parallel n-split k-stratified-fold continuous SVM Scikitlearn MVPA #")
print("# (c) Jean-Remi King 2012, jeanremi.king [at] gmail [dot] com        #")
print("######################################################################")

# Implementation of a multivariate pattern analysis based on the scikit-learn
# toolbox (http://scikit-learn.org/stable/). It reads two .mat files
# (filenameX, filenamey) created by 'jr_classify.m'
#
# Function:
#     skl_king_parallel.py filenameX filenamey [number_of_cores]
#
# Inputs:
#     in filenameX:
#       Xm:     samples x features x classification matrix (e.g. trials x
#               chans x time)
#     in filenamey:
#       y:      vector indicating the class of each sample. Negative values
#               will be used for generalization only. 0 indicates to-be-
#               ignored samples.
#       y2:     cost/weights applied on each sample
#       path:   export directory
#       nameX:  export filename X
#       namey:  export filename y
#       folding:type of folding(e.g. stratified)
#       n_splits:number of splits
#       n_folds: number of folds
#       C:       SVM penalization parameter
#       compute_probas:     compute logit fit
#       compute_predict:    compute traditional SVM
#       fs_n:   number of univariate features selected for classification
#       dims:   classification performed on dims dimensions
#       dims_tg:classification generalized on dims_tg dimensions
#
# Ouputs:
#       predict: prediction matrix (split x samples x dims x dimsg)
#       predictg:same as predict for generalized samples
#       probas:  probas matrix (split x samples x dims x dimsg x class)
#       probasg: same as probas for generalized samples
#       coef:    weight hyperplan vector
#       all_folds:folding report (split x fold x samples)
#       y_all:   original y
#       y:       training y
#       yg:      generalized y
#       filenameX:
#       filenamey:
#
# Results are reported in: path + nameX + '_' + namey + "_results.mat"
###############################################################################
# (c) Jean-Remi King: jeanremi.king [at] gmail [dot] com
###############################################################################

# update    2013 01 03: input binary format
# update    2012 12 20: remove np.copy, add compute distance
# update    2012 11 29: fix 3rd dimension issue
# update    2012 11 13: fix bug str output on some python versions
# update    2012 11 02: change stratified kfolding y by y2
# update    2012 11 02: add np.copy to Xtrain and Xtest
# update    2012 11 01: correct feature selection coef bug when at 100 %
# update    2012 10 23: correct leaveoneout bug
# update    2012 10 23: correct major n_split new_order error
# update    2012 10 18: correct python/matlab dim incompatibility
# update    2012 10 18: correct error fs between 99 and 100 && remove Kbest
# update    2012 10 17: correct error n_features shape and add nice
# update    2012 10 01: correct prediction error+change loading results option
# update    2012 09 14: handle fs float error
# update    2012 09 14: pass n_cores to sys.arg
# version   2012 09 13: implementation of parallelization

###############################################################################
print("LIBRARY")
import sys as sys
import numpy as np
from scipy import stats
from sklearn import svm
from sklearn.cross_validation import StratifiedKFold, LeaveOneOut, KFold
from sklearn.feature_selection import SelectPercentile, SelectKBest, f_classif
from sklearn.externals.joblib import Parallel, delayed
import scipy.io as sio
from sklearn.preprocessing import Scaler
from sklearn.grid_search import GridSearchCV
from sklearn.metrics import precision_score

###############################################################################
print("INPUT DATA")
#-- get argument to load specific file
filenameX = str(sys.argv[1])
filenamey = str(sys.argv[2])
if len(sys.argv) <= 3:
    n_cores = -1
else:
    n_cores = int(sys.argv[3])

print("cores: " + str(n_cores))
print(filenameX)
print(filenamey)

#-- Load data into python
if filenameX[-4:] == ".mat":  # backward compatibility
    mat = sio.loadmat(filenameX)
    Xm_all = mat["Xm"]  # data
    if np.size(Xm_all.shape) == 2:  # fix 3rd dimension issue
        X = np.zeros(np.append(Xm_all.shape, 1))
        X[:, :, 0] = Xm_all
        Xm_all = X
else:  # load binary
    mat = sio.loadmat(filenameX[0:-4] + "_dims.mat")
    Xdim = mat["Xdim"].reshape(3)
    Xm_all = np.fromfile(filenameX, dtype=np.float64)
    Xm_all = Xm_all.reshape(Xdim[2], Xdim[1], Xdim[0]).transpose([2, 1, 0])

#-- load classification parameters
mat = sio.loadmat(filenamey)
dims = mat["dims"]  # select time windows to compute
dims = np.reshape(dims, dims.size) - 1  # reshape for skl compatibility
dims_tg = mat["dims_tg"] - 1

mat = sio.loadmat(filenamey, squeeze_me=True)
path = mat["path"]
nameX = mat["nameX"]
namey = mat["namey"]
folding = mat["folding"]
n_splits = mat["n_splits"]
n_folds = mat["n_folds"]  # fold number
svm_C = mat["C"]  # svm penalization parameter
compute_probas = mat["compute_probas"]
compute_predict = mat["compute_predict"]
compute_distance = mat["compute_distance"]
fs_n = mat["fs"]  # feature selection
y_all = mat["y"]  # class used for train and test
print(Xm_all.shape)
print(y_all.shape)
y2_all = mat["y2"]  # class used for sample weights

#-- build training and generalizing classes
Xm = Xm_all[y_all > 0, :, :]  # training categories
Xmg = Xm_all[y_all < 0, :, :]  # generalization categories
y = y_all[y_all > 0]
yg = y_all[y_all < 0]
y2 = y2_all[y_all > 0]

n_samples, n_features, unused = Xm.shape
n_samplesg, unused, unused = Xmg.shape
n_featuresg = n_features
n_dims = dims.shape[0]
n_dimsg = n_dims

n_dims_tg = dims_tg.shape[1]
n_dimsg_tg = dims_tg.shape[1]

n_classes = np.unique(y).shape[0]
#deal with sample_weight
sample_weight = np.ones(y.shape[0])
classes = np.unique(y2)
for c in range(classes.shape[0]):
    sample_weight[y2 == classes[c]] = 1. / (np.sum(y2 == classes[c]))

###############################################################################
print("PREPARE CLASSIFICATION")

#-- classifier

clf = GridSearchCV(svm.SVC(kernel='linear', probability=True),
    {'C': svm_C}, score_func=precision_score)

#-- normalizer
scaler = Scaler()

#-- feature selection
if fs_n > 1:
    fs = SelectKBest(f_classif, k=fs_n)
elif fs_n == -1:
    fs = SelectKBest(f_classif, k=1)
else:
    fs = SelectPercentile(f_classif, percentile=fs_n * 100)

#-- results initialization
if compute_predict:
    predict = np.zeros([n_splits, n_samples, n_dims, n_dims_tg]) ** np.nan
    predictg = np.zeros([n_splits, n_samplesg, n_dimsg, n_dimsg_tg, n_folds]) ** np.nan
else:
    predict = []
    predictg = []

if compute_probas:
    probas = np.zeros([n_splits, n_samples, n_dims, n_dims_tg, n_classes]) ** np.nan
    probasg = np.zeros([n_splits, n_samplesg, n_dimsg, n_dimsg_tg, n_classes, n_folds]) ** np.nan
else:
    probas = []
    probasg = []

if compute_distance:
    distance = np.zeros([n_splits, n_samples, n_dims, n_dims_tg, n_classes]) ** np.nan
    distanceg = np.zeros([n_splits, n_samplesg, n_dimsg, n_dimsg_tg, n_classes, n_folds]) ** np.nan
else:
    distance = []
    distanceg = []

coef = np.empty([n_splits, n_folds, n_dims, n_classes * (n_classes - 1) / 2, n_features]) ** 0
all_folds = np.zeros([n_splits, n_folds, n_samples]) ** np.nan

###############################################################################
#-- Define parallel cross validation

def my_pipeline(train, test,
    Xm_shfl, y_shfl, sw_shfl, Xmg,
    dims, fs, scaler, clf,
    n_samples, n_dims, n_dims_tg, n_classes):
    # indicate opened fold
    sys.stdout.write("<")
    sys.stdout.flush()
    # initialize results within a given fold
    if compute_predict:
        predict = np.zeros([n_samples, n_dims, n_dims_tg]) ** np.nan
        predictg = np.zeros([n_samplesg, n_dimsg, n_dimsg_tg]) ** np.nan
    else:
        predict = []
        predictg = []
    if compute_probas:
        probas = np.zeros([n_samples, n_dims, n_dims_tg, n_classes]) ** np.nan
        probasg = np.zeros([n_samplesg, n_dimsg, n_dimsg_tg, n_classes]) ** np.nan
    else:
        probas = []
        probasg = []
    if compute_distance:
        distance = np.zeros([n_samples, n_dims, n_dims_tg, n_classes]) ** np.nan
        distanceg = np.zeros([n_samplesg, n_dimsg, n_dimsg_tg, n_classes]) ** np.nan
    else:
        distance = []
        distanceg = []
    coef = np.empty([n_dims, n_classes * (n_classes - 1) / 2, n_features]) ** 0
    # apply different classification along dimension 0
    for d in range(0, dims.shape[0]):
        Xtrain = Xm_shfl[train, :, dims[d]]
        ytrain = y_shfl[train]
        sw_train = sw_shfl[train]
        # (deal with NaN samples in training)
        ytrain = ytrain[~np.isnan(np.nansum(Xtrain, axis=1))]
        sw_train = sw_train[~np.isnan(np.nansum(Xtrain, axis=1))]
        Xtrain = Xtrain[~np.isnan(np.nansum(Xtrain, axis=1)), :]
        if np.unique(ytrain).shape[0] > 1:
            # feature selection
            fs.fit(Xtrain, ytrain)
            Xtrain = fs.transform(Xtrain)
            # normalization
            scaler.fit(Xtrain)
            Xtrain = scaler.transform(Xtrain)
            # Grid search
            clf.fit(Xtrain, ytrain, sample_weight=sw_train, cv=5)
            # select best classifier
            _clf = clf.best_estimator_
            # print(clf.best_params_)
            # retrieve features selected during univariate selection
            if fs_n == -1 or (fs_n > 1):
                # uni_features = sorted(range(len(fs.pvalues_)),key=lambda x:fs.pvalues_[x])
                uni_features = range(0, _clf.coef_.shape[1])
            else:
                uni_features = fs.pvalues_ <= stats.scoreatpercentile(fs.pvalues_, fs.percentile)
            # retrieve hyperplan (unselected features as 0)
            coef[d, :, uni_features] = scaler.inverse_transform(_clf.coef_).T
            # generalize across all time points
            for d_tg in range(0, n_dims_tg):
                # select data
                Xtest = Xm_shfl[test, :, dims_tg[d, d_tg]]
                # handles NaNs
                test_nan = np.isnan(np.nansum(Xtest, axis=1))
                Xtest = Xtest[~test_nan, :]
                # feature selection from training
                Xtest = fs.transform(Xtest)
                # normalize from training
                Xtest = scaler.transform(Xtest)
                # generalize test samples
                if (Xtest.shape[0] - np.sum(test_nan)) > 0:
                    if compute_predict:
                        predict[test[~test_nan], d, d_tg] = _clf.predict(Xtest)
                    if compute_probas:
                        probas[test[~test_nan], d, d_tg, :] = _clf.predict_proba(Xtest)
                    if compute_distance:
                        distance[test[~test_nan], d, d_tg, :] = _clf.decision_function(Xtest)  # correct!
                # predict on generalization sample
                # select data
                Xtestg = Xmg[:, :, dims_tg[d, d_tg]]
                # handles NaNs
                test_nan = np.isnan(np.nansum(Xtestg, axis=1))
                if (Xtestg.shape[0] - np.sum(test_nan)) > 0:
                    Xtestg = Xtestg[~test_nan, :]
                    # preproc feature selection and normalization
                    Xtestg = fs.transform(Xtestg)
                    Xtestg = scaler.transform(Xtestg)
                    # compute prediction
                    if compute_predict:
                        predictg[~test_nan, d, d_tg] = clf.predict(Xtestg)
                    if compute_probas:
                        probasg[~test_nan, d, d_tg, :] = clf.predict_proba(Xtestg)
                    if compute_distance:
                        distanceg[~test_nan, d, d_tg, :] = clf.decision_function(Xtestg)  # correct!
    # summarize fold results
    out = {
        'coef': coef,
        'predict': predict,
        'predictg': predictg,
        'probas': probas,
        'probasg': probasg,
        'distance': distance,
        'distanceg': distanceg}
    # indicate end of fold
    sys.stdout.write(">")
    sys.stdout.flush()
    return out

###############################################################################
print("CLASSIFY")
#-- Shuffle split
for split in range(n_splits):
    print("split " + str(split))
    #-- shuffle order in case this is not the first split
    new_order = np.array(range(y.shape[0]))
    if split > 0:
        np.random.shuffle(new_order)
        y_shfl = y
        y_shfl = y_shfl[new_order]
        y2_shfl = y2
        y2_shfl = y2_shfl[new_order]
        Xm_shfl = Xm
        Xm_shfl = Xm_shfl[new_order, :, :]
        sw_shfl = sample_weight
        sw_shfl = sw_shfl[new_order]
    else:
        y_shfl = y
        y2_shfl = y2
        Xm_shfl = Xm
        sw_shfl = sample_weight
    #-- define crossvalidation
    if folding == 'stratified':
        cv = StratifiedKFold(y2_shfl, k=n_folds)
    elif folding == 'kfolding':
        cv = KFold(n=y2_shfl.shape[0], k=n_folds)
    elif folding == 'leaveoneout':
        n_folds = y_shfl.shape[0]
        cv = LeaveOneOut(n=y_shfl.shape[0])
    else:
        print("unknown crossvalidation method!")
    # Cross-validation computed in parallel
    out = Parallel(n_jobs=n_cores)(delayed(my_pipeline)(
        train=train,
        test=test,
        Xm_shfl=Xm_shfl,
        y_shfl=y_shfl,
        sw_shfl=sw_shfl,
        Xmg=Xmg,
        dims=dims,
        fs=fs,
        scaler=scaler,
        clf=clf,
        n_samples=n_samples,
        n_dims=n_dims,
        n_dims_tg=n_dims_tg,
        n_classes=n_classes) for fold, (train, test) in enumerate(cv))
    # reorder results folds and splits
    for fold, (train, test) in enumerate(cv):
        all_folds[split, fold, train] = 1
        all_folds[split, fold, test] = 0
        coef[split, fold, :, :, :] = out[fold]['coef']
        if compute_predict:
            predict[split, new_order[test], :, :] = out[fold]['predict'][test, :, :]
            predictg[split, :, :, :, fold] = out[fold]['predictg']
        if compute_probas:
            probas[split, new_order[test], :, :, :] = out[fold]['probas'][test, :, :, :]
            probasg[split, :, :, :, :, fold] = out[fold]['probasg']
        if compute_distance:
            distance[split, new_order[test], :, :, :] = out[fold]['distance'][test, :, :, :]
            distanceg[split, :, :, :, :, fold] = out[fold]['distanceg']
    all_folds[split, :, new_order] = all_folds[split, :, :].T

###############################################################################
print("EXPORT DATA")
mat['predict'] = predict
mat['predictg'] = predictg
mat['probas'] = probas
mat['probasg'] = probasg
mat['distance'] = distance
mat['distanceg'] = distanceg
mat['coef'] = coef
mat['all_folds'] = all_folds
mat['y_all'] = y_all
mat['y'] = y
mat['yg'] = yg
mat['filenameX'] = filenameX
mat['filenamey'] = filenamey

print nameX
print namey
print path

output = str(path) + str(nameX) + '_' + str(namey) + "_results.mat"

print(output)
sio.savemat(output, mat)