from __future__ import print_function, absolute_import, division
import sys
import inspect
import socket

import numpy as np
from sklearn.utils import check_random_state
from sklearn.model_selection import ParameterGrid
try:
    from hyperopt import (Trials, tpe, fmin, STATUS_OK, STATUS_RUNNING,
                          STATUS_FAIL)
except ImportError:
    # hyperopt is optional, but required for hyperopt_tpe()
    pass

from .search_space import EnumVariable
from .acquisition_functions import AcquisitionFunction
from .surrogate_models import (MaximumLikelihoodGaussianProcess,
                               GaussianProcessKernel)

DEFAULT_TIMEOUT = socket._GLOBAL_DEFAULT_TIMEOUT

class BaseStrategy(object):
    short_name = None

    def suggest(self, history, searchspace):
        """
        Parameters
        ----------
        history : list of 3-tuples
            History of past function evaluations. Each element in history
            should be a tuple `(params, score, status)`, where `params` is a
            dict mapping parameter names to values
        searchspace : SearchSpace
            Instance of search_space.SearchSpace
        random_state :i nteger or numpy.RandomState, optional
            The random seed for sampling. If an integer is given, it fixes the
            seed. Defaults to the global numpy random number generator.

        Returns
        -------
        new_params : dict
        """
        raise NotImplementedError()

    @staticmethod
    def is_repeated_suggestion(params, history):
        """
        Parameters
        ----------
        params : dict
            Trial param set
        history : list of 3-tuples
            History of past function evaluations. Each element in history
            should be a tuple `(params, score, status)`, where `params` is a
            dict mapping parameter names to values

        Returns
        -------
        is_repeated_suggestion : bool
        """
        if any(params == hparams and hstatus == 'SUCCEEDED'
               for hparams, hscore, hstatus in history):
            return True
        else:
            return False

class RandomSearch(BaseStrategy):
    short_name = 'random'

    def __init__(self, seed=None):
        self.seed = seed

    def suggest(self, history, searchspace):
        """Randomly suggest params from searchspace.
        """
        return searchspace.rvs(self.seed)

class HyperoptTPE(BaseStrategy):
    short_name = 'hyperopt_tpe'

    def __init__(self, seed=None, gamma=0.25, seeds=20):
        self.seed = seed
        self.gamma = gamma
        self.seeds = seeds

    def suggest(self, history, searchspace):
        """
        Suggest params to maximize an objective function based on the
        function evaluation history using a tree of Parzen estimators (TPE),
        as implemented in the hyperopt package.

        Use of this function requires that hyperopt be installed.
        """
        # This function is very odd, because as far as I can tell there's
        # no real documented API for any of the internals of hyperopt. Its
        # execution model is that hyperopt calls your objective function
        # (instead of merely providing you with suggested points, and then
        # you calling the function yourself), and its very tricky (for me)
        # to use the internal hyperopt data structures to get these predictions
        # out directly.

        # so they path we take in this function is to construct a synthetic
        # hyperopt.Trials database which from the `history`, and then call
        # hyoperopt.fmin with a dummy objective function that logs the value
        # used, and then return that value to our client.

        # The form of the hyperopt.Trials database isn't really documented in
        # the code -- most of this comes from reverse engineering it, by
        # running fmin() on a simple function and then inspecting the form of
        # the resulting trials object.
        if 'hyperopt' not in sys.modules:
            raise ImportError('No module named hyperopt')

        random = check_random_state(self.seed)
        hp_searchspace = searchspace.to_hyperopt()

        trials = Trials()
        for i, (params, scores, status) in enumerate(history):
            if status == 'SUCCEEDED':
                # we're doing maximization, hyperopt.fmin() does minimization,
                # so we need to swap the sign
                result = {'loss': -np.mean(scores), 'status': STATUS_OK}
            elif status == 'PENDING':
                result = {'status': STATUS_RUNNING}
            elif status == 'FAILED':
                result = {'status': STATUS_FAIL}
            else:
                raise RuntimeError('unrecognized status: %s' % status)

            # the vals key in the trials dict is basically just the params
            # dict, but enum variables (hyperopt hp.choice() nodes) are
            # different, because the index of the parameter is specified
            # in vals, not the parameter itself.

            vals = {}
            for var in searchspace:
                if isinstance(var, EnumVariable):
                    # get the index in the choices of the parameter, and use
                    # that.
                    matches = [
                        i for i, c in enumerate(var.choices)
                        if c == params[var.name]
                    ]
                    assert len(matches) == 1
                    vals[var.name] = matches
                else:
                    # the other big difference is that all of the param values
                    # are wrapped in length-1 lists.
                    vals[var.name] = [params[var.name]]

            trials.insert_trial_doc({
                'misc': {
                    'cmd': ('domain_attachment', 'FMinIter_Domain'),
                    'idxs': dict((k, [i]) for k in hp_searchspace.keys()),
                    'tid': i,
                    'vals': vals,
                    'workdir': None
                },
                'result': result,
                'tid': i,
                # bunch of fixed fields that hyperopt seems to require
                'owner': None,
                'spec': None,
                'state': 2,
                'book_time': None,
                'exp_key': None,
                'refresh_time': None,
                'version': 0
            })

        trials.refresh()
        chosen_params_container = []

        def suggest(*args, **kwargs):
            return tpe.suggest(*args,
                               **kwargs,
                               gamma=self.gamma,
                               n_startup_jobs=self.seeds)

        def mock_fn(x):
            # http://stackoverflow.com/a/3190783/1079728
            # to get around no nonlocal keywork in python2
            chosen_params_container.append(x)
            return 0

        fmin(fn=mock_fn,
             algo=tpe.suggest,
             space=hp_searchspace,
             trials=trials,
             max_evals=len(trials.trials) + 1,
             **self._hyperopt_fmin_random_kwarg(random))
        chosen_params = chosen_params_container[0]

        return chosen_params

    @staticmethod
    def _hyperopt_fmin_random_kwarg(random):
        if 'rstate' in inspect.getargspec(fmin).args:
            # 0.0.3-dev version uses this argument
            kwargs = {'rstate': random, 'allow_trials_fmin': False}
        elif 'rseed' in inspect.getargspec(fmin).args:
            # 0.0.2 version uses different argument
            kwargs = {'rseed': random.randint(2**32 - 1)}
        return kwargs

class Bayes(BaseStrategy):
    short_name = 'bayes'

    def __init__(self,
                 acquisition=None,
                 surrogate=None,
                 kernels=None,
                 seed=None,
                 seeds=1,
                 max_feval=5E4,
                 max_iter=1E5,
                 n_iter=50):
        self.seed = seed
        self.seeds = seeds
        self.max_feval = max_feval
        self.max_iter = max_iter
        self.n_iter = n_iter
        self.n_dims = None
        if surrogate is None:
            surrogate = 'gp'
        self.surrogate = surrogate
        if kernels is None:
            kernels = [{
                'name': 'GPy.kern.Matern52',
                'params': {
                    'ARD': True
                },
                'options': {
                    'independent': False
                }
            }]
        self.kernel_params = kernels
        if acquisition is None:
            acquisition = {'name': 'osprey', 'params': {}}
        self.acquisition_params = acquisition

    def _get_data(self, history, searchspace):
        X = []
        Y = []
        V = []
        ignore = []
        for param_dict, scores, status in history:
            # transform points into the GP domain. This invloves bringing
            # int and enum variables to floating point, etc.
            if status == 'FAILED':
                # not sure how to deal with these yet
                continue

            point = searchspace.point_to_unit(param_dict)
            if status == 'SUCCEEDED':
                X.append(point)
                Y.append(np.mean(scores))
                V.append(np.var(scores))

            elif status == 'PENDING':
                ignore.append(point)
            else:
                raise RuntimeError('unrecognized status: %s' % status)

        return (np.array(X).reshape(-1, self.n_dims),
                np.array(Y).reshape(-1, 1), np.array(V).reshape(-1, 1),
                np.array(ignore).reshape(-1, self.n_dims))

    def _from_unit(self, result, searchspace):

        # Note that GP only deals with float-valued variables, so we have
        # a transform step on either side, where int and enum valued variables
        # are transformed before calling gp, and then the result suggested by
        # GP needs to be reverse-transformed.
        out = {}
        for gpvalue, var in zip(result, searchspace):
            out[var.name] = var.point_from_unit(float(gpvalue))
        return out

    def _is_within(self, point, X, tol=1E-2):
        if True in (np.sqrt(((point - X)**2).sum(axis=0)) <= tol):
            return True
        return False

    def suggest(self, history, searchspace, max_tries=5):
        if len(history) < self.seeds:
            return RandomSearch().suggest(history, searchspace)

        self.n_dims = searchspace.n_dims

        X, Y, V, ignore = self._get_data(history, searchspace)
        # TODO make _create_kernel accept optional args.
        # Define and fit model
        if self.surrogate == 'gp':
            kernel = GaussianProcessKernel(self.kernel_params, self.n_dims)
            model = MaximumLikelihoodGaussianProcess(X=X,
                                                     Y=Y,
                                                     kernel=kernel.kernel,
                                                     max_feval=self.max_feval)
        else:
            raise NotImplementedError(
                'Surrogate model not recognised.  Please choose from: gp')
        model.fit()

        # Define acquisition function and get best candidate
        af = AcquisitionFunction(surrogate=model,
                                 acquisition_params=self.acquisition_params,
                                 n_dims=self.n_dims,
                                 n_iter=self.n_iter,
                                 max_iter=self.max_iter)
        suggestion = af.get_best_candidate()

        if suggestion in ignore or self._is_within(suggestion, X):
            return RandomSearch().suggest(history, searchspace)

        return self._from_unit(suggestion, searchspace)

class GridSearch(BaseStrategy):
    short_name = 'grid'

    def __init__(self):
        self.param_grid = None
        self.current = -1

    def suggest(self, history, searchspace):
        # Convert searchspace to param_grid
        if self.param_grid is None:
            if not all(isinstance(v, EnumVariable) for v in searchspace):
                raise RuntimeError(
                    "GridSearchStrategy is defined only for all-enum search space"
                )

            self.param_grid = ParameterGrid(
                dict((v.name, v.choices) for v in searchspace))

        # NOTE: there is no way of signaling end of parameters to be searched against
        # so user should pick correctly number of evaluations
        self.current += 1
        return self.param_grid[self.current % len(self.param_grid)]


import numpy as np
from scipy.spatial import Delaunay

points = np.random.rand(30, 2) # 30 points in 2-d
tri = Delaunay(points)

# Make a list of line segments:
# edge_points = [ ((x1_1, y1_1), (x2_1, y2_1)),
#                 ((x1_2, y1_2), (x2_2, y2_2)),
#                 ... ]
edge_points = []
edges = set()

def add_edge(i, j):
    """Add a line between the i-th and j-th points, if not in the list already"""
    if (i, j) in edges or (j, i) in edges:
        # already added
        return
    edges.add( (i, j) )
    edge_points.append(points[ [i, j] ])

# loop over triangles:
# ia, ib, ic = indices of corner points of the triangle
for ia, ib, ic in tri.vertices:
    add_edge(ia, ib)
    add_edge(ib, ic)
    add_edge(ic, ia)

# plot it: the LineCollection is just a (maybe) faster way to plot lots of
# lines at once
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

lines = LineCollection(edge_points)
plt.figure()
plt.title('Delaunay triangulation')
plt.gca().add_collection(lines)
plt.plot(points[:,0], points[:,1], 'o', hold=1)
plt.xlim(-1, 2)
plt.ylim(-1, 2)

# -- the same stuff for the convex hull

edges = set()
edge_points = []

for ia, ib in tri.convex_hull:
    add_edge(ia, ib)

lines = LineCollection(edge_points)
plt.figure()
plt.title('Convex hull')
plt.gca().add_collection(lines)
plt.plot(points[:,0], points[:,1], 'o', hold=1)
plt.xlim(-1, 2)
plt.ylim(-1, 2)
plt.show()


# Copyright 2012 the rootpy developers
# distributed under the terms of the GNU General Public License
from __future__ import absolute_import
from ...extern.six.moves import range
from ...extern.six import string_types

__all__ = [
    'plot_corrcoef_matrix',
    'corrcoef',
    'cov',
]

def plot_corrcoef_matrix(matrix, names=None,
                         cmap=None, cmap_text=None,
                         fontsize=12, grid=False,
                         axes=None):
    """
    This function will draw a lower-triangular correlation matrix

    Parameters
    ----------

    matrix : 2-dimensional numpy array/matrix
        A correlation coefficient matrix

    names : list of strings, optional (default=None)
        List of the parameter names corresponding to the rows in ``matrix``.

    cmap : matplotlib color map, optional (default=None)
        Color map used to color the matrix cells.

    cmap_text : matplotlib color map, optional (default=None)
        Color map used to color the cell value text. If None, then
        all values will be black.

    fontsize : int, optional (default=12)
        Font size of parameter name and correlation value text.

    grid : bool, optional (default=False)
        If True, then draw dashed grid lines around the matrix elements.

    axes : matplotlib Axes instance, optional (default=None)
        The axes to plot on. If None then use the global current axes.

    Notes
    -----

    NumPy and matplotlib are required

    Examples
    --------

    >>> matrix = corrcoef(data.T, weights=weights)
    >>> plot_corrcoef_matrix(matrix, names)

    """
    import numpy as np
    from matplotlib import pyplot as plt
    from matplotlib import cm

    if axes is None:
        axes = plt.gca()

    matrix = np.asarray(matrix)

    if matrix.ndim != 2:
        raise ValueError("matrix is not a 2-dimensional array or matrix")
    if matrix.shape[0] != matrix.shape[1]:
        raise ValueError("matrix is not square")
    if names is not None and len(names) != matrix.shape[0]:
        raise ValueError("the number of names does not match the number of "
                         "rows/columns in the matrix")

    # mask out the upper triangular matrix
    matrix[np.triu_indices(matrix.shape[0])] = np.nan

    if isinstance(cmap_text, string_types):
        cmap_text = cm.get_cmap(cmap_text, 201)
    if cmap is None:
        cmap = cm.get_cmap('jet', 201)
    elif isinstance(cmap, string_types):
        cmap = cm.get_cmap(cmap, 201)
    # make NaN pixels white
    cmap.set_bad('w')

    axes.imshow(matrix, interpolation='nearest',
                cmap=cmap, origin='upper',
                vmin=-1, vmax=1)

    axes.set_frame_on(False)
    plt.setp(axes.get_yticklabels(), visible=False)
    plt.setp(axes.get_yticklines(), visible=False)
    plt.setp(axes.get_xticklabels(), visible=False)
    plt.setp(axes.get_xticklines(), visible=False)

    if grid:
        # draw grid lines
        for slot in range(1, matrix.shape[0] - 1):
            # vertical
            axes.plot((slot - 0.5, slot - 0.5),
                      (slot - 0.5, matrix.shape[0] - 0.5), 'k:', linewidth=1)
            # horizontal
            axes.plot((-0.5, slot + 0.5),
                      (slot + 0.5, slot + 0.5), 'k:', linewidth=1)
        if names is not None:
            for slot in range(1, matrix.shape[0]):
                # diagonal
                axes.plot((slot - 0.5, slot + 1.5),
                          (slot - 0.5, slot - 2.5), 'k:', linewidth=1)

    # label cell values
    for row, col in zip(*np.tril_indices(matrix.shape[0], k=-1)):
        value = matrix[row][col]
        if cmap_text is not None:
            color = cmap_text((value + 1.) / 2.)
        else:
            color = 'black'
        axes.text(
            col, row,
            "{0:d}%".format(int(value * 100)),
            color=color,
            ha='center', va='center',
            fontsize=fontsize)

    if names is not None:
        # write parameter names
        for i, name in enumerate(names):
            axes.annotate(
                name, (i, i),
                rotation=45,
                ha='left', va='bottom',
                transform=axes.transData,
                fontsize=fontsize)

def cov(m, y=None, rowvar=1, bias=0, ddof=None, weights=None, repeat_weights=0):
    """
    Estimate a covariance matrix, given data.

    Covariance indicates the level to which two variables vary together.
    If we examine N-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`,
    then the covariance matrix element :math:`C_{ij}` is the covariance of
    :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance
    of :math:`x_i`.

    Parameters
    ----------
    m : array_like
        A 1-D or 2-D array containing multiple variables and observations.
        Each row of `m` represents a variable, and each column a single
        observation of all those variables. Also see `rowvar` below.
    y : array_like, optional
        An additional set of variables and observations. `y` has the same
        form as that of `m`.
    rowvar : int, optional
        If `rowvar` is non-zero (default), then each row represents a
        variable, with observations in the columns. Otherwise, the relationship
        is transposed: each column represents a variable, while the rows
        contain observations.
    bias : int, optional
        Default normalization is by ``(N - 1)``, where ``N`` is the number of
        observations given (unbiased estimate). If `bias` is 1, then
        normalization is by ``N``. These values can be overridden by using
        the keyword ``ddof`` in numpy versions >= 1.5.
    ddof : int, optional
        .. versionadded:: 1.5
        If not ``None`` normalization is by ``(N - ddof)``, where ``N`` is
        the number of observations; this overrides the value implied by
        ``bias``. The default value is ``None``.
    weights : array-like, optional
        A 1-D array of weights with a length equal to the number of
        observations.
    repeat_weights : int, optional
        The default treatment of weights in the weighted covariance is to first
        normalize them to unit sum and use the biased weighted covariance
        equation. If `repeat_weights` is 1 then the weights must represent an
        integer number of occurrences of each observation and both a biased and
        unbiased weighted covariance is defined because the total sample size
        can be determined.

    Returns
    -------
    out : ndarray
        The covariance matrix of the variables.

    See Also
    --------
    corrcoef : Normalized covariance matrix

    Examples
    --------
    Consider two variables, :math:`x_0` and :math:`x_1`, which
    correlate perfectly, but in opposite directions:

    >>> x = np.array([[0, 2], [1, 1], [2, 0]]).T
    >>> x
    array([[0, 1, 2],
           [2, 1, 0]])

    Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance
    matrix shows this clearly:

    >>> np.cov(x)
    array([[ 1., -1.],
           [-1.,  1.]])

    Note that element :math:`C_{0,1}`, which shows the correlation between
    :math:`x_0` and :math:`x_1`, is negative.

    Further, note how `x` and `y` are combined:

    >>> x = [-2.1, -1,  4.3]
    >>> y = [3,  1.1,  0.12]
    >>> X = np.vstack((x,y))
    >>> print np.cov(X)
    [[ 11.71        -4.286     ]
     [ -4.286        2.14413333]]
    >>> print np.cov(x, y)
    [[ 11.71        -4.286     ]
     [ -4.286        2.14413333]]
    >>> print np.cov(x)
    11.71

    """
    import numpy as np
    # Check inputs
    if ddof is not None and ddof != int(ddof):
        raise ValueError(
            "ddof must be integer")

    X = np.array(m, ndmin=2, dtype=float)
    if X.size == 0:
        # handle empty arrays
        return np.array(m)
    if X.shape[0] == 1:
        rowvar = 1
    if rowvar:
        axis = 0
        tup = (slice(None), np.newaxis)
    else:
        axis = 1
        tup = (np.newaxis, slice(None))

    if y is not None:
        y = np.array(y, copy=False, ndmin=2, dtype=float)
        X = np.concatenate((X, y), axis)

    if ddof is None:
        if bias == 0:
            ddof = 1
        else:
            ddof = 0

    if weights is not None:
        weights = np.array(weights, dtype=float)
        weights_sum = weights.sum()
        if weights_sum <= 0:
            raise ValueError(
                "sum of weights is non-positive")
        X -= np.average(X, axis=1-axis, weights=weights)[tup]

        if repeat_weights:
            # each weight represents a number of repetitions of an observation
            # the total sample size can be determined in this case and we have
            # both an unbiased and biased weighted covariance
            fact = weights_sum - ddof
        else:
            # normalize weights so they sum to unity
            weights /= weights_sum
            # unbiased weighted covariance is not defined if the weights are
            # not integral frequencies (repeat-type)
            fact = (1. - np.power(weights, 2).sum())
    else:
        weights = 1
        X -= X.mean(axis=1-axis)[tup]
        if rowvar:
            N = X.shape[1]
        else:
            N = X.shape[0]
        fact = float(N - ddof)

    if not rowvar:
        return (np.dot(weights * X.T, X.conj()) / fact).squeeze()
    else:
        return (np.dot(weights * X, X.T.conj()) / fact).squeeze()

def corrcoef(x, y=None, rowvar=1, bias=0, ddof=None, weights=None,
             repeat_weights=0):
    """
    Return correlation coefficients.

    Please refer to the documentation for `cov` for more detail.  The
    relationship between the correlation coefficient matrix, `P`, and the
    covariance matrix, `C`, is

    .. math:: P_{ij} = \\frac{ C_{ij} } { \\sqrt{ C_{ii} * C_{jj} } }

    The values of `P` are between -1 and 1, inclusive.

    Parameters
    ----------
    x : array_like
        A 1-D or 2-D array containing multiple variables and observations.
        Each row of `m` represents a variable, and each column a single
        observation of all those variables. Also see `rowvar` below.
    y : array_like, optional
        An additional set of variables and observations. `y` has the same
        shape as `m`.
    rowvar : int, optional
        If `rowvar` is non-zero (default), then each row represents a
        variable, with observations in the columns. Otherwise, the relationship
        is transposed: each column represents a variable, while the rows
        contain observations.
    bias : int, optional
        Default normalization is by ``(N - 1)``, where ``N`` is the number of
        observations (unbiased estimate). If `bias` is 1, then
        normalization is by ``N``. These values can be overridden by using
        the keyword ``ddof`` in numpy versions >= 1.5.
    ddof : {None, int}, optional
        .. versionadded:: 1.5
        If not ``None`` normalization is by ``(N - ddof)``, where ``N`` is
        the number of observations; this overrides the value implied by
        ``bias``. The default value is ``None``.
    weights : array-like, optional
        A 1-D array of weights with a length equal to the number of
        observations.
    repeat_weights : int, optional
        The default treatment of weights in the weighted covariance is to first
        normalize them to unit sum and use the biased weighted covariance
        equation. If `repeat_weights` is 1 then the weights must represent an
        integer number of occurrences of each observation and both a biased and
        unbiased weighted covariance is defined because the total sample size
        can be determined.

    Returns
    -------
    out : ndarray
        The correlation coefficient matrix of the variables.

    See Also
    --------
    cov : Covariance matrix

    """
    import numpy as np
    c = cov(x, y, rowvar, bias, ddof, weights, repeat_weights)
    if c.size == 0:
        # handle empty arrays
        return c
    try:
        d = np.diag(c)
    except ValueError:  # scalar covariance
        return 1
    return c / np.sqrt(np.multiply.outer(d, d))


"""
This script evaluates a embedding model in a semantic similarity perspective.
It uses the dataset of ASSIN sentence similarity shared task and the method
of Hartmann which achieved the best results in the competition.

ASSIN shared-task website:
http://propor2016.di.fc.ul.pt/?page_id=381

Paper of Hartmann can be found at:
http://www.linguamatica.com/index.php/linguamatica/article/download/v8n2-6/365
"""

from sklearn.linear_model import LinearRegression
from sentence_similarity.utils.assin_eval import read_xml, eval_similarity
from gensim.models import KeyedVectors
from xml.dom import minidom
from numpy import array
from os import path
import pickle
import argparse

DATA_DIR = 'sentence_similarity/data/'
TEST_DIR = path.join(DATA_DIR, 'assin-test-gold/')

def gensim_embedding_difference(data, field1, field2):
    """Calculate the similarity between the sum of all embeddings."""
    distances = []
    for pair in data:
        e1 = [i if i in embeddings else 'unk' for i in pair[field1]]
        e2 = [i if i in embeddings else 'unk' for i in pair[field2]]
        distances.append([embeddings.n_similarity(e1, e2)])
    return distances

def evaluate_testset(x, y, test):
    """Docstring."""
    l_reg = LinearRegression()
    l_reg.fit(x, y)
    test_predict = l_reg.predict(test)
    return test_predict

def write_xml(filename, pred):
    """Docstring."""
    with open(filename) as fp:
        xml = minidom.parse(fp)
    pairs = xml.getElementsByTagName('pair')
    for pair in pairs:
        pair.setAttribute('similarity', str(pred[pairs.index(pair)]))
    with open(filename, 'w') as fp:
        fp.write(xml.toxml())

if __name__ == '__main__':
    # Parser descriptors
    parser = argparse.ArgumentParser(
        description='''Sentence similarity evaluation for word embeddings in
        brazilian and european variants of Portuguese language. It is expected
        a word embedding model in text format.''')

    parser.add_argument('embedding',
                        type=str,
                        help='embedding model')

    parser.add_argument('lang',
                        choices=['br', 'eu'],
                        help='{br, eu} choose PT-BR or PT-EU testset')

    args = parser.parse_args()
    lang = args.lang
    emb = args.embedding

    # Loading embedding model
    embeddings = KeyedVectors.load_word2vec_format(emb,
                                                   binary=False,
                                                   unicode_errors="ignore")

    # Loading evaluation data and parsing it
    with open('%sassin-pt%s-train.pkl' % (DATA_DIR, lang), 'rb') as fp:
        data = pickle.load(fp)

    with open('%sassin-pt%s-test-gold.pkl' % (DATA_DIR, lang), 'rb') as fp:
        test = pickle.load(fp)

    # Getting features
    features = gensim_embedding_difference(data, 'tokens_t1', 'tokens_t2')
    features_test = gensim_embedding_difference(test, 'tokens_t1', 'tokens_t2')

    # Predicting similarities
    results = array([float(i['result']) for i in data])
    results_test = evaluate_testset(features, results, features_test)

    write_xml('%soutput.xml' % DATA_DIR, results_test)

    # Evaluating
    pairs_gold = read_xml('%sassin-pt%s-test.xml' % (TEST_DIR, lang), True)
    pairs_sys = read_xml('%soutput.xml' % DATA_DIR, True)
    eval_similarity(pairs_gold, pairs_sys)


# encoding: UTF-8

'''
本文件中包含的是CTA模块的回测引擎，回测引擎的API和CTA引擎一致，
可以使用和实盘相同的代码进行回测。
History
<id>            <author>        <description>
2017051200      hetajen         样例：策略回测和优化
'''
from __future__ import division

'''2017051200 Add by hetajen begin'''
import time
'''2017051200 Add by hetajen end'''
from datetime import datetime, timedelta
from collections import OrderedDict
from itertools import product
import multiprocessing
import pymongo

from ctaBase import *
from vtConstant import *
from vtGateway import VtOrderData, VtTradeData
from vtFunction import loadMongoSetting

########################################################################
class BacktestingEngine(object):
    """
    CTA回测引擎
    函数接口和策略引擎保持一样，
    从而实现同一套代码从回测到实盘。
    """

    TICK_MODE = 'tick'
    BAR_MODE = 'bar'

    #----------------------------------------------------------------------
    def __init__(self):
        """Constructor"""
        # 本地停止单编号计数
        self.stopOrderCount = 0
        # stopOrderID = STOPORDERPREFIX + str(stopOrderCount)

        # 本地停止单字典
        # key为stopOrderID，value为stopOrder对象
        self.stopOrderDict = {}             # 停止单撤销后不会从本字典中删除
        self.workingStopOrderDict = {}      # 停止单撤销后会从本字典中删除

        # 引擎类型为回测
        self.engineType = ENGINETYPE_BACKTESTING

        # 回测相关
        self.strategy = None        # 回测策略
        self.mode = self.BAR_MODE   # 回测模式，默认为K线

        self.startDate = ''
        self.initDays = 0
        self.endDate = ''

        self.slippage = 0           # 回测时假设的滑点
        self.rate = 0               # 回测时假设的佣金比例（适用于百分比佣金）
        self.size = 1               # 合约大小，默认为1
        self.priceTick = 0          # 价格最小变动

        self.dbClient = None        # 数据库客户端
        self.dbCursor = None        # 数据库指针

        #self.historyData = []       # 历史数据的列表，回测用
        self.initData = []          # 初始化用的数据
        #self.backtestingData = []   # 回测用的数据

        self.dbName = ''            # 回测数据库名
        self.symbol = ''            # 回测集合名

        self.dataStartDate = None       # 回测数据开始日期，datetime对象
        self.dataEndDate = None         # 回测数据结束日期，datetime对象
        self.strategyStartDate = None   # 策略启动日期（即前面的数据用于初始化），datetime对象

        self.limitOrderDict = OrderedDict()         # 限价单字典
        self.workingLimitOrderDict = OrderedDict()  # 活动限价单字典，用于进行撮合用
        self.limitOrderCount = 0                    # 限价单编号

        self.tradeCount = 0             # 成交编号
        self.tradeDict = OrderedDict()  # 成交字典

        self.logList = []               # 日志记录

        # 当前最新数据，用于模拟成交用
        self.tick = None
        self.bar = None
        self.dt = None      # 最新的时间

    #----------------------------------------------------------------------
    def setStartDate(self, startDate='20100416', initDays=10):
        """设置回测的启动日期"""
        self.startDate = startDate
        self.initDays = initDays

        self.dataStartDate = datetime.strptime(startDate, '%Y%m%d')

        initTimeDelta = timedelta(initDays)
        self.strategyStartDate = self.dataStartDate + initTimeDelta

    #----------------------------------------------------------------------
    def setEndDate(self, endDate=''):
        """设置回测的结束日期"""
        self.endDate = endDate
        if endDate:
            self.dataEndDate= datetime.strptime(endDate, '%Y%m%d')
            # 若不修改时间则会导致不包含dataEndDate当天数据
            self.dataEndDate.replace(hour=23, minute=59)

    #----------------------------------------------------------------------
    def setBacktestingMode(self, mode):
        """设置回测模式"""
        self.mode = mode

    #----------------------------------------------------------------------
    def setDatabase(self, dbName, symbol):
        """设置历史数据所用的数据库"""
        self.dbName = dbName
        self.symbol = symbol

    #----------------------------------------------------------------------
    def loadHistoryData(self):
        """载入历史数据"""
        host, port, logging = loadMongoSetting()

        self.dbClient = pymongo.MongoClient(host, port)
        collection = self.dbClient[self.dbName][self.symbol]

        self.output(u'开始载入数据')

        # 首先根据回测模式，确认要使用的数据类
        if self.mode == self.BAR_MODE:
            dataClass = CtaBarData
            func = self.newBar
        else:
            dataClass = CtaTickData
            func = self.newTick

        # 载入初始化需要用的数据
        flt = {'datetime':{'$gte':self.dataStartDate,
                           '$lt':self.strategyStartDate}}
        initCursor = collection.find(flt)

        # 将数据从查询指针中读取出，并生成列表
        self.initData = []              # 清空initData列表
        for d in initCursor:
            data = dataClass()
            data.__dict__ = d
            self.initData.append(data)

        # 载入回测数据
        if not self.dataEndDate:
            flt = {'datetime':{'$gte':self.strategyStartDate}}   # 数据过滤条件
        else:
            flt = {'datetime':{'$gte':self.strategyStartDate,
                               '$lte':self.dataEndDate}}
        self.dbCursor = collection.find(flt)

        self.output(u'载入完成，数据量：%s' %(initCursor.count() + self.dbCursor.count()))

    #----------------------------------------------------------------------
    def runBacktesting(self):
        """运行回测"""
        # 载入历史数据
        self.loadHistoryData()

        # 首先根据回测模式，确认要使用的数据类
        if self.mode == self.BAR_MODE:
            dataClass = CtaBarData
            func = self.newBar
        else:
            dataClass = CtaTickData
            func = self.newTick

        self.output(u'开始回测')

        self.strategy.inited = True
        self.strategy.onInit()
        self.output(u'策略初始化完成')

        self.strategy.trading = True
        self.strategy.onStart()
        self.output(u'策略启动完成')

        self.output(u'开始回放数据')

        for d in self.dbCursor:
            data = dataClass()
            data.__dict__ = d
            func(data)

        self.output(u'数据回放结束')

    #----------------------------------------------------------------------
    def newBar(self, bar):
        """新的K线"""
        self.bar = bar
        self.dt = bar.datetime
        self.crossLimitOrder()      # 先撮合限价单
        self.crossStopOrder()       # 再撮合停止单
        self.strategy.onBar(bar)    # 推送K线到策略中

    #----------------------------------------------------------------------
    def newTick(self, tick):
        """新的Tick"""
        self.tick = tick
        self.dt = tick.datetime
        self.crossLimitOrder()
        self.crossStopOrder()
        self.strategy.onTick(tick)

    #----------------------------------------------------------------------
    def initStrategy(self, strategyClass, setting=None):
        """
        初始化策略
        setting是策略的参数设置，如果使用类中写好的默认设置则可以不传该参数
        """
        self.strategy = strategyClass(self, setting)
        self.strategy.name = self.strategy.className

    #----------------------------------------------------------------------
    def sendOrder(self, vtSymbol, orderType, price, volume, strategy):
        """发单"""
        self.limitOrderCount += 1
        orderID = str(self.limitOrderCount)

        order = VtOrderData()
        order.vtSymbol = vtSymbol
        order.price = self.roundToPriceTick(price)
        order.totalVolume = volume
        order.status = STATUS_NOTTRADED     # 刚提交尚未成交
        order.orderID = orderID
        order.vtOrderID = orderID
        order.orderTime = str(self.dt)

        # CTA委托类型映射
        if orderType == CTAORDER_BUY:
            order.direction = DIRECTION_LONG
            order.offset = OFFSET_OPEN
        elif orderType == CTAORDER_SELL:
            order.direction = DIRECTION_SHORT
            order.offset = OFFSET_CLOSE
        elif orderType == CTAORDER_SHORT:
            order.direction = DIRECTION_SHORT
            order.offset = OFFSET_OPEN
        elif orderType == CTAORDER_COVER:
            order.direction = DIRECTION_LONG
            order.offset = OFFSET_CLOSE

        # 保存到限价单字典中
        self.workingLimitOrderDict[orderID] = order
        self.limitOrderDict[orderID] = order

        return orderID

    #----------------------------------------------------------------------
    def cancelOrder(self, vtOrderID):
        """撤单"""
        if vtOrderID in self.workingLimitOrderDict:
            order = self.workingLimitOrderDict[vtOrderID]
            order.status = STATUS_CANCELLED
            order.cancelTime = str(self.dt)
            del self.workingLimitOrderDict[vtOrderID]

    #----------------------------------------------------------------------
    def sendStopOrder(self, vtSymbol, orderType, price, volume, strategy):
        """发停止单（本地实现）"""
        self.stopOrderCount += 1
        stopOrderID = STOPORDERPREFIX + str(self.stopOrderCount)

        so = StopOrder()
        so.vtSymbol = vtSymbol
        so.price = self.roundToPriceTick(price)
        so.volume = volume
        so.strategy = strategy
        so.stopOrderID = stopOrderID
        so.status = STOPORDER_WAITING

        if orderType == CTAORDER_BUY:
            so.direction = DIRECTION_LONG
            so.offset = OFFSET_OPEN
        elif orderType == CTAORDER_SELL:
            so.direction = DIRECTION_SHORT
            so.offset = OFFSET_CLOSE
        elif orderType == CTAORDER_SHORT:
            so.direction = DIRECTION_SHORT
            so.offset = OFFSET_OPEN
        elif orderType == CTAORDER_COVER:
            so.direction = DIRECTION_LONG
            so.offset = OFFSET_CLOSE

        # 保存stopOrder对象到字典中
        self.stopOrderDict[stopOrderID] = so
        self.workingStopOrderDict[stopOrderID] = so

        return stopOrderID

    #----------------------------------------------------------------------
    def cancelStopOrder(self, stopOrderID):
        """撤销停止单"""
        # 检查停止单是否存在
        if stopOrderID in self.workingStopOrderDict:
            so = self.workingStopOrderDict[stopOrderID]
            so.status = STOPORDER_CANCELLED
            del self.workingStopOrderDict[stopOrderID]

    #----------------------------------------------------------------------
    def crossLimitOrder(self):
        """基于最新数据撮合限价单"""
        # 先确定会撮合成交的价格
        if self.mode == self.BAR_MODE:
            buyCrossPrice = self.bar.low        # 若买入方向限价单价格高于该价格，则会成交
            sellCrossPrice = self.bar.high      # 若卖出方向限价单价格低于该价格，则会成交
            buyBestCrossPrice = self.bar.open   # 在当前时间点前发出的买入委托可能的最优成交价
            sellBestCrossPrice = self.bar.open  # 在当前时间点前发出的卖出委托可能的最优成交价
        else:
            buyCrossPrice = self.tick.askPrice1
            sellCrossPrice = self.tick.bidPrice1
            buyBestCrossPrice = self.tick.askPrice1
            sellBestCrossPrice = self.tick.bidPrice1

        # 遍历限价单字典中的所有限价单
        for orderID, order in self.workingLimitOrderDict.items():
            # 判断是否会成交
            buyCross = (order.direction==DIRECTION_LONG and
                        order.price>=buyCrossPrice and
                        buyCrossPrice > 0)      # 国内的tick行情在涨停时askPrice1为0，此时买无法成交

            sellCross = (order.direction==DIRECTION_SHORT and
                         order.price<=sellCrossPrice and
                         sellCrossPrice > 0)    # 国内的tick行情在跌停时bidPrice1为0，此时卖无法成交

            # 如果发生了成交
            if buyCross or sellCross:
                # 推送成交数据
                self.tradeCount += 1            # 成交编号自增1
                tradeID = str(self.tradeCount)
                trade = VtTradeData()
                trade.vtSymbol = order.vtSymbol
                trade.tradeID = tradeID
                trade.vtTradeID = tradeID
                trade.orderID = order.orderID
                trade.vtOrderID = order.orderID
                trade.direction = order.direction
                trade.offset = order.offset

                # 以买入为例：
                # 1. 假设当根K线的OHLC分别为：100, 125, 90, 110
                # 2. 假设在上一根K线结束(也是当前K线开始)的时刻，策略发出的委托为限价105
                # 3. 则在实际中的成交价会是100而不是105，因为委托发出时市场的最优价格是100
                if buyCross:
                    trade.price = min(order.price, buyBestCrossPrice)
                    self.strategy.pos += order.totalVolume
                else:
                    trade.price = max(order.price, sellBestCrossPrice)
                    self.strategy.pos -= order.totalVolume

                trade.volume = order.totalVolume
                trade.tradeTime = str(self.dt)
                trade.dt = self.dt
                self.strategy.onTrade(trade)

                self.tradeDict[tradeID] = trade

                # 推送委托数据
                order.tradedVolume = order.totalVolume
                order.status = STATUS_ALLTRADED
                self.strategy.onOrder(order)

                # 从字典中删除该限价单
                del self.workingLimitOrderDict[orderID]

    #----------------------------------------------------------------------
    def crossStopOrder(self):
        """基于最新数据撮合停止单"""
        # 先确定会撮合成交的价格，这里和限价单规则相反
        if self.mode == self.BAR_MODE:
            buyCrossPrice = self.bar.high    # 若买入方向停止单价格低于该价格，则会成交
            sellCrossPrice = self.bar.low    # 若卖出方向限价单价格高于该价格，则会成交
            bestCrossPrice = self.bar.open   # 最优成交价，买入停止单不能低于，卖出停止单不能高于
        else:
            buyCrossPrice = self.tick.lastPrice
            sellCrossPrice = self.tick.lastPrice
            bestCrossPrice = self.tick.lastPrice

        # 遍历停止单字典中的所有停止单
        for stopOrderID, so in self.workingStopOrderDict.items():
            # 判断是否会成交
            buyCross = so.direction==DIRECTION_LONG and so.price<=buyCrossPrice
            sellCross = so.direction==DIRECTION_SHORT and so.price>=sellCrossPrice

            # 如果发生了成交
            if buyCross or sellCross:
                # 推送成交数据
                self.tradeCount += 1            # 成交编号自增1
                tradeID = str(self.tradeCount)
                trade = VtTradeData()
                trade.vtSymbol = so.vtSymbol
                trade.tradeID = tradeID
                trade.vtTradeID = tradeID

                if buyCross:
                    self.strategy.pos += so.volume
                    trade.price = max(bestCrossPrice, so.price)
                else:
                    self.strategy.pos -= so.volume
                    trade.price = min(bestCrossPrice, so.price)

                self.limitOrderCount += 1
                orderID = str(self.limitOrderCount)
                trade.orderID = orderID
                trade.vtOrderID = orderID

                trade.direction = so.direction
                trade.offset = so.offset
                trade.volume = so.volume
                trade.tradeTime = str(self.dt)
                trade.dt = self.dt
                self.strategy.onTrade(trade)

                self.tradeDict[tradeID] = trade

                # 推送委托数据
                so.status = STOPORDER_TRIGGERED

                order = VtOrderData()
                order.vtSymbol = so.vtSymbol
                order.symbol = so.vtSymbol
                order.orderID = orderID
                order.vtOrderID = orderID
                order.direction = so.direction
                order.offset = so.offset
                order.price = so.price
                order.totalVolume = so.volume
                order.tradedVolume = so.volume
                order.status = STATUS_ALLTRADED
                order.orderTime = trade.tradeTime
                self.strategy.onOrder(order)

                self.limitOrderDict[orderID] = order

                # 从字典中删除该限价单
                if stopOrderID in self.workingStopOrderDict:
                    del self.workingStopOrderDict[stopOrderID]

    #----------------------------------------------------------------------
    def insertData(self, dbName, collectionName, data):
        """考虑到回测中不允许向数据库插入数据，防止实盘交易中的一些代码出错"""
        pass

    #----------------------------------------------------------------------
    def loadBar(self, dbName, collectionName, startDate):
        """直接返回初始化数据列表中的Bar"""
        return self.initData

    #----------------------------------------------------------------------
    def loadTick(self, dbName, collectionName, startDate):
        """直接返回初始化数据列表中的Tick"""
        return self.initData

    #----------------------------------------------------------------------
    def writeCtaLog(self, content):
        """记录日志"""
        log = str(self.dt) + ' ' + content
        self.logList.append(log)

    #----------------------------------------------------------------------
    def output(self, content):
        """输出内容"""
        print str(datetime.now()) + "\t" + content

    #----------------------------------------------------------------------
    def calculateBacktestingResult(self):
        """
        计算回测结果
        """
        self.output(u'计算回测结果')

        # 首先基于回测后的成交记录，计算每笔交易的盈亏
        resultList = []             # 交易结果列表

        longTrade = []              # 未平仓的多头交易
        shortTrade = []             # 未平仓的空头交易

        tradeTimeList = []          # 每笔成交时间戳
        posList = [0]               # 每笔成交后的持仓情况

        for trade in self.tradeDict.values():
            # 多头交易
            if trade.direction == DIRECTION_LONG:
                # 如果尚无空头交易
                if not shortTrade:
                    longTrade.append(trade)
                # 当前多头交易为平空
                else:
                    while True:
                        entryTrade = shortTrade[0]
                        exitTrade = trade

                        # 清算开平仓交易
                        closedVolume = min(exitTrade.volume, entryTrade.volume)
                        result = TradingResult(entryTrade.price, entryTrade.dt,
                                               exitTrade.price, exitTrade.dt,
                                               -closedVolume, self.rate, self.slippage, self.size)
                        resultList.append(result)

                        posList.extend([-1,0])
                        tradeTimeList.extend([result.entryDt, result.exitDt])

                        # 计算未清算部分
                        entryTrade.volume -= closedVolume
                        exitTrade.volume -= closedVolume

                        # 如果开仓交易已经全部清算，则从列表中移除
                        if not entryTrade.volume:
                            shortTrade.pop(0)

                        # 如果平仓交易已经全部清算，则退出循环
                        if not exitTrade.volume:
                            break

                        # 如果平仓交易未全部清算，
                        if exitTrade.volume:
                            # 且开仓交易已经全部清算完，则平仓交易剩余的部分
                            # 等于新的反向开仓交易，添加到队列中
                            if not shortTrade:
                                longTrade.append(exitTrade)
                                break
                            # 如果开仓交易还有剩余，则进入下一轮循环
                            else:
                                pass

            # 空头交易
            else:
                # 如果尚无多头交易
                if not longTrade:
                    shortTrade.append(trade)
                # 当前空头交易为平多
                else:
                    while True:
                        entryTrade = longTrade[0]
                        exitTrade = trade

                        # 清算开平仓交易
                        closedVolume = min(exitTrade.volume, entryTrade.volume)
                        result = TradingResult(entryTrade.price, entryTrade.dt,
                                               exitTrade.price, exitTrade.dt,
                                               closedVolume, self.rate, self.slippage, self.size)
                        resultList.append(result)

                        posList.extend([1,0])
                        tradeTimeList.extend([result.entryDt, result.exitDt])

                        # 计算未清算部分
                        entryTrade.volume -= closedVolume
                        exitTrade.volume -= closedVolume

                        # 如果开仓交易已经全部清算，则从列表中移除
                        if not entryTrade.volume:
                            longTrade.pop(0)

                        # 如果平仓交易已经全部清算，则退出循环
                        if not exitTrade.volume:
                            break

                        # 如果平仓交易未全部清算，
                        if exitTrade.volume:
                            # 且开仓交易已经全部清算完，则平仓交易剩余的部分
                            # 等于新的反向开仓交易，添加到队列中
                            if not longTrade:
                                shortTrade.append(exitTrade)
                                break
                            # 如果开仓交易还有剩余，则进入下一轮循环
                            else:
                                pass

        # 检查是否有交易
        if not resultList:
            self.output(u'无交易结果')
            return {}

        # 然后基于每笔交易的结果，我们可以计算具体的盈亏曲线和最大回撤等
        capital = 0             # 资金
        maxCapital = 0          # 资金最高净值
        drawdown = 0            # 回撤

        totalResult = 0         # 总成交数量
        totalTurnover = 0       # 总成交金额（合约面值）
        totalCommission = 0     # 总手续费
        totalSlippage = 0       # 总滑点

        timeList = []           # 时间序列
        pnlList = []            # 每笔盈亏序列
        capitalList = []        # 盈亏汇总的时间序列
        drawdownList = []       # 回撤的时间序列

        winningResult = 0       # 盈利次数
        losingResult = 0        # 亏损次数
        totalWinning = 0        # 总盈利金额
        totalLosing = 0         # 总亏损金额

        for result in resultList:
            capital += result.pnl
            maxCapital = max(capital, maxCapital)
            drawdown = capital - maxCapital

            pnlList.append(result.pnl)
            timeList.append(result.exitDt)      # 交易的时间戳使用平仓时间
            capitalList.append(capital)
            drawdownList.append(drawdown)

            totalResult += 1
            totalTurnover += result.turnover
            totalCommission += result.commission
            totalSlippage += result.slippage

            if result.pnl >= 0:
                winningResult += 1
                totalWinning += result.pnl
            else:
                losingResult += 1
                totalLosing += result.pnl

        # 计算盈亏相关数据
        winningRate = winningResult/totalResult*100         # 胜率

        averageWinning = 0                                  # 这里把数据都初始化为0
        averageLosing = 0
        profitLossRatio = 0

        if winningResult:
            averageWinning = totalWinning/winningResult     # 平均每笔盈利
        if losingResult:
            averageLosing = totalLosing/losingResult        # 平均每笔亏损
        if averageLosing:
            profitLossRatio = -averageWinning/averageLosing # 盈亏比

        # 返回回测结果
        d = {}
        d['capital'] = capital
        d['maxCapital'] = maxCapital
        d['drawdown'] = drawdown
        d['totalResult'] = totalResult
        d['totalTurnover'] = totalTurnover
        d['totalCommission'] = totalCommission
        d['totalSlippage'] = totalSlippage
        d['timeList'] = timeList
        d['pnlList'] = pnlList
        d['capitalList'] = capitalList
        d['drawdownList'] = drawdownList
        d['winningRate'] = winningRate
        d['averageWinning'] = averageWinning
        d['averageLosing'] = averageLosing
        d['profitLossRatio'] = profitLossRatio
        d['posList'] = posList
        d['tradeTimeList'] = tradeTimeList

        return d

    #----------------------------------------------------------------------
    def showBacktestingResult(self):
        """显示回测结果"""
        d = self.calculateBacktestingResult()

        # 输出
        self.output('-' * 30)
        self.output(u'第一笔交易：\t%s' % d['timeList'][0])
        self.output(u'最后一笔交易：\t%s' % d['timeList'][-1])

        self.output(u'总交易次数：\t%s' % formatNumber(d['totalResult']))
        self.output(u'总盈亏：\t%s' % formatNumber(d['capital']))
        self.output(u'最大回撤: \t%s' % formatNumber(min(d['drawdownList'])))

        self.output(u'平均每笔盈利：\t%s' %formatNumber(d['capital']/d['totalResult']))
        self.output(u'平均每笔滑点：\t%s' %formatNumber(d['totalSlippage']/d['totalResult']))
        self.output(u'平均每笔佣金：\t%s' %formatNumber(d['totalCommission']/d['totalResult']))

        self.output(u'胜率\t\t%s%%' %formatNumber(d['winningRate']))
        self.output(u'盈利交易平均值\t%s' %formatNumber(d['averageWinning']))
        self.output(u'亏损交易平均值\t%s' %formatNumber(d['averageLosing']))
        self.output(u'盈亏比：\t%s' %formatNumber(d['profitLossRatio']))

        # 绘图
        import matplotlib.pyplot as plt
        import numpy as np

        try:
            import seaborn as sns       # 如果安装了seaborn则设置为白色风格
            sns.set_style('whitegrid')
        except ImportError:
            pass

        pCapital = plt.subplot(4, 1, 1)
        pCapital.set_ylabel("capital")
        pCapital.plot(d['capitalList'], color='r', lw=0.8)

        pDD = plt.subplot(4, 1, 2)
        pDD.set_ylabel("DD")
        pDD.bar(range(len(d['drawdownList'])), d['drawdownList'], color='g')

        pPnl = plt.subplot(4, 1, 3)
        pPnl.set_ylabel("pnl")
        pPnl.hist(d['pnlList'], bins=50, color='c')

        pPos = plt.subplot(4, 1, 4)
        pPos.set_ylabel("Position")
        if d['posList'][-1] == 0:
            del d['posList'][-1]
        tradeTimeIndex = [item.strftime("%m/%d %H:%M:%S") for item in d['tradeTimeList']]
        xindex = np.arange(0, len(tradeTimeIndex), np.int(len(tradeTimeIndex)/10))
        tradeTimeIndex = map(lambda i: tradeTimeIndex[i], xindex)
        pPos.plot(d['posList'], color='k', drawstyle='steps-pre')
        pPos.set_ylim(-1.2, 1.2)
        plt.sca(pPos)
        plt.tight_layout()
        plt.xticks(xindex, tradeTimeIndex, rotation=30)  # 旋转15

        plt.show()

    #----------------------------------------------------------------------
    def putStrategyEvent(self, name):
        """发送策略更新事件，回测中忽略"""
        pass

    #----------------------------------------------------------------------
    def setSlippage(self, slippage):
        """设置滑点点数"""
        self.slippage = slippage

    #----------------------------------------------------------------------
    def setSize(self, size):
        """设置合约大小"""
        self.size = size

    #----------------------------------------------------------------------
    def setRate(self, rate):
        """设置佣金比例"""
        self.rate = rate

    #----------------------------------------------------------------------
    def setPriceTick(self, priceTick):
        """设置价格最小变动"""
        self.priceTick = priceTick

    #----------------------------------------------------------------------
    def runOptimization(self, strategyClass, optimizationSetting):
        """优化参数"""
        # 获取优化设置
        settingList = optimizationSetting.generateSetting()
        targetName = optimizationSetting.optimizeTarget

        # 检查参数设置问题
        if not settingList or not targetName:
            self.output(u'优化设置有问题，请检查')

        # 遍历优化
        resultList = []
        for setting in settingList:
            self.clearBacktestingResult()
            self.output('-' * 30)
            self.output('setting: %s' %str(setting))
            self.initStrategy(strategyClass, setting)
            self.runBacktesting()
            d = self.calculateBacktestingResult()
            try:
                targetValue = d[targetName]
            except KeyError:
                targetValue = 0
            resultList.append(([str(setting)], targetValue))

        # 显示结果
        resultList.sort(reverse=True, key=lambda result:result[1])
        self.output('-' * 30)
        self.output(u'优化结果：')
        for result in resultList:
            self.output(u'%s: %s' %(result[0], result[1]))
        return result

    #----------------------------------------------------------------------
    def clearBacktestingResult(self):
        """清空之前回测的结果"""
        # 清空限价单相关
        self.limitOrderCount = 0
        self.limitOrderDict.clear()
        self.workingLimitOrderDict.clear()

        # 清空停止单相关
        self.stopOrderCount = 0
        self.stopOrderDict.clear()
        self.workingStopOrderDict.clear()

        # 清空成交相关
        self.tradeCount = 0
        self.tradeDict.clear()

    #----------------------------------------------------------------------
    def runParallelOptimization(self, strategyClass, optimizationSetting):
        """并行优化参数"""
        # 获取优化设置
        settingList = optimizationSetting.generateSetting()
        targetName = optimizationSetting.optimizeTarget

        # 检查参数设置问题
        if not settingList or not targetName:
            self.output(u'优化设置有问题，请检查')

        # 多进程优化，启动一个对应CPU核心数量的进程池
        pool = multiprocessing.Pool(multiprocessing.cpu_count())
        l = []

        for setting in settingList:
            l.append(pool.apply_async(optimize, (strategyClass, setting,
                                                 targetName, self.mode,
                                                 self.startDate, self.initDays, self.endDate,
                                                 self.slippage, self.rate, self.size,
                                                 self.dbName, self.symbol)))
        pool.close()
        pool.join()

        # 显示结果
        resultList = [res.get() for res in l]
        resultList.sort(reverse=True, key=lambda result:result[1])
        self.output('-' * 30)
        self.output(u'优化结果：')
        for result in resultList:
            self.output(u'%s: %s' %(result[0], result[1]))

    #----------------------------------------------------------------------
    def roundToPriceTick(self, price):
        """取整价格到合约最小价格变动"""
        if not self.priceTick:
            return price

        newPrice = round(price/self.priceTick, 0) * self.priceTick
        return newPrice

########################################################################
class TradingResult(object):
    """每笔交易的结果"""

    #----------------------------------------------------------------------
    def __init__(self, entryPrice, entryDt, exitPrice,
                 exitDt, volume, rate, slippage, size):
        """Constructor"""
        self.entryPrice = entryPrice    # 开仓价格
        self.exitPrice = exitPrice      # 平仓价格

        self.entryDt = entryDt          # 开仓时间datetime
        self.exitDt = exitDt            # 平仓时间

        self.volume = volume    # 交易数量（+/-代表方向）

        self.turnover = (self.entryPrice+self.exitPrice)*size*abs(volume)   # 成交金额
        self.commission = self.turnover*rate                                # 手续费成本
        self.slippage = slippage*2*size*abs(volume)                         # 滑点成本
        self.pnl = ((self.exitPrice - self.entryPrice) * volume * size
                    - self.commission - self.slippage)                      # 净盈亏

########################################################################
class OptimizationSetting(object):
    """优化设置"""

    #----------------------------------------------------------------------
    def __init__(self):
        """Constructor"""
        self.paramDict = OrderedDict()

        self.optimizeTarget = ''        # 优化目标字段

    #----------------------------------------------------------------------
    def addParameter(self, name, start, end=None, step=None):
        """增加优化参数"""
        if end is None and step is None:
            self.paramDict[name] = [start]
            return

        if end < start:
            print u'参数起始点必须不大于终止点'
            return

        if step <= 0:
            print u'参数布进必须大于0'
            return

        l = []
        param = start

        while param <= end:
            l.append(param)
            param += step

        self.paramDict[name] = l

    #----------------------------------------------------------------------
    def generateSetting(self):
        """生成优化参数组合"""
        # 参数名的列表
        nameList = self.paramDict.keys()
        paramList = self.paramDict.values()

        # 使用迭代工具生产参数对组合
        productList = list(product(*paramList))

        # 把参数对组合打包到一个个字典组成的列表中
        settingList = []
        for p in productList:
            d = dict(zip(nameList, p))
            settingList.append(d)

        return settingList

    #----------------------------------------------------------------------
    def setOptimizeTarget(self, target):
        """设置优化目标字段"""
        self.optimizeTarget = target

#----------------------------------------------------------------------
def formatNumber(n):
    """格式化数字到字符串"""
    rn = round(n, 2)        # 保留两位小数
    return format(rn, ',')  # 加上千分符

#----------------------------------------------------------------------
def optimize(strategyClass, setting, targetName,
             mode, startDate, initDays, endDate,
             slippage, rate, size,
             dbName, symbol):
    """多进程优化时跑在每个进程中运行的函数"""
    engine = BacktestingEngine()
    engine.setBacktestingMode(mode)
    engine.setStartDate(startDate, initDays)
    engine.setEndDate(endDate)
    engine.setSlippage(slippage)
    engine.setRate(rate)
    engine.setSize(size)
    engine.setDatabase(dbName, symbol)

    engine.initStrategy(strategyClass, setting)
    engine.runBacktesting()
    d = engine.calculateBacktestingResult()
    try:
        targetValue = d[targetName]
    except KeyError:
        targetValue = 0
    return (str(setting), targetValue)

'''2017051200 Modify by hetajen begin'''
from strategy.strategyAtrRsi import AtrRsiStrategy

def getEngine():
    engine = BacktestingEngine()
    engine.setBacktestingMode(engine.BAR_MODE)  # 引擎的回测模式为K线
    engine.setStartDate('20120101')  # 回测用的数据起始日期
    engine.setSlippage(0.2)  # 股指1跳
    engine.setRate(0.3 / 10000)  # 万0.3
    engine.setSize(300)  # 股指合约大小
    engine.setDatabase(MINUTE_DB_NAME, 'IF0000')
    return engine

def getParam(type=0):
    if type == 0:
        setting = {'atrLength': 11}
    else:
        setting = OptimizationSetting()  # 新建一个优化任务设置对象
        setting.setOptimizeTarget('capital')  # 设置优化排序的目标是策略净盈利
        setting.addParameter('atrLength', 12, 20, 2)  # 增加第一个优化参数atrLength，起始11，结束12，步进1
        setting.addParameter('atrMa', 20, 30, 5)  # 增加第二个优化参数atrMa，起始20，结束30，步进1
        setting.addParameter('rsiLength', 5)  # 增加一个固定数值的参数
    return setting

def xh_backTesting():
    engine = getEngine()
    setting = getParam()

    engine.initStrategy(AtrRsiStrategy, setting)
    engine.runBacktesting() # 开始跑回测
    engine.showBacktestingResult() # 显示回测结果

def xh_optimize():
    engine = getEngine()
    setting = getParam(1)

    engine.runOptimization(AtrRsiStrategy, setting) # 单进程优化。耗时：xxx秒
    #engine.runParallelOptimization(AtrRsiStrategy, setting) # 多进程优化。耗时：xx秒

if __name__ == '__main__':
    start = time.time()
    xh_backTesting()
    xh_optimize()
    print u'耗时：%s' % (time.time() - start)  # 性能测试
'''2017051200 Modify by hetajen end'''


"""Testing for bicluster metrics module"""

import numpy as np

from sklearn.utils.testing import assert_equal, assert_almost_equal

from sklearn.metrics.cluster.bicluster import _jaccard
from sklearn.metrics import consensus_score

def test_jaccard():
    a1 = np.array([True, True, False, False])
    a2 = np.array([True, True, True, True])
    a3 = np.array([False, True, True, False])
    a4 = np.array([False, False, True, True])

    assert_equal(_jaccard(a1, a1, a1, a1), 1)
    assert_equal(_jaccard(a1, a1, a2, a2), 0.25)
    assert_equal(_jaccard(a1, a1, a3, a3), 1.0 / 7)
    assert_equal(_jaccard(a1, a1, a4, a4), 0)

def test_consensus_score():
    a = [[True, True, False, False],
         [False, False, True, True]]
    b = a[::-1]

    assert_equal(consensus_score((a, a), (a, a)), 1)
    assert_equal(consensus_score((a, a), (b, b)), 1)
    assert_equal(consensus_score((a, b), (a, b)), 1)
    assert_equal(consensus_score((a, b), (b, a)), 1)

    assert_equal(consensus_score((a, a), (b, a)), 0)
    assert_equal(consensus_score((a, a), (a, b)), 0)
    assert_equal(consensus_score((b, b), (a, b)), 0)
    assert_equal(consensus_score((b, b), (b, a)), 0)

def test_consensus_score_issue2445():
    ''' Different number of biclusters in A and B'''
    a_rows = np.array([[True, True, False, False],
                       [False, False, True, True],
                       [False, False, False, True]])
    a_cols = np.array([[True, True, False, False],
                       [False, False, True, True],
                       [False, False, False, True]])
    idx = [0, 2]
    s = consensus_score((a_rows, a_cols), (a_rows[idx], a_cols[idx]))
    # B contains 2 of the 3 biclusters in A, so score should be 2/3
    assert_almost_equal(s, 2.0/3.0)


#!/usr/bin/env python
# encoding: utf-8
"""
Example of creating a block model using the blockmodel function in NX.  Data used is the Hartford, CT drug users network:

@article{,
	title = {Social Networks of Drug Users in {High-Risk} Sites: Finding the Connections},
	volume = {6},
	shorttitle = {Social Networks of Drug Users in {High-Risk} Sites},
	url = {http://dx.doi.org/10.1023/A:1015457400897},
	doi = {10.1023/A:1015457400897},
	number = {2},
	journal = {{AIDS} and Behavior},
	author = {Margaret R. Weeks and Scott Clair and Stephen P. Borgatti and Kim Radda and Jean J. Schensul},
	month = jun,
	year = {2002},
	pages = {193--206}
}

"""

__author__ = """\n""".join(['Drew Conway <drew.conway@nyu.edu>',
                            'Aric Hagberg <hagberg@lanl.gov>'])

from collections import defaultdict
import networkx as nx
import numpy
from scipy.cluster import hierarchy
from scipy.spatial import distance
import matplotlib.pyplot as plt

def create_hc(G):
    """Creates hierarchical cluster of graph G from distance matrix"""
    path_length=nx.all_pairs_shortest_path_length(G)
    distances=numpy.zeros((len(G),len(G)))
    for u,p in path_length.items():
        for v,d in p.items():
            distances[u][v]=d
    # Create hierarchical cluster
    Y=distance.squareform(distances)
    Z=hierarchy.complete(Y)  # Creates HC using farthest point linkage
    # This partition selection is arbitrary, for illustrive purposes
    membership=list(hierarchy.fcluster(Z,t=1.15))
    # Create collection of lists for blockmodel
    partition=defaultdict(list)
    for n,p in zip(list(range(len(G))),membership):
        partition[p].append(n)
    return list(partition.values())

if __name__ == '__main__':
    G=nx.read_edgelist("hartford_drug.edgelist")

    # Extract largest connected component into graph H
    H = next(nx.connected_component_subgraphs(G))
    # Makes life easier to have consecutively labeled integer nodes
    H=nx.convert_node_labels_to_integers(H)
    # Create parititions with hierarchical clustering
    partitions=create_hc(H)
    # Build blockmodel graph
    BM=nx.blockmodel(H,partitions)

    # Draw original graph
    pos=nx.spring_layout(H,iterations=100)
    fig=plt.figure(1,figsize=(6,10))
    ax=fig.add_subplot(211)
    nx.draw(H,pos,with_labels=False,node_size=10)
    plt.xlim(0,1)
    plt.ylim(0,1)

    # Draw block model with weighted edges and nodes sized by number of internal nodes
    node_size=[BM.node[x]['nnodes']*10 for x in BM.nodes()]
    edge_width=[(2*d['weight']) for (u,v,d) in BM.edges(data=True)]
    # Set positions to mean of positions of internal nodes from original graph
    posBM={}
    for n in BM:
        xy=numpy.array([pos[u] for u in BM.node[n]['graph']])
        posBM[n]=xy.mean(axis=0)
    ax=fig.add_subplot(212)
    nx.draw(BM,posBM,node_size=node_size,width=edge_width,with_labels=False)
    plt.xlim(0,1)
    plt.ylim(0,1)
    plt.axis('off')
    plt.savefig('hartford_drug_block_model.png')


"""
Provides a collection of utilities for comparing (image) results.

"""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six

import hashlib
import os
import shutil

import numpy as np

import matplotlib
from matplotlib.compat import subprocess
from matplotlib.testing.noseclasses import ImageComparisonFailure
from matplotlib import _png
from matplotlib import _get_cachedir
from matplotlib import cbook
from distutils import version

__all__ = ['compare_float', 'compare_images', 'comparable_formats']

def make_test_filename(fname, purpose):
    """
    Make a new filename by inserting `purpose` before the file's
    extension.
    """
    base, ext = os.path.splitext(fname)
    return '%s-%s%s' % (base, purpose, ext)

def compare_float(expected, actual, relTol=None, absTol=None):
    """
    Fail if the floating point values are not close enough, with
    the given message.

    You can specify a relative tolerance, absolute tolerance, or both.

    """
    if relTol is None and absTol is None:
        raise ValueError("You haven't specified a 'relTol' relative "
                         "tolerance or a 'absTol' absolute tolerance "
                         "function argument. You must specify one.")
    msg = ""

    if absTol is not None:
        absDiff = abs(expected - actual)
        if absTol < absDiff:
            template = ['',
                        'Expected: {expected}',
                        'Actual:   {actual}',
                        'Abs diff: {absDiff}',
                        'Abs tol:  {absTol}']
            msg += '\n  '.join([line.format(**locals()) for line in template])

    if relTol is not None:
        # The relative difference of the two values.  If the expected value is
        # zero, then return the absolute value of the difference.
        relDiff = abs(expected - actual)
        if expected:
            relDiff = relDiff / abs(expected)

        if relTol < relDiff:
            # The relative difference is a ratio, so it's always unit-less.
            template = ['',
                        'Expected: {expected}',
                        'Actual:   {actual}',
                        'Rel diff: {relDiff}',
                        'Rel tol:  {relTol}']
            msg += '\n  '.join([line.format(**locals()) for line in template])

    return msg or None

def get_cache_dir():
    cachedir = _get_cachedir()
    if cachedir is None:
        raise RuntimeError('Could not find a suitable configuration directory')
    cache_dir = os.path.join(cachedir, 'test_cache')
    if not os.path.exists(cache_dir):
        try:
            cbook.mkdirs(cache_dir)
        except IOError:
            return None
    if not os.access(cache_dir, os.W_OK):
        return None
    return cache_dir

def get_file_hash(path, block_size=2 ** 20):
    md5 = hashlib.md5()
    with open(path, 'rb') as fd:
        while True:
            data = fd.read(block_size)
            if not data:
                break
            md5.update(data)
    return md5.hexdigest()

def make_external_conversion_command(cmd):
    def convert(old, new):
        cmdline = cmd(old, new)
        pipe = subprocess.Popen(
            cmdline, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = pipe.communicate()
        errcode = pipe.wait()
        if not os.path.exists(new) or errcode:
            msg = "Conversion command failed:\n%s\n" % ' '.join(cmdline)
            if stdout:
                msg += "Standard output:\n%s\n" % stdout
            if stderr:
                msg += "Standard error:\n%s\n" % stderr
            raise IOError(msg)

    return convert

def _update_converter():
    gs, gs_v = matplotlib.checkdep_ghostscript()
    if gs_v is not None:
        cmd = lambda old, new: \
            [gs, '-q', '-sDEVICE=png16m', '-dNOPAUSE', '-dBATCH',
             '-sOutputFile=' + new, old]
        converter['pdf'] = make_external_conversion_command(cmd)
        converter['eps'] = make_external_conversion_command(cmd)

    if matplotlib.checkdep_inkscape() is not None:
        cmd = lambda old, new: \
            ['inkscape', '-z', old, '--export-png', new]
        converter['svg'] = make_external_conversion_command(cmd)

#: A dictionary that maps filename extensions to functions which
#: themselves map arguments `old` and `new` (filenames) to a list of strings.
#: The list can then be passed to Popen to convert files with that
#: extension to png format.
converter = {}
_update_converter()

def comparable_formats():
    """
    Returns the list of file formats that compare_images can compare
    on this system.

    """
    return ['png'] + list(six.iterkeys(converter))

def convert(filename, cache):
    """
    Convert the named file into a png file.  Returns the name of the
    created file.

    If *cache* is True, the result of the conversion is cached in
    `matplotlib._get_cachedir() + '/test_cache/'`.  The caching is based
    on a hash of the exact contents of the input file.  The is no limit
    on the size of the cache, so it may need to be manually cleared
    periodically.

    """
    base, extension = filename.rsplit('.', 1)
    if extension not in converter:
        raise ImageComparisonFailure(
            "Don't know how to convert %s files to png" % extension)
    newname = base + '_' + extension + '.png'
    if not os.path.exists(filename):
        raise IOError("'%s' does not exist" % filename)

    # Only convert the file if the destination doesn't already exist or
    # is out of date.
    if (not os.path.exists(newname) or
            os.stat(newname).st_mtime < os.stat(filename).st_mtime):
        if cache:
            cache_dir = get_cache_dir()
        else:
            cache_dir = None

        if cache_dir is not None:
            hash_value = get_file_hash(filename)
            new_ext = os.path.splitext(newname)[1]
            cached_file = os.path.join(cache_dir, hash_value + new_ext)
            if os.path.exists(cached_file):
                shutil.copyfile(cached_file, newname)
                return newname

        converter[extension](filename, newname)

        if cache_dir is not None:
            shutil.copyfile(newname, cached_file)

    return newname

#: Maps file extensions to a function which takes a filename as its
#: only argument to return a list suitable for execution with Popen.
#: The purpose of this is so that the result file (with the given
#: extension) can be verified with tools such as xmllint for svg.
verifiers = {}

# Turning this off, because it seems to cause multiprocessing issues
if matplotlib.checkdep_xmllint() and False:
    verifiers['svg'] = lambda filename: [
        'xmllint', '--valid', '--nowarning', '--noout', filename]

def verify(filename):
    """Verify the file through some sort of verification tool."""
    if not os.path.exists(filename):
        raise IOError("'%s' does not exist" % filename)
    base, extension = filename.rsplit('.', 1)
    verifier = verifiers.get(extension, None)
    if verifier is not None:
        cmd = verifier(filename)
        pipe = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = pipe.communicate()
        errcode = pipe.wait()
        if errcode != 0:
            msg = "File verification command failed:\n%s\n" % ' '.join(cmd)
            if stdout:
                msg += "Standard output:\n%s\n" % stdout
            if stderr:
                msg += "Standard error:\n%s\n" % stderr
            raise IOError(msg)

def crop_to_same(actual_path, actual_image, expected_path, expected_image):
    # clip the images to the same size -- this is useful only when
    # comparing eps to pdf
    if actual_path[-7:-4] == 'eps' and expected_path[-7:-4] == 'pdf':
        aw, ah = actual_image.shape
        ew, eh = expected_image.shape
        actual_image = actual_image[int(aw / 2 - ew / 2):int(
            aw / 2 + ew / 2), int(ah / 2 - eh / 2):int(ah / 2 + eh / 2)]
    return actual_image, expected_image

def calculate_rms(expectedImage, actualImage):
    "Calculate the per-pixel errors, then compute the root mean square error."
    num_values = np.prod(expectedImage.shape)
    abs_diff_image = abs(expectedImage - actualImage)

    # On Numpy 1.6, we can use bincount with minlength, which is much
    # faster than using histogram
    expected_version = version.LooseVersion("1.6")
    found_version = version.LooseVersion(np.__version__)
    if found_version >= expected_version:
        histogram = np.bincount(abs_diff_image.ravel(), minlength=256)
    else:
        histogram = np.histogram(abs_diff_image, bins=np.arange(257))[0]

    sum_of_squares = np.sum(histogram * np.arange(len(histogram)) ** 2)
    rms = np.sqrt(float(sum_of_squares) / num_values)

    return rms

def compare_images(expected, actual, tol, in_decorator=False):
    """
    Compare two "image" files checking differences within a tolerance.

    The two given filenames may point to files which are convertible to
    PNG via the `.converter` dictionary. The underlying RMS is calculated
    with the `.calculate_rms` function.

    Parameters
    ----------
    expected : str
        The filename of the expected image.
    actual :str
        The filename of the actual image.
    tol : float
        The tolerance (a color value difference, where 255 is the
        maximal difference).  The test fails if the average pixel
        difference is greater than this value.
    in_decorator : bool
        If called from image_comparison decorator, this should be
        True. (default=False)

    Example
    -------
    img1 = "./baseline/plot.png"
    img2 = "./output/plot.png"
    compare_images( img1, img2, 0.001 ):

    """
    if not os.path.exists(actual):
        msg = "Output image %s does not exist." % actual
        raise Exception(msg)

    if os.stat(actual).st_size == 0:
        msg = "Output image file %s is empty." % actual
        raise Exception(msg)

    verify(actual)

    # Convert the image to png
    extension = expected.split('.')[-1]

    if not os.path.exists(expected):
        raise IOError('Baseline image %r does not exist.' % expected)

    if extension != 'png':
        actual = convert(actual, False)
        expected = convert(expected, True)

    # open the image files and remove the alpha channel (if it exists)
    expectedImage = _png.read_png_int(expected)
    actualImage = _png.read_png_int(actual)
    expectedImage = expectedImage[:, :, :3]
    actualImage = actualImage[:, :, :3]

    actualImage, expectedImage = crop_to_same(
        actual, actualImage, expected, expectedImage)

    # convert to signed integers, so that the images can be subtracted without
    # overflow
    expectedImage = expectedImage.astype(np.int16)
    actualImage = actualImage.astype(np.int16)

    rms = calculate_rms(expectedImage, actualImage)

    diff_image = make_test_filename(actual, 'failed-diff')

    if rms <= tol:
        if os.path.exists(diff_image):
            os.unlink(diff_image)
        return None

    save_diff_image(expected, actual, diff_image)

    results = dict(rms=rms, expected=str(expected),
                   actual=str(actual), diff=str(diff_image), tol=tol)

    if not in_decorator:
        # Then the results should be a string suitable for stdout.
        template = ['Error: Image files did not match.',
                    'RMS Value: {rms}',
                    'Expected:  \n    {expected}',
                    'Actual:    \n    {actual}',
                    'Difference:\n    {diff}',
                    'Tolerance: \n    {tol}', ]
        results = '\n  '.join([line.format(**results) for line in template])
    return results

def save_diff_image(expected, actual, output):
    expectedImage = _png.read_png(expected)
    actualImage = _png.read_png(actual)
    actualImage, expectedImage = crop_to_same(
        actual, actualImage, expected, expectedImage)
    expectedImage = np.array(expectedImage).astype(np.float)
    actualImage = np.array(actualImage).astype(np.float)
    assert expectedImage.ndim == actualImage.ndim
    assert expectedImage.shape == actualImage.shape
    absDiffImage = abs(expectedImage - actualImage)

    # expand differences in luminance domain
    absDiffImage *= 255 * 10
    save_image_np = np.clip(absDiffImage, 0, 255).astype(np.uint8)
    height, width, depth = save_image_np.shape

    # The PDF renderer doesn't produce an alpha channel, but the
    # matplotlib PNG writer requires one, so expand the array
    if depth == 3:
        with_alpha = np.empty((height, width, 4), dtype=np.uint8)
        with_alpha[:, :, 0:3] = save_image_np
        save_image_np = with_alpha

    # Hard-code the alpha channel to fully solid
    save_image_np[:, :, 3] = 255

    _png.write_png(save_image_np.tostring(), width, height, output)


import os
import shutil
import tempfile
import warnings
import numpy
from pickle import loads
from pickle import dumps

from sklearn.datasets import get_data_home
from sklearn.datasets import clear_data_home
from sklearn.datasets import load_files
from sklearn.datasets import load_sample_images
from sklearn.datasets import load_sample_image
from sklearn.datasets import load_digits
from sklearn.datasets import load_diabetes
from sklearn.datasets import load_linnerud
from sklearn.datasets import load_iris
from sklearn.datasets import load_breast_cancer
from sklearn.datasets import load_boston
from sklearn.datasets import load_wine
from sklearn.datasets.base import Bunch

from sklearn.externals.six import b, u

from sklearn.utils.testing import assert_false
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import with_setup

DATA_HOME = tempfile.mkdtemp(prefix="scikit_learn_data_home_test_")
LOAD_FILES_ROOT = tempfile.mkdtemp(prefix="scikit_learn_load_files_test_")
TEST_CATEGORY_DIR1 = ""
TEST_CATEGORY_DIR2 = ""

def _remove_dir(path):
    if os.path.isdir(path):
        shutil.rmtree(path)

def teardown_module():
    """Test fixture (clean up) run once after all tests of this module"""
    for path in [DATA_HOME, LOAD_FILES_ROOT]:
        _remove_dir(path)

def setup_load_files():
    global TEST_CATEGORY_DIR1
    global TEST_CATEGORY_DIR2
    TEST_CATEGORY_DIR1 = tempfile.mkdtemp(dir=LOAD_FILES_ROOT)
    TEST_CATEGORY_DIR2 = tempfile.mkdtemp(dir=LOAD_FILES_ROOT)
    sample_file = tempfile.NamedTemporaryFile(dir=TEST_CATEGORY_DIR1,
                                              delete=False)
    sample_file.write(b("Hello World!\n"))
    sample_file.close()

def teardown_load_files():
    _remove_dir(TEST_CATEGORY_DIR1)
    _remove_dir(TEST_CATEGORY_DIR2)

def test_data_home():
    # get_data_home will point to a pre-existing folder
    data_home = get_data_home(data_home=DATA_HOME)
    assert_equal(data_home, DATA_HOME)
    assert_true(os.path.exists(data_home))

    # clear_data_home will delete both the content and the folder it-self
    clear_data_home(data_home=data_home)
    assert_false(os.path.exists(data_home))

    # if the folder is missing it will be created again
    data_home = get_data_home(data_home=DATA_HOME)
    assert_true(os.path.exists(data_home))

def test_default_empty_load_files():
    res = load_files(LOAD_FILES_ROOT)
    assert_equal(len(res.filenames), 0)
    assert_equal(len(res.target_names), 0)
    assert_equal(res.DESCR, None)

@with_setup(setup_load_files, teardown_load_files)
def test_default_load_files():
    res = load_files(LOAD_FILES_ROOT)
    assert_equal(len(res.filenames), 1)
    assert_equal(len(res.target_names), 2)
    assert_equal(res.DESCR, None)
    assert_equal(res.data, [b("Hello World!\n")])

@with_setup(setup_load_files, teardown_load_files)
def test_load_files_w_categories_desc_and_encoding():
    category = os.path.abspath(TEST_CATEGORY_DIR1).split('/').pop()
    res = load_files(LOAD_FILES_ROOT, description="test",
                     categories=category, encoding="utf-8")
    assert_equal(len(res.filenames), 1)
    assert_equal(len(res.target_names), 1)
    assert_equal(res.DESCR, "test")
    assert_equal(res.data, [u("Hello World!\n")])

@with_setup(setup_load_files, teardown_load_files)
def test_load_files_wo_load_content():
    res = load_files(LOAD_FILES_ROOT, load_content=False)
    assert_equal(len(res.filenames), 1)
    assert_equal(len(res.target_names), 2)
    assert_equal(res.DESCR, None)
    assert_equal(res.get('data'), None)

def test_load_sample_images():
    try:
        res = load_sample_images()
        assert_equal(len(res.images), 2)
        assert_equal(len(res.filenames), 2)
        assert_true(res.DESCR)
    except ImportError:
        warnings.warn("Could not load sample images, PIL is not available.")

def test_load_digits():
    digits = load_digits()
    assert_equal(digits.data.shape, (1797, 64))
    assert_equal(numpy.unique(digits.target).size, 10)

    # test return_X_y option
    X_y_tuple = load_digits(return_X_y=True)
    bunch = load_digits()
    assert_true(isinstance(X_y_tuple, tuple))
    assert_array_equal(X_y_tuple[0], bunch.data)
    assert_array_equal(X_y_tuple[1], bunch.target)

def test_load_digits_n_class_lt_10():
    digits = load_digits(9)
    assert_equal(digits.data.shape, (1617, 64))
    assert_equal(numpy.unique(digits.target).size, 9)

def test_load_sample_image():
    try:
        china = load_sample_image('china.jpg')
        assert_equal(china.dtype, 'uint8')
        assert_equal(china.shape, (427, 640, 3))
    except ImportError:
        warnings.warn("Could not load sample images, PIL is not available.")

def test_load_missing_sample_image_error():
    have_PIL = True
    try:
        try:
            from scipy.misc import imread
        except ImportError:
            from scipy.misc.pilutil import imread
    except ImportError:
        have_PIL = False
    if have_PIL:
        assert_raises(AttributeError, load_sample_image,
                      'blop.jpg')
    else:
        warnings.warn("Could not load sample images, PIL is not available.")

def test_load_diabetes():
    res = load_diabetes()
    assert_equal(res.data.shape, (442, 10))
    assert_true(res.target.size, 442)
    assert_equal(len(res.feature_names), 10)

    # test return_X_y option
    X_y_tuple = load_diabetes(return_X_y=True)
    bunch = load_diabetes()
    assert_true(isinstance(X_y_tuple, tuple))
    assert_array_equal(X_y_tuple[0], bunch.data)
    assert_array_equal(X_y_tuple[1], bunch.target)

def test_load_linnerud():
    res = load_linnerud()
    assert_equal(res.data.shape, (20, 3))
    assert_equal(res.target.shape, (20, 3))
    assert_equal(len(res.target_names), 3)
    assert_true(res.DESCR)

    # test return_X_y option
    X_y_tuple = load_linnerud(return_X_y=True)
    bunch = load_linnerud()
    assert_true(isinstance(X_y_tuple, tuple))
    assert_array_equal(X_y_tuple[0], bunch.data)
    assert_array_equal(X_y_tuple[1], bunch.target)

def test_load_iris():
    res = load_iris()
    assert_equal(res.data.shape, (150, 4))
    assert_equal(res.target.size, 150)
    assert_equal(res.target_names.size, 3)
    assert_true(res.DESCR)

    # test return_X_y option
    X_y_tuple = load_iris(return_X_y=True)
    bunch = load_iris()
    assert_true(isinstance(X_y_tuple, tuple))
    assert_array_equal(X_y_tuple[0], bunch.data)
    assert_array_equal(X_y_tuple[1], bunch.target)

def test_load_wine():
    res = load_wine()
    assert_equal(res.data.shape, (178, 13))
    assert_equal(res.target.size, 178)
    assert_equal(res.target_names.size, 3)
    assert_true(res.DESCR)

    # test return_X_y option
    X_y_tuple = load_wine(return_X_y=True)
    bunch = load_wine()
    assert_true(isinstance(X_y_tuple, tuple))
    assert_array_equal(X_y_tuple[0], bunch.data)
    assert_array_equal(X_y_tuple[1], bunch.target)

def test_load_breast_cancer():
    res = load_breast_cancer()
    assert_equal(res.data.shape, (569, 30))
    assert_equal(res.target.size, 569)
    assert_equal(res.target_names.size, 2)
    assert_true(res.DESCR)

    # test return_X_y option
    X_y_tuple = load_breast_cancer(return_X_y=True)
    bunch = load_breast_cancer()
    assert_true(isinstance(X_y_tuple, tuple))
    assert_array_equal(X_y_tuple[0], bunch.data)
    assert_array_equal(X_y_tuple[1], bunch.target)

def test_load_boston():
    res = load_boston()
    assert_equal(res.data.shape, (506, 13))
    assert_equal(res.target.size, 506)
    assert_equal(res.feature_names.size, 13)
    assert_true(res.DESCR)

    # test return_X_y option
    X_y_tuple = load_boston(return_X_y=True)
    bunch = load_boston()
    assert_true(isinstance(X_y_tuple, tuple))
    assert_array_equal(X_y_tuple[0], bunch.data)
    assert_array_equal(X_y_tuple[1], bunch.target)

def test_loads_dumps_bunch():
    bunch = Bunch(x="x")
    bunch_from_pkl = loads(dumps(bunch))
    bunch_from_pkl.x = "y"
    assert_equal(bunch_from_pkl['x'], bunch_from_pkl.x)

def test_bunch_pickle_generated_with_0_16_and_read_with_0_17():
    bunch = Bunch(key='original')
    # This reproduces a problem when Bunch pickles have been created
    # with scikit-learn 0.16 and are read with 0.17. Basically there
    # is a suprising behaviour because reading bunch.key uses
    # bunch.__dict__ (which is non empty for 0.16 Bunch objects)
    # whereas assigning into bunch.key uses bunch.__setattr__. See
    # https://github.com/scikit-learn/scikit-learn/issues/6196 for
    # more details
    bunch.__dict__['key'] = 'set from __dict__'
    bunch_from_pkl = loads(dumps(bunch))
    # After loading from pickle the __dict__ should have been ignored
    assert_equal(bunch_from_pkl.key, 'original')
    assert_equal(bunch_from_pkl['key'], 'original')
    # Making sure that changing the attr does change the value
    # associated with __getitem__ as well
    bunch_from_pkl.key = 'changed'
    assert_equal(bunch_from_pkl.key, 'changed')
    assert_equal(bunch_from_pkl['key'], 'changed')

def test_bunch_dir():
    # check that dir (important for autocomplete) shows attributes
    data = load_iris()
    assert_true("data" in dir(data))