#!/usr/bin/env python

'''
	File name: main_make_movie.py
	Author: Guillaume Viejo
	Date created: 09/10/2017
	Python Version: 3.5.2

To make shank mapping

'''

import numpy as np
import pandas as pd
# from matplotlib.pyplot import plot,show,draw
import scipy.io
from functions import *
from pylab import *
from sklearn.decomposition import PCA
import _pickle as cPickle
import neuroseries as nts
import sys

sys.exit()

###############################################################################################################
# LOADING DATA
###############################################################################################################
data_directory 	= '/mnt/DataGuillaume/MergedData/'
datasets 		= np.loadtxt(data_directory+'datasets_ThalHpc.list', delimiter = '\n', dtype = str, comments = '#')

theta_mod, theta_ses 	= loadThetaMod('/mnt/DataGuillaume/MergedData/THETA_THAL_mod.pickle', datasets, return_index=True)
swr_mod, swr_ses 		= loadSWRMod('/mnt/DataGuillaume/MergedData/SWR_THAL_corr.pickle', datasets, return_index=True)
spind_mod, spind_ses 	= loadSpindMod('/mnt/DataGuillaume/MergedData/SPINDLE_mod.pickle', datasets, return_index=True)
spike_spindle_phase 	= cPickle.load(open('/mnt/DataGuillaume/MergedData/SPIKE_SPINDLE_PHASE.pickle', 'rb'))
spike_theta_phase 		= cPickle.load(open('/mnt/DataGuillaume/MergedData/SPIKE_THETA_PHASE.pickle', 'rb'))

nbins 					= 400
binsize					= 5
times 					= np.arange(0, binsize*(nbins+1), binsize) - (nbins*binsize)/2

theta 					= pd.DataFrame(	index = theta_ses['rem'],
									columns = ['phase', 'pvalue', 'kappa'],
									data = theta_mod['rem'])

# filtering swr_mod
swr 					= pd.DataFrame(	columns = swr_ses,
										index = times,
										data = gaussFilt(swr_mod, (10,)).transpose())

# Cut swr_mod from -500 to 500
swr = swr.loc[-500:500]
# CHECK FOR NAN
tmp1 			= swr.columns[swr.isnull().any()].values
tmp2 			= theta.index[theta.isnull().any(1)].values
# CHECK P-VALUE
tmp3	 		= theta.index[(theta['pvalue'] > 1).values].values
tmp 			= np.unique(np.concatenate([tmp1,tmp2,tmp3]))
# copy and delete
if len(tmp):
	swr_modth 	= swr.drop(tmp, axis = 1)
	theta_modth = theta.drop(tmp, axis = 0)

swr_modth_copy 	= swr_modth.copy()
neuron_index = swr_modth.columns
times = swr_modth.loc[-500:500].index.values

###############################################################################################################
# MOVIE + jPCA for each animal
###############################################################################################################
mouses 	= ['Mouse12', 'Mouse17', 'Mouse20', 'Mouse32']
# times 	= np.arange(0, 1005, 5) - 500 # BAD

interval_to_cut = {	'Mouse12':[89,128],
					'Mouse17':[84,123],
					'Mouse20':[92,131],
					'Mouse32':[80,125]}

movies 		= dict.fromkeys(mouses)
rXX 		= dict.fromkeys(mouses)
maps 		= dict.fromkeys(mouses)
headdir 	= dict.fromkeys(mouses)
adnloc 		= dict.fromkeys(mouses)
xpos 		= dict.fromkeys(mouses)
ypos 		= dict.fromkeys(mouses)
xpos_shank 	= dict.fromkeys(mouses)
ypos_shank 	= dict.fromkeys(mouses)
xpos_phase 	= dict.fromkeys(mouses)
ypos_phase 	= dict.fromkeys(mouses)
theta_dens	= dict.fromkeys(mouses)
hd_neurons_index = []

for m in mouses:
	print(m)
	depth = pd.DataFrame(index = np.genfromtxt(data_directory+m+"/"+m+".depth", dtype = 'str', usecols = 0),
						data = np.genfromtxt(data_directory+m+"/"+m+".depth", usecols = 1),
						columns = ['depth'])
	neurons 		= np.array([n for n in neuron_index if m in n])
	sessions 		= np.unique([n.split("_")[0] for n in neuron_index if m in n])
	nb_bins 		= 201
	swr_shank 		= np.zeros((len(sessions),8,nb_bins))
	# nb_bins			= interval_to_cut[m][1] - interval_to_cut[m][0]
	theta_shank 	= np.zeros((len(sessions),8,30)) # that's radian bins here
	spindle_shank 	= np.zeros((len(sessions),8,30)) # that's radian bins here
	bins_phase 		= np.linspace(0.0, 2*np.pi+0.00001, 31)
	count_total 	= np.zeros((len(sessions),8))
	hd_neurons 		= np.zeros((len(sessions),8))
	amplitute		= np.zeros((len(sessions),8))
	mod_theta 		= np.zeros((len(sessions),8))
###########################################################################################################
# JPCA
###########################################################################################################
	rX,phi_swr,dynamical_system = jPCA(swr_modth[neurons].values.transpose(), times)
	phi_swr 		= pd.DataFrame(index = neurons, data = phi_swr)
###########################################################################################################
# VARIOUS
###########################################################################################################
	for s in sessions:
		generalinfo 		= scipy.io.loadmat(data_directory+m+"/"+s+'/Analysis/GeneralInfo.mat')
		shankStructure 		= loadShankStructure(generalinfo)
		spikes,shank		= loadSpikeData(data_directory+m+"/"+s+'/Analysis/SpikeData.mat', shankStructure['thalamus'])
		hd_info 			= scipy.io.loadmat(data_directory+m+'/'+s+'/Analysis/HDCells.mat')['hdCellStats'][:,-1]
		hd_info_neuron		= np.array([hd_info[n] for n in spikes.keys()])
		shankIndex 			= np.array([shank[n] for n in spikes.keys()]).flatten()
		if np.max(shankIndex) > 8 : sys.exit("Invalid shank index for thalamus" + s)
		shank_to_neurons 	= {k:np.array(list(spikes.keys()))[shankIndex == k] for k in np.unique(shankIndex)}

		for k in shank_to_neurons.keys():
			count_total[np.where(sessions== s)[0][0],k] = len(shank_to_neurons[k])
			hd_neurons[np.where(sessions== s)[0][0],k] = np.sum(hd_info_neuron[shankIndex == k])
			mod_theta[np.where(sessions== s)[0][0],k] = (theta.loc[[s+'_'+str(i) for i in shank_to_neurons[k]]]['pvalue'] < 0.05).sum()
			# amplitute[np.where(sessions==s)[0][0],k] = (swr.loc[shank_to_neurons[k]].var(1)).mean()
###########################################################################################################
# SWR MOD
###########################################################################################################
		neurons_mod_in_s 	= np.array([n for n in neurons if s in n])
		shank_to_neurons 	= {k:np.array([n for n in neurons_mod_in_s if shankIndex[int(n.split("_")[1])] == k]) for k in np.unique(shankIndex)}

		for k in shank_to_neurons.keys():
			# if np.sum(hd_info_neuron[[int(n.split("_")[1]) for n in shank_to_neurons[k]]]):
			# print(s, k, len(shank_to_neurons[k]))
			# if s == 'Mouse17-130204': sys.exit()
			if len(shank_to_neurons[k]):
				swr_shank[np.where(sessions== s)[0][0],k] = swr_modth[shank_to_neurons[k]].mean(1).values

###########################################################################################################
# THETA MOD
###########################################################################################################
		for k in shank_to_neurons.keys():
			if len(shank_to_neurons[k]):
				for n in shank_to_neurons[k]:
					phi = spike_theta_phase['rem'][n]
					phi[phi<0.0] += 2*np.pi
					index = np.digitize(phi, bins_phase)-1
					for t in index:
					 	theta_shank[np.where(sessions == s)[0][0],k,t] += 1.0
###########################################################################################################
# SPIND HPC MOD
###########################################################################################################
		for k in shank_to_neurons.keys():
			if len(shank_to_neurons[k]):
				for n in shank_to_neurons[k]:
					if n in list(spike_spindle_phase.keys()):
						phi = spike_spindle_phase['hpc'][n]
						phi[phi<0.0] += 2*np.pi
						index = np.digitize(phi, bins_phase)-1
						for t in index:
						 	spindle_shank[np.where(sessions == s)[0][0],k,t] += 1.0

	for t in range(len(times)):
		swr_shank[:,:,t] = np.flip(swr_shank[:,:,t], 1)
	for t in range(theta_shank.shape[-1]):
		theta_shank[:,:,t] = np.flip(theta_shank[:,:,t], 1)
		spindle_shank[:,:,t] = np.flip(spindle_shank[:,:,t], 1)

	# saving
	movies[m] = {	'swr'	:	swr_shank		,
					'theta'	:	theta_shank		,
					'spindle':	spindle_shank	}

	hd_neurons	= hd_neurons/(count_total+1.0)
	mod_theta 	= mod_theta/(count_total+1.0)
	rXX[m]		= rX
	maps[m] 	= {	'total':		np.flip(count_total,1),
					'x'			: 	np.arange(0.0, 8*0.2, 0.2),
					'y'			: 	depth.loc[sessions].values.flatten()
					}
	headdir[m] 	= np.flip(hd_neurons, 1)
	theta_dens[m] = np.flip(mod_theta, 1)

for m in movies.keys():
	datatosave = {	'movies':movies[m],
					'total':maps[m]['total'],
					'x':maps[m]['x'],
					'y':maps[m]['y'],
					'headdir':headdir[m],
					'jpc':rXX[m],
					'theta_dens':theta_dens[m]
					}

	cPickle.dump(datatosave, open("../data/maps/"+m+".pickle", 'wb'))

sys.exit()

m = 'Mouse12'
space = 0.01

thl_lines = np.load("../figures/thalamus_lines.mat.npy").sum(2)
xlines, ylines, thl_lines = interpolate(thl_lines, 	np.linspace(maps[m]['x'].min(), maps[m]['x'].max(), thl_lines.shape[1]),
 													np.linspace(maps[m]['y'].min(), maps[m]['y'].max(), thl_lines.shape[0]), 0.001)
thl_lines -= thl_lines.min()
thl_lines /= thl_lines.max()
thl_lines[thl_lines>0.6] = 1.0
thl_lines[thl_lines<=0.6] = 0.0
xnew, ynew, total = interpolate(maps[m]['total'].copy(), maps[m]['x'], maps[m]['y'], space)
# total -= total.min()
# total /= total.max()
total = softmax(total, 20.0, 0.2)

for k in movies[m].keys():
	movies[m][k] = filter_(movies[m][k], (2,2,5))

filmov = dict.fromkeys(movies[m].keys())
for k in filmov:
	tmp = []
	for t in range(movies[m][k].shape[-1]):
		# frame = movies[m][k][:,:,t] / (maps[m]['total']+1.0)
		frame = movies[m][k][:,:,t]
		xnew, ynew, frame = interpolate(frame, maps[m]['x'], maps[m]['y'], space)
		tmp.append(frame)
	tmp = np.array(tmp)
	filmov[k] = filter_(tmp, 5)
	filmov[k] = filmov[k] - np.min(filmov[k])
	filmov[k] = filmov[k] / np.max(filmov[k] + 1e-8)
	filmov[k] = softmax(filmov[k], 10, 0.5)

xnew, ynew, head = interpolate(headdir[m].copy(), maps[m]['x'], maps[m]['y'], space)
head[head < np.percentile(head, 90)] = 0.0

# sys.exit()

# figure()
# index = np.arange(0,20,1)+90

# for i in range(len(index)):
# 	subplot(4,5,i+1)
# 	# imshow(get_rgb(filmov['swr'][index[i]].copy(), total.copy(), np.ones_like(total), 0.83),
# 	imshow(filmov['swr'][index[i]].copy(),
# 			aspect = 'auto',
# 			origin = 'upper',
# 			cmap = 'jet', vmin = 0.0, vmax = 1.0)
# 			# extent = (xnew[0], xnew[-1], ynew[-1], ynew[0]))
# 	title("t = "+str(times[index[i]])+" ms")
# 	# contour(head, aspect = 'equal',origin = 'upper', extent = (xnew[0], xnew[-1], ynew[-1], ynew[0]))
# 	# contour(thl_lines, aspect = 'equal', origin = 'upper', extent = (xlines[0], xlines[-1], ylines[-1], ylines[0]), colors = 'white')
# 	# show(thl_lines, aspect = 'equal', origin = 'upper', extent = (xlines[0], xlines[-1], ylines[-1], ylines[0]))
# show()

from matplotlib import animation, rc
from IPython.display import HTML, Image

rc('animation', html='html5')
fig, axes = plt.subplots(1,1)
images = [axes.imshow(get_rgb(filmov['swr'][0].copy(), np.ones_like(total), total, 0.65), vmin = 0.0, vmax = 1.0, aspect = 'equal', origin = 'upper', extent = (xnew[0], xnew[-1], ynew[-1], ynew[0]))]
# images = [axes.imshow(filmov['swr'][0], aspect = 'equal', origin = 'upper', cmap = 'jet', vmin = 0.0, vmax = 1.0, extent = (xnew[0], xnew[-1], ynew[-1], ynew[0]))]
axes.contour(head, aspect = 'equal',origin = 'upper', extent = (xnew[0], xnew[-1], ynew[-1], ynew[0]), cmap = 'gist_gray')
axes.contour(thl_lines, aspect = 'equal', origin = 'upper', extent = (xlines[0], xlines[-1], ylines[-1], ylines[0]), colors = 'white')
def init():
	images[0].set_data(get_rgb(filmov['swr'][0].copy(), np.ones_like(total), total, 0.65))
	# images[0].set_data(filmov['swr'][0])
	return images

def animate(t):
	images[0].set_data(get_rgb(filmov['swr'][t].copy(), np.ones_like(total), total, 0.65))
	# images[0].set_data(filmov['swr'][t])
	return images

anim = animation.FuncAnimation(fig, animate, init_func=init,
						   frames=range(len(times)), interval=0, blit=False, repeat_delay = 5000)

anim.save('../figures/swr_mod_'+m+'.gif', writer='imagemagick', fps=60)
show()
sys.exit()

sys.exit()

from matplotlib import animation, rc
from IPython.display import HTML, Image

rc('animation', html='html5')
fig, axes = plt.subplots(1,3)
images = []
for k, i in zip(['swr', 'theta', 'spindle'], range(3)):
	images.append(axes[i].imshow(filmov[k][0], aspect = 'auto', cmap = 'jet', origin = 'upper', extent = (xnew[0], xnew[-1], ynew[-1], ynew[0])))
	contour(head, aspect = 'equal',origin = 'upper', extent = (xnew[0], xnew[-1], ynew[-1], ynew[0]))
def init():
	for i in range(3): images[i].set_data(filmov[k][0])
	contour(head, aspect = 'equal',origin = 'upper', extent = (xnew[0], xnew[-1], ynew[-1], ynew[0]))
	return images

def animate(t):
	for i in range(3): images[i].set_data(filmov[k][t])
	contour(head, aspect = 'equal',origin = 'upper', extent = (xnew[0], xnew[-1], ynew[-1], ynew[0]))
	return images

anim = animation.FuncAnimation(fig, animate, init_func=init,
						   frames=range(len(times)), interval=0, blit=True, repeat_delay = 0)

sys.exit()

m = 'Mouse12'

images = []
# for i in range(len(mouses)):
# 	lines1.append(axes[0,i].plot([],[],'o-')[0])
# 	lines2.append(axes[0,i].plot([],[],'o-')[0])
# 	axes[0,i].set_xlim(-500, 500)
# 	axes[0,i].set_ylim(rXX[mouses[i]].min(), rXX[mouses[i]].max())
images.append(axes.imshow(movies[m]['spindle'][:,:,0], aspect = 'auto', cmap = 'jet'))

def init():
	# for i, m in zip(range(len(mouses)), mouses):
	# 	images[i].set_data(movies[m][0])
	# 	lines1[i].set_data(times[0], rXX[m][0,0])
	# 	lines2[i].set_data(times[0], rXX[m][0,1])
	# return images+lines1+lines2
	images[0].set_data(movies[m]['spindle'][:,:,0])
	return images

def animate(t):
	# for i, m in zip(range(len(mouses)), mouses):
	# 	images[i].set_data(movies[m][t])
	# 	lines1[i].set_data(times[0:t], rXX[m][0:t,0])
	# 	lines2[i].set_data(times[0:t], rXX[m][0:t,1])
	images[0].set_data(movies[m]['spindle'][:,:,t])
	return images

anim = animation.FuncAnimation(fig, animate, init_func=init,
						   frames=movies[m]['spindle'].shape[-1], interval=0, blit=True, repeat_delay = 1)

show()

# anim.save('../figures/animation_swr_mod_jpca.gif', writer='imagemagick', fps=60)


from __future__ import division, print_function, absolute_import

from numpy.testing import dec, assert_, assert_array_equal

try:
    import matplotlib
    matplotlib.rcParams['backend'] = 'Agg'
    import matplotlib.pyplot as plt
    has_matplotlib = True
except:
    has_matplotlib = False

from scipy.spatial import \
     delaunay_plot_2d, voronoi_plot_2d, convex_hull_plot_2d, \
     Delaunay, Voronoi, ConvexHull

class TestPlotting:
    points = [(0,0), (0,1), (1,0), (1,1)]

    @dec.skipif(not has_matplotlib, "Matplotlib not available")
    def test_delaunay(self):
        # Smoke test
        fig = plt.figure()
        obj = Delaunay(self.points)
        s_before = obj.simplices.copy()
        r = delaunay_plot_2d(obj, ax=fig.gca())
        assert_array_equal(obj.simplices, s_before)  # shouldn't modify
        assert_(r is fig)
        delaunay_plot_2d(obj, ax=fig.gca())

    @dec.skipif(not has_matplotlib, "Matplotlib not available")
    def test_voronoi(self):
        # Smoke test
        fig = plt.figure()
        obj = Voronoi(self.points)
        r = voronoi_plot_2d(obj, ax=fig.gca())
        assert_(r is fig)
        voronoi_plot_2d(obj)

    @dec.skipif(not has_matplotlib, "Matplotlib not available")
    def test_convex_hull(self):
        # Smoke test
        fig = plt.figure()
        tri = ConvexHull(self.points)
        r = convex_hull_plot_2d(tri, ax=fig.gca())
        assert_(r is fig)
        convex_hull_plot_2d(tri)


# Author: Gael Varoquaux
# License: BSD

import numpy as np
import scipy.sparse as sp
from numpy.testing import assert_array_equal

from nose.tools import assert_true
from nose.tools import assert_false
from nose.tools import assert_equal
from nose.tools import assert_raises
from sklearn.base import BaseEstimator, clone, is_classifier
from sklearn.svm import SVC
from sklearn.pipeline import Pipeline
from sklearn.grid_search import GridSearchCV

#############################################################################
# A few test classes
class MyEstimator(BaseEstimator):

    def __init__(self, l1=0, empty=None):
        self.l1 = l1
        self.empty = empty

class K(BaseEstimator):
    def __init__(self, c=None, d=None):
        self.c = c
        self.d = d

class T(BaseEstimator):
    def __init__(self, a=None, b=None):
        self.a = a
        self.b = b

class Buggy(BaseEstimator):
    " A buggy estimator that does not set its parameters right. "

    def __init__(self, a=None):
        self.a = 1

#############################################################################
# The tests

def test_clone():
    """Tests that clone creates a correct deep copy.

    We create an estimator, make a copy of its original state
    (which, in this case, is the current state of the setimator),
    and check that the obtained copy is a correct deep copy.

    """
    from sklearn.feature_selection import SelectFpr, f_classif

    selector = SelectFpr(f_classif, alpha=0.1)
    new_selector = clone(selector)
    assert_true(selector is not new_selector)
    assert_equal(selector.get_params(), new_selector.get_params())

    selector = SelectFpr(f_classif, alpha=np.zeros((10, 2)))
    new_selector = clone(selector)
    assert_true(selector is not new_selector)

def test_clone_2():
    """Tests that clone doesn't copy everything.

    We first create an estimator, give it an own attribute, and
    make a copy of its original state. Then we check that the copy doesn't
    have the specific attribute we manually added to the initial estimator.
    """
    from sklearn.feature_selection import SelectFpr, f_classif

    selector = SelectFpr(f_classif, alpha=0.1)
    selector.own_attribute = "test"
    new_selector = clone(selector)
    assert_false(hasattr(new_selector, "own_attribute"))

def test_clone_buggy():
    """Check that clone raises an error on buggy estimators."""
    buggy = Buggy()
    buggy.a = 2
    assert_raises(RuntimeError, clone, buggy)

def test_clone_empty_array():
    """Regression test for cloning estimators with empty arrays"""
    clf = MyEstimator(empty=np.array([]))
    clf2 = clone(clf)
    assert_array_equal(clf.empty, clf2.empty)

    clf = MyEstimator(empty=sp.csr_matrix(np.array([[0]])))
    clf2 = clone(clf)
    assert_array_equal(clf.empty.data, clf2.empty.data)

def test_repr():
    """Smoke test the repr of the base estimator."""
    my_estimator = MyEstimator()
    repr(my_estimator)
    test = T(K(), K())
    assert_equal(
        repr(test),
        "T(a=K(c=None, d=None), b=K(c=None, d=None))"
    )

def test_str():
    """Smoke test the str of the base estimator"""
    my_estimator = MyEstimator()
    str(my_estimator)

def test_get_params():
    test = T(K(), K())

    assert_true('a__d' in test.get_params(deep=True))
    assert_true('a__d' not in test.get_params(deep=False))

    test.set_params(a__d=2)
    assert_true(test.a.d == 2)
    assert_raises(ValueError, test.set_params, a__a=2)

def test_is_classifier():
    svc = SVC()
    assert_true(is_classifier(svc))
    assert_true(is_classifier(GridSearchCV(svc, {'C': [0.1, 1]})))
    assert_true(is_classifier(Pipeline([('svc', svc)])))
    assert_true(is_classifier(Pipeline([('svc_cv',
                              GridSearchCV(svc, {'C': [0.1, 1]}))])))


"""Common tests for metaestimators"""

import functools

import numpy as np

from sklearn.base import BaseEstimator
from sklearn.externals.six import iterkeys
from sklearn.datasets import make_classification
from sklearn.utils.testing import assert_true, assert_false, assert_raises
from sklearn.pipeline import Pipeline
from sklearn.model_selection import GridSearchCV, RandomizedSearchCV
from sklearn.feature_selection import RFE, RFECV
from sklearn.ensemble import BaggingClassifier

class DelegatorData(object):
    def __init__(self, name, construct, skip_methods=(),
                 fit_args=make_classification()):
        self.name = name
        self.construct = construct
        self.fit_args = fit_args
        self.skip_methods = skip_methods

DELEGATING_METAESTIMATORS = [
    DelegatorData('Pipeline', lambda est: Pipeline([('est', est)])),
    DelegatorData('GridSearchCV',
                  lambda est: GridSearchCV(
                      est, param_grid={'param': [5]}, cv=2),
                  skip_methods=['score']),
    DelegatorData('RandomizedSearchCV',
                  lambda est: RandomizedSearchCV(
                      est, param_distributions={'param': [5]}, cv=2, n_iter=1),
                  skip_methods=['score']),
    DelegatorData('RFE', RFE,
                  skip_methods=['transform', 'inverse_transform', 'score']),
    DelegatorData('RFECV', RFECV,
                  skip_methods=['transform', 'inverse_transform', 'score']),
    DelegatorData('BaggingClassifier', BaggingClassifier,
                  skip_methods=['transform', 'inverse_transform', 'score',
                                'predict_proba', 'predict_log_proba', 'predict'])
]

def test_metaestimator_delegation():
    # Ensures specified metaestimators have methods iff subestimator does
    def hides(method):
        @property
        def wrapper(obj):
            if obj.hidden_method == method.__name__:
                raise AttributeError('%r is hidden' % obj.hidden_method)
            return functools.partial(method, obj)
        return wrapper

    class SubEstimator(BaseEstimator):
        def __init__(self, param=1, hidden_method=None):
            self.param = param
            self.hidden_method = hidden_method

        def fit(self, X, y=None, *args, **kwargs):
            self.coef_ = np.arange(X.shape[1])
            return True

        def _check_fit(self):
            if not hasattr(self, 'coef_'):
                raise RuntimeError('Estimator is not fit')

        @hides
        def inverse_transform(self, X, *args, **kwargs):
            self._check_fit()
            return X

        @hides
        def transform(self, X, *args, **kwargs):
            self._check_fit()
            return X

        @hides
        def predict(self, X, *args, **kwargs):
            self._check_fit()
            return np.ones(X.shape[0])

        @hides
        def predict_proba(self, X, *args, **kwargs):
            self._check_fit()
            return np.ones(X.shape[0])

        @hides
        def predict_log_proba(self, X, *args, **kwargs):
            self._check_fit()
            return np.ones(X.shape[0])

        @hides
        def decision_function(self, X, *args, **kwargs):
            self._check_fit()
            return np.ones(X.shape[0])

        @hides
        def score(self, X, *args, **kwargs):
            self._check_fit()
            return 1.0

    methods = [k for k in iterkeys(SubEstimator.__dict__)
               if not k.startswith('_') and not k.startswith('fit')]
    methods.sort()

    for delegator_data in DELEGATING_METAESTIMATORS:
        delegate = SubEstimator()
        delegator = delegator_data.construct(delegate)
        for method in methods:
            if method in delegator_data.skip_methods:
                continue
            assert_true(hasattr(delegate, method))
            assert_true(hasattr(delegator, method),
                        msg="%s does not have method %r when its delegate does"
                            % (delegator_data.name, method))
            # delegation before fit raises an exception
            assert_raises(Exception, getattr(delegator, method),
                          delegator_data.fit_args[0])

        delegator.fit(*delegator_data.fit_args)
        for method in methods:
            if method in delegator_data.skip_methods:
                continue
            # smoke test delegation
            getattr(delegator, method)(delegator_data.fit_args[0])

        for method in methods:
            if method in delegator_data.skip_methods:
                continue
            delegate = SubEstimator(hidden_method=method)
            delegator = delegator_data.construct(delegate)
            assert_false(hasattr(delegate, method))
            assert_false(hasattr(delegator, method),
                         msg="%s has method %r when its delegate does not"
                             % (delegator_data.name, method))


"""
======================================================
Imputing missing values before building an estimator
======================================================

This example shows that imputing the missing values can give better results
than discarding the samples containing any missing value.
Imputing does not always improve the predictions, so please check via cross-validation.
Sometimes dropping rows or using marker values is more effective.

Missing values can be replaced by the mean, the median or the most frequent
value using the ``strategy`` hyper-parameter.
The median is a more robust estimator for data with high magnitude variables
which could dominate results (otherwise known as a 'long tail').

Script output::

  Score with the entire dataset = 0.56
  Score without the samples containing missing values = 0.48
  Score after imputation of the missing values = 0.55

In this case, imputing helps the classifier get close to the original score.

"""
import numpy as np

from sklearn.datasets import load_boston
from sklearn.ensemble import RandomForestRegressor
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import Imputer
from sklearn.model_selection import cross_val_score

rng = np.random.RandomState(0)

dataset = load_boston()
X_full, y_full = dataset.data, dataset.target
n_samples = X_full.shape[0]
n_features = X_full.shape[1]

# Estimate the score on the entire dataset, with no missing values
estimator = RandomForestRegressor(random_state=0, n_estimators=100)
score = cross_val_score(estimator, X_full, y_full).mean()
print("Score with the entire dataset = %.2f" % score)

# Add missing values in 75% of the lines
missing_rate = 0.75
n_missing_samples = np.floor(n_samples * missing_rate)
missing_samples = np.hstack((np.zeros(n_samples - n_missing_samples,
                                      dtype=np.bool),
                             np.ones(n_missing_samples,
                                     dtype=np.bool)))
rng.shuffle(missing_samples)
missing_features = rng.randint(0, n_features, n_missing_samples)

# Estimate the score without the lines containing missing values
X_filtered = X_full[~missing_samples, :]
y_filtered = y_full[~missing_samples]
estimator = RandomForestRegressor(random_state=0, n_estimators=100)
score = cross_val_score(estimator, X_filtered, y_filtered).mean()
print("Score without the samples containing missing values = %.2f" % score)

# Estimate the score after imputation of the missing values
X_missing = X_full.copy()
X_missing[np.where(missing_samples)[0], missing_features] = 0
y_missing = y_full.copy()
estimator = Pipeline([("imputer", Imputer(missing_values=0,
                                          strategy="mean",
                                          axis=0)),
                      ("forest", RandomForestRegressor(random_state=0,
                                                       n_estimators=100))])
score = cross_val_score(estimator, X_missing, y_missing).mean()
print("Score after imputation of the missing values = %.2f" % score)


import numpy as np
import scipy.sparse as sp

from scipy import linalg
from numpy.testing import assert_array_almost_equal, assert_array_equal

from sklearn.datasets import make_classification
from sklearn.utils.sparsefuncs import (mean_variance_axis,
                                       inplace_column_scale,
                                       inplace_row_scale,
                                       inplace_swap_row, inplace_swap_column,
                                       min_max_axis,
                                       count_nonzero, csc_median_axis_0)
from sklearn.utils.sparsefuncs_fast import assign_rows_csr
from sklearn.utils.testing import assert_raises

def test_mean_variance_axis0():
    X, _ = make_classification(5, 4, random_state=0)
    # Sparsify the array a little bit
    X[0, 0] = 0
    X[2, 1] = 0
    X[4, 3] = 0
    X_lil = sp.lil_matrix(X)
    X_lil[1, 0] = 0
    X[1, 0] = 0
    X_csr = sp.csr_matrix(X_lil)

    X_means, X_vars = mean_variance_axis(X_csr, axis=0)
    assert_array_almost_equal(X_means, np.mean(X, axis=0))
    assert_array_almost_equal(X_vars, np.var(X, axis=0))

    X_csc = sp.csc_matrix(X_lil)
    X_means, X_vars = mean_variance_axis(X_csc, axis=0)

    assert_array_almost_equal(X_means, np.mean(X, axis=0))
    assert_array_almost_equal(X_vars, np.var(X, axis=0))
    assert_raises(TypeError, mean_variance_axis, X_lil, axis=0)

    X = X.astype(np.float32)
    X_csr = X_csr.astype(np.float32)
    X_csc = X_csr.astype(np.float32)
    X_means, X_vars = mean_variance_axis(X_csr, axis=0)
    assert_array_almost_equal(X_means, np.mean(X, axis=0))
    assert_array_almost_equal(X_vars, np.var(X, axis=0))
    X_means, X_vars = mean_variance_axis(X_csc, axis=0)
    assert_array_almost_equal(X_means, np.mean(X, axis=0))
    assert_array_almost_equal(X_vars, np.var(X, axis=0))
    assert_raises(TypeError, mean_variance_axis, X_lil, axis=0)

def test_mean_variance_illegal_axis():
    X, _ = make_classification(5, 4, random_state=0)
    # Sparsify the array a little bit
    X[0, 0] = 0
    X[2, 1] = 0
    X[4, 3] = 0
    X_csr = sp.csr_matrix(X)
    assert_raises(ValueError, mean_variance_axis, X_csr, axis=-3)
    assert_raises(ValueError, mean_variance_axis, X_csr, axis=2)
    assert_raises(ValueError, mean_variance_axis, X_csr, axis=-1)

def test_mean_variance_axis1():
    X, _ = make_classification(5, 4, random_state=0)
    # Sparsify the array a little bit
    X[0, 0] = 0
    X[2, 1] = 0
    X[4, 3] = 0
    X_lil = sp.lil_matrix(X)
    X_lil[1, 0] = 0
    X[1, 0] = 0
    X_csr = sp.csr_matrix(X_lil)

    X_means, X_vars = mean_variance_axis(X_csr, axis=1)
    assert_array_almost_equal(X_means, np.mean(X, axis=1))
    assert_array_almost_equal(X_vars, np.var(X, axis=1))

    X_csc = sp.csc_matrix(X_lil)
    X_means, X_vars = mean_variance_axis(X_csc, axis=1)

    assert_array_almost_equal(X_means, np.mean(X, axis=1))
    assert_array_almost_equal(X_vars, np.var(X, axis=1))
    assert_raises(TypeError, mean_variance_axis, X_lil, axis=1)

    X = X.astype(np.float32)
    X_csr = X_csr.astype(np.float32)
    X_csc = X_csr.astype(np.float32)
    X_means, X_vars = mean_variance_axis(X_csr, axis=1)
    assert_array_almost_equal(X_means, np.mean(X, axis=1))
    assert_array_almost_equal(X_vars, np.var(X, axis=1))
    X_means, X_vars = mean_variance_axis(X_csc, axis=1)
    assert_array_almost_equal(X_means, np.mean(X, axis=1))
    assert_array_almost_equal(X_vars, np.var(X, axis=1))
    assert_raises(TypeError, mean_variance_axis, X_lil, axis=1)

def test_densify_rows():
    X = sp.csr_matrix([[0, 3, 0],
                       [2, 4, 0],
                       [0, 0, 0],
                       [9, 8, 7],
                       [4, 0, 5]], dtype=np.float64)
    X_rows = np.array([0, 2, 3], dtype=np.intp)
    out = np.ones((6, X.shape[1]), dtype=np.float64)
    out_rows = np.array([1, 3, 4], dtype=np.intp)

    expect = np.ones_like(out)
    expect[out_rows] = X[X_rows, :].toarray()

    assign_rows_csr(X, X_rows, out_rows, out)
    assert_array_equal(out, expect)

def test_inplace_column_scale():
    rng = np.random.RandomState(0)
    X = sp.rand(100, 200, 0.05)
    Xr = X.tocsr()
    Xc = X.tocsc()
    XA = X.toarray()
    scale = rng.rand(200)
    XA *= scale

    inplace_column_scale(Xc, scale)
    inplace_column_scale(Xr, scale)
    assert_array_almost_equal(Xr.toarray(), Xc.toarray())
    assert_array_almost_equal(XA, Xc.toarray())
    assert_array_almost_equal(XA, Xr.toarray())
    assert_raises(TypeError, inplace_column_scale, X.tolil(), scale)

    X = X.astype(np.float32)
    scale = scale.astype(np.float32)
    Xr = X.tocsr()
    Xc = X.tocsc()
    XA = X.toarray()
    XA *= scale
    inplace_column_scale(Xc, scale)
    inplace_column_scale(Xr, scale)
    assert_array_almost_equal(Xr.toarray(), Xc.toarray())
    assert_array_almost_equal(XA, Xc.toarray())
    assert_array_almost_equal(XA, Xr.toarray())
    assert_raises(TypeError, inplace_column_scale, X.tolil(), scale)

def test_inplace_row_scale():
    rng = np.random.RandomState(0)
    X = sp.rand(100, 200, 0.05)
    Xr = X.tocsr()
    Xc = X.tocsc()
    XA = X.toarray()
    scale = rng.rand(100)
    XA *= scale.reshape(-1, 1)

    inplace_row_scale(Xc, scale)
    inplace_row_scale(Xr, scale)
    assert_array_almost_equal(Xr.toarray(), Xc.toarray())
    assert_array_almost_equal(XA, Xc.toarray())
    assert_array_almost_equal(XA, Xr.toarray())
    assert_raises(TypeError, inplace_column_scale, X.tolil(), scale)

    X = X.astype(np.float32)
    scale = scale.astype(np.float32)
    Xr = X.tocsr()
    Xc = X.tocsc()
    XA = X.toarray()
    XA *= scale.reshape(-1, 1)
    inplace_row_scale(Xc, scale)
    inplace_row_scale(Xr, scale)
    assert_array_almost_equal(Xr.toarray(), Xc.toarray())
    assert_array_almost_equal(XA, Xc.toarray())
    assert_array_almost_equal(XA, Xr.toarray())
    assert_raises(TypeError, inplace_column_scale, X.tolil(), scale)

def test_inplace_swap_row():
    X = np.array([[0, 3, 0],
                  [2, 4, 0],
                  [0, 0, 0],
                  [9, 8, 7],
                  [4, 0, 5]], dtype=np.float64)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)

    swap = linalg.get_blas_funcs(('swap',), (X,))
    swap = swap[0]
    X[0], X[-1] = swap(X[0], X[-1])
    inplace_swap_row(X_csr, 0, -1)
    inplace_swap_row(X_csc, 0, -1)
    assert_array_equal(X_csr.toarray(), X_csc.toarray())
    assert_array_equal(X, X_csc.toarray())
    assert_array_equal(X, X_csr.toarray())

    X[2], X[3] = swap(X[2], X[3])
    inplace_swap_row(X_csr, 2, 3)
    inplace_swap_row(X_csc, 2, 3)
    assert_array_equal(X_csr.toarray(), X_csc.toarray())
    assert_array_equal(X, X_csc.toarray())
    assert_array_equal(X, X_csr.toarray())
    assert_raises(TypeError, inplace_swap_row, X_csr.tolil())

    X = np.array([[0, 3, 0],
                  [2, 4, 0],
                  [0, 0, 0],
                  [9, 8, 7],
                  [4, 0, 5]], dtype=np.float32)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)
    swap = linalg.get_blas_funcs(('swap',), (X,))
    swap = swap[0]
    X[0], X[-1] = swap(X[0], X[-1])
    inplace_swap_row(X_csr, 0, -1)
    inplace_swap_row(X_csc, 0, -1)
    assert_array_equal(X_csr.toarray(), X_csc.toarray())
    assert_array_equal(X, X_csc.toarray())
    assert_array_equal(X, X_csr.toarray())
    X[2], X[3] = swap(X[2], X[3])
    inplace_swap_row(X_csr, 2, 3)
    inplace_swap_row(X_csc, 2, 3)
    assert_array_equal(X_csr.toarray(), X_csc.toarray())
    assert_array_equal(X, X_csc.toarray())
    assert_array_equal(X, X_csr.toarray())
    assert_raises(TypeError, inplace_swap_row, X_csr.tolil())

def test_inplace_swap_column():
    X = np.array([[0, 3, 0],
                  [2, 4, 0],
                  [0, 0, 0],
                  [9, 8, 7],
                  [4, 0, 5]], dtype=np.float64)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)

    swap = linalg.get_blas_funcs(('swap',), (X,))
    swap = swap[0]
    X[:, 0], X[:, -1] = swap(X[:, 0], X[:, -1])
    inplace_swap_column(X_csr, 0, -1)
    inplace_swap_column(X_csc, 0, -1)
    assert_array_equal(X_csr.toarray(), X_csc.toarray())
    assert_array_equal(X, X_csc.toarray())
    assert_array_equal(X, X_csr.toarray())

    X[:, 0], X[:, 1] = swap(X[:, 0], X[:, 1])
    inplace_swap_column(X_csr, 0, 1)
    inplace_swap_column(X_csc, 0, 1)
    assert_array_equal(X_csr.toarray(), X_csc.toarray())
    assert_array_equal(X, X_csc.toarray())
    assert_array_equal(X, X_csr.toarray())
    assert_raises(TypeError, inplace_swap_column, X_csr.tolil())

    X = np.array([[0, 3, 0],
                  [2, 4, 0],
                  [0, 0, 0],
                  [9, 8, 7],
                  [4, 0, 5]], dtype=np.float32)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)
    swap = linalg.get_blas_funcs(('swap',), (X,))
    swap = swap[0]
    X[:, 0], X[:, -1] = swap(X[:, 0], X[:, -1])
    inplace_swap_column(X_csr, 0, -1)
    inplace_swap_column(X_csc, 0, -1)
    assert_array_equal(X_csr.toarray(), X_csc.toarray())
    assert_array_equal(X, X_csc.toarray())
    assert_array_equal(X, X_csr.toarray())
    X[:, 0], X[:, 1] = swap(X[:, 0], X[:, 1])
    inplace_swap_column(X_csr, 0, 1)
    inplace_swap_column(X_csc, 0, 1)
    assert_array_equal(X_csr.toarray(), X_csc.toarray())
    assert_array_equal(X, X_csc.toarray())
    assert_array_equal(X, X_csr.toarray())
    assert_raises(TypeError, inplace_swap_column, X_csr.tolil())

def test_min_max_axis0():
    X = np.array([[0, 3, 0],
                  [2, -1, 0],
                  [0, 0, 0],
                  [9, 8, 7],
                  [4, 0, 5]], dtype=np.float64)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)

    mins_csr, maxs_csr = min_max_axis(X_csr, axis=0)
    assert_array_equal(mins_csr, X.min(axis=0))
    assert_array_equal(maxs_csr, X.max(axis=0))

    mins_csc, maxs_csc = min_max_axis(X_csc, axis=0)
    assert_array_equal(mins_csc, X.min(axis=0))
    assert_array_equal(maxs_csc, X.max(axis=0))

    X = X.astype(np.float32)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)
    mins_csr, maxs_csr = min_max_axis(X_csr, axis=0)
    assert_array_equal(mins_csr, X.min(axis=0))
    assert_array_equal(maxs_csr, X.max(axis=0))
    mins_csc, maxs_csc = min_max_axis(X_csc, axis=0)
    assert_array_equal(mins_csc, X.min(axis=0))
    assert_array_equal(maxs_csc, X.max(axis=0))

def test_min_max_axis1():
    X = np.array([[0, 3, 0],
                  [2, -1, 0],
                  [0, 0, 0],
                  [9, 8, 7],
                  [4, 0, 5]], dtype=np.float64)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)

    mins_csr, maxs_csr = min_max_axis(X_csr, axis=1)
    assert_array_equal(mins_csr, X.min(axis=1))
    assert_array_equal(maxs_csr, X.max(axis=1))

    mins_csc, maxs_csc = min_max_axis(X_csc, axis=1)
    assert_array_equal(mins_csc, X.min(axis=1))
    assert_array_equal(maxs_csc, X.max(axis=1))

    X = X.astype(np.float32)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)
    mins_csr, maxs_csr = min_max_axis(X_csr, axis=1)
    assert_array_equal(mins_csr, X.min(axis=1))
    assert_array_equal(maxs_csr, X.max(axis=1))
    mins_csc, maxs_csc = min_max_axis(X_csc, axis=1)
    assert_array_equal(mins_csc, X.min(axis=1))
    assert_array_equal(maxs_csc, X.max(axis=1))

def test_min_max_axis_errors():
    X = np.array([[0, 3, 0],
                  [2, -1, 0],
                  [0, 0, 0],
                  [9, 8, 7],
                  [4, 0, 5]], dtype=np.float64)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)
    assert_raises(TypeError, min_max_axis, X_csr.tolil(), axis=0)
    assert_raises(ValueError, min_max_axis, X_csr, axis=2)
    assert_raises(ValueError, min_max_axis, X_csc, axis=-3)

def test_count_nonzero():
    X = np.array([[0, 3, 0],
                  [2, -1, 0],
                  [0, 0, 0],
                  [9, 8, 7],
                  [4, 0, 5]], dtype=np.float64)
    X_csr = sp.csr_matrix(X)
    X_csc = sp.csc_matrix(X)
    X_nonzero = X != 0
    sample_weight = [.5, .2, .3, .1, .1]
    X_nonzero_weighted = X_nonzero * np.array(sample_weight)[:, None]

    for axis in [0, 1, -1, -2, None]:
        assert_array_almost_equal(count_nonzero(X_csr, axis=axis),
                                  X_nonzero.sum(axis=axis))
        assert_array_almost_equal(count_nonzero(X_csr, axis=axis,
                                                sample_weight=sample_weight),
                                  X_nonzero_weighted.sum(axis=axis))

    assert_raises(TypeError, count_nonzero, X_csc)
    assert_raises(ValueError, count_nonzero, X_csr, axis=2)

def test_csc_row_median():
    # Test csc_row_median actually calculates the median.

    # Test that it gives the same output when X is dense.
    rng = np.random.RandomState(0)
    X = rng.rand(100, 50)
    dense_median = np.median(X, axis=0)
    csc = sp.csc_matrix(X)
    sparse_median = csc_median_axis_0(csc)
    assert_array_equal(sparse_median, dense_median)

    # Test that it gives the same output when X is sparse
    X = rng.rand(51, 100)
    X[X < 0.7] = 0.0
    ind = rng.randint(0, 50, 10)
    X[ind] = -X[ind]
    csc = sp.csc_matrix(X)
    dense_median = np.median(X, axis=0)
    sparse_median = csc_median_axis_0(csc)
    assert_array_equal(sparse_median, dense_median)

    # Test for toy data.
    X = [[0, -2], [-1, -1], [1, 0], [2, 1]]
    csc = sp.csc_matrix(X)
    assert_array_equal(csc_median_axis_0(csc), np.array([0.5, -0.5]))
    X = [[0, -2], [-1, -5], [1, -3]]
    csc = sp.csc_matrix(X)
    assert_array_equal(csc_median_axis_0(csc), np.array([0., -3]))

    # Test that it raises an Error for non-csc matrices.
    assert_raises(TypeError, csc_median_axis_0, sp.csr_matrix(X))


# encoding: UTF-8
"""
This file tweaks ctaTemplate Module to suit multi-TimeFrame strategies.
"""

from ctaBase import *
from ctaTemplate import CtaTemplate
import numpy as np

########################################################################
class BreakOut(CtaTemplate):

    """
    "infoArray" 字典是用来储存辅助品种信息的, 可以是同品种的不同分钟k线, 也可以是不同品种的价格。

    调用的方法:
    价格序列:
    self.infoArray["数据库名 + 空格 + collection名"]["close"]
    self.infoArray["数据库名 + 空格 + collection名"]["high"]
    self.infoArray["数据库名 + 空格 + collection名"]["low"]

    单个价格:
    self.infoBar["数据库名 + 空格 + collection名"]
    返回的值为一个ctaBarData 或 None
    """

    #----------------------------------------------------------------------
    def __init__(self, ctaEngine, setting):
        """日内突破交易策略, 出场方式非常多, 本文件使用指标出场"""

        className = 'BreakOut'
        author = 'Joe'
        super(BreakOut, self).__init__(ctaEngine, setting)

        # 设置辅助品种数据字典
        self.infoArray = {}
        self.initInfobar = {}
        self.infoBar = {}

        # 缓存数据量
        self.bufferSize = 100
        self.bufferCount = 0
        self.initDays = 10

        # 设置参数
        self.pOBO_Mult = 0.5        # 计算突破点位
        # self.pProtMult = 2          # 止损的ATR倍数
        # self.pProfitMult = 2        # 止盈相对于止损的倍数
        # self.SlTp_On = False        # 止损止盈功能
        # self.EODTime = 15           # 设置日内平仓时间

        self.vOBO_stretch = EMPTY_FLOAT
        self.vOBO_initialpoint = EMPTY_FLOAT
        self.vOBO_level_L = EMPTY_FLOAT
        self.vOBO_level_S = EMPTY_FLOAT

        self.orderList = []

        # 参数列表，保存了参数的名称
        paramList = ['name',
                     'className',
                     'author',
                     'pOBO_Mult',
                     'pProtMult',
                     'pProfitMult',
                     'SlTp_On',
                     'EODTime']

        # 变量列表，保存了变量的名称
        varList = ['vOBO_stretch',
                   'vOBO_initialpoint',
                   'vOBO_level_L',
                   'vOBO_level_S']

        # ----------------------------------------------------------------------
    def onInit(self):
        """初始化策略（必须由用户继承实现）"""
        self.writeCtaLog(u'%s策略初始化' % self.name)

        # 载入历史数据，并采用回放计算的方式初始化策略数值
        initData = self.loadBar(self.initDays)
        for bar in initData:

            # 推送新数据, 同时检查是否有information bar需要推送
            # Update new bar, check whether the Time Stamp matching any information bar
            ibar = self.checkInfoBar(bar)
            self.onBar(bar, infobar=ibar)

        self.putEvent()

    #----------------------------------------------------------------------
    def onStart(self):
        """启动策略（必须由用户继承实现）"""
        self.writeCtaLog(u'%s策略启动' %self.name)
        self.putEvent()

    #----------------------------------------------------------------------
    def onStop(self):
        """停止策略（必须由用户继承实现）"""
        self.writeCtaLog(u'%s策略停止' %self.name)
        self.putEvent()

    # ----------------------------------------------------------------------
    def checkInfoBar(self, bar):
        """在初始化时, 检查辅助品种数据的推送(初始化结束后, 回测时不会调用)"""

        initInfoCursorDict = self.ctaEngine.initInfoCursor

        # 如果"initInfobar"字典为空, 初始化字典, 插入第一个数据
        # If dictionary "initInfobar" is empty, insert first data record
        if self.initInfobar == {}:
            for info_symbol in initInfoCursorDict:
                try:
                    self.initInfobar[info_symbol] = next(initInfoCursorDict[info_symbol])
                except StopIteration:
                    print "Data of information symbols is empty! Input is a list, not str."
                    raise

        # 若有某一品种的 TimeStamp 和执行报价的 TimeStamp 匹配, 则将"initInfobar"中的数据推送,
        # 然后更新该品种的数据
        # If any symbol's TimeStamp is matched with execution symbol's TimeStamp, return data
        # in "initInfobar", and update new data.
        temp = {}
        for info_symbol in self.initInfobar:

            data = self.initInfobar[info_symbol]

            # Update data only when Time Stamp is matched

            if (data is not None) and (data['datetime'] <= bar.datetime):

                try:
                    temp[info_symbol] = CtaBarData()
                    temp[info_symbol].__dict__ = data
                    self.initInfobar[info_symbol] = next(initInfoCursorDict[info_symbol])
                except StopIteration:
                    self.initInfobar[info_symbol] = None
                    self.ctaEngine.output("No more data for initializing %s." % (info_symbol,))
            else:
                temp[info_symbol] = None

        return temp

    # ----------------------------------------------------------------------
    def updateInfoArray(self, infobar):
        """收到Infomation Data, 更新辅助品种缓存字典"""

        for name in infobar:

            data = infobar[name]

            # Construct empty array
            if len(self.infoArray) < len(infobar) :
                self.infoArray[name] = {
                    "close": np.zeros(self.bufferSize),
                    "high": np.zeros(self.bufferSize),
                    "low": np.zeros(self.bufferSize),
                    "open": np.zeros(self.bufferSize)
                }

            if data is None:
                pass

            else:
                self.infoArray[name]["close"][0:self.bufferSize - 1] = \
                    self.infoArray[name]["close"][1:self.bufferSize]
                self.infoArray[name]["high"][0:self.bufferSize - 1] = \
                    self.infoArray[name]["high"][1:self.bufferSize]
                self.infoArray[name]["low"][0:self.bufferSize - 1] = \
                    self.infoArray[name]["low"][1:self.bufferSize]
                self.infoArray[name]["open"][0:self.bufferSize - 1] = \
                    self.infoArray[name]["open"][1:self.bufferSize]

                self.infoArray[name]["close"][-1] = data.close
                self.infoArray[name]["high"][-1] = data.high
                self.infoArray[name]["low"][-1] = data.low
                self.infoArray[name]["open"][-1] = data.open

    # ----------------------------------------------------------------------
    def onBar(self, bar, **kwargs):
        """收到Bar推送（必须由用户继承实现）"""

        # Update infomation data
        # "infobar"是由不同时间或不同品种的品种数据组成的字典, 如果和执行品种的 TimeStamp 不匹配,
        # 则传入的是"None", 当time stamp和执行品种匹配时, 传入的是"Bar"
        if "infobar" in kwargs:
            self.infoBar = kwargs["infobar"]
            self.updateInfoArray(kwargs["infobar"])

        # 若读取的缓存数据不足, 不考虑交易
        self.bufferCount += 1
        if self.bufferCount < self.bufferSize:
            return

        # 计算指标数值
        a = np.sum(self.infoArray["TestData @GC_1D"]["close"])
        if a == 0.0:
            return

        # Only updating indicators when information bar changes
        # 只有在30min或者1d K线更新后才更新指标
        TradeOn = False
        if any([i is not None for i in self.infoBar]):
            TradeOn = True
            self.vRange = self.infoArray["TestData @GC_1D"]["high"][-1] -\
                          self.infoArray["TestData @GC_1D"]["low"][-1]
            self.vOBO_stretch = self.vRange * self.pOBO_Mult
            self.vOBO_initialpoint = self.infoArray["TestData @GC_1D"]["close"][-1]
            self.vOBO_level_L = self.vOBO_initialpoint + self.vOBO_stretch
            self.vOBO_level_S = self.vOBO_initialpoint - self.vOBO_stretch

            self.atrValue30M = talib.abstract.ATR(self.infoArray["TestData @GC_30M"])[-1]

        # 判断是否要进行交易

        # 当前无仓位
        if (self.pos == 0 and TradeOn == True):

            # 撤销之前发出的尚未成交的委托（包括限价单和停止单）
            for orderID in self.orderList:
                self.cancelOrder(orderID)
            self.orderList = []

            # 若上一个30分钟K线的最高价大于OBO_level_L
            # 且当前的价格大于OBO_level_L, 则买入
            if self.infoArray["TestData @GC_30M"]["high"][-1] > self.vOBO_level_L:

                if bar.close > self.vOBO_level_L:

                    self.buy(bar.close + 0.5, 1)

                    # 下单后, 在下一个30Min K线之前不交易
                    TradeOn = False

            # 若上一个30分钟K线的最高价低于OBO_level_S
            # 且当前的价格小于OBO_level_S, 则卖出
            elif self.infoArray["TestData @GC_30M"]["low"][-1] < self.vOBO_level_S:

                if bar.close < self.vOBO_level_S:

                    self.short(bar.close - 0.5, 1)

                    # 下单后, 在下一个30Min K线之前不交易
                    TradeOn = False

        # 持有多头仓位
        elif self.pos > 0:

            # 当价格低于initialpoint水平, 出场
            if bar.close < self.vOBO_initialpoint:
                self.sell(bar.close - 0.5 , 1)

        # 持有空头仓位
        elif self.pos < 0:

            # 当价格高于initialpoint水平, 出场
            if bar.close > self.vOBO_initialpoint:
                self.cover(bar.close + 0.5, 1)

        # 发出状态更新事件
        self.putEvent()

    # ----------------------------------------------------------------------
    def onOrder(self, order):
        """收到委托变化推送（必须由用户继承实现）"""
        pass

    # ----------------------------------------------------------------------
    def onTrade(self, trade):
        pass

if __name__ == '__main__':
    # 提供直接双击回测的功能
    # 导入PyQt4的包是为了保证matplotlib使用PyQt4而不是PySide，防止初始化出错
    from ctaBacktestMultiTF import *
    from PyQt4 import QtCore, QtGui
    import time

    '''
    创建回测引擎
    设置引擎的回测模式为K线
    设置回测用的数据起始日期
    载入历史数据到引擎中
    在引擎中创建策略对象

    Create backtesting engine
    Set backtest mode as "Bar"
    Set "Start Date" of data range
    Load historical data to engine
    Create strategy instance in engine
    '''
    engine = BacktestEngineMultiTF()
    engine.setBacktestingMode(engine.BAR_MODE)
    engine.setStartDate('20120101')
    engine.setEndDate('20150101')
    engine.setDatabase("TestData", "@GC_1M", info_symbol=[("TestData","@GC_30M"),
                                                          ("TestData","@GC_1D")])

    # Set parameters for strategy
    engine.initStrategy(BreakOut, {})

    # 设置产品相关参数
    engine.setSlippage(0.2)  # 股指1跳
    engine.setCommission(0.3 / 10000)  # 万0.3
    engine.setSize(1)  # 股指合约大小

    # 开始跑回测
    start = time.time()

    engine.runBacktesting()

    # 显示回测结果
    engine.showBacktestingResult()

    print 'Time consumed：%s' % (time.time() - start)


import pandas as pd
import numpy as np

from utils.webdata import get_close_of_symbols

class TradeSimulator(object):

    def __init__(self, start_val=1000000, leverage=2.0, allow_short=True):
        """

        Parameters
        ----------
        start_val: float
            start value of the portfolio
        leverage: float
            max leverage
        allow_short: boolean
            allows to sell short
        """
        self.start_val = start_val
        self.leverage = leverage
        self.allow_short = allow_short

    def compute_leverage(self, prices, shares, cash, order_type, order_share, order_symbol):
        """
        Compute the leverage of the shares

        Parameters
        ----------
        prices: Series
        shares: dict
            contains the current shares for each symbol
        cash: float
            current cash
        order_type: [BUY or SELL]
            the type of the order
        order_share: int
            the number of shares of the order
        order_symbol: string
            the symbol of the order

        Returns
        ----------
        leverage: float
        """
        if order_type == 'BUY':
            shares[order_symbol] += order_share
            cash -= prices[order_symbol] * order_share
        else:
            shares[order_symbol] -= order_share
            cash += prices[order_symbol] * order_share

        longs = shorts = 0
        for symbol in shares.keys():
            if shares[symbol] >= 0:
                longs += shares[symbol] * prices[symbol]
            else:
                shorts -= shares[symbol] * prices[symbol]

        leverage = (longs + shorts) / (longs - shorts + cash)
        return leverage

    def simulate(self, start_date=None, end_date=None, prices=None, orders=None, orders_file=None):
        """Simulate the trades with the given orders and the prices.

        Parameters
        ----------
        start_date: string
        end_date: string
        prices: DataFrame
        orders: DataFrame
        orders_file: string

        Returns
        ----------
        portvals: DataFrame
            the daily portfolio values in the simulation
        """
        if orders is None:
            orders = pd.read_csv(orders_file, parse_dates=True)

        symbols = list(set(orders['Symbol']))
        if prices is None:
            prices = get_close_of_symbols(symbols, start_date, end_date, add_spy=True) # add SPY so as to remove no-trade days
            if prices is None:
                return None

            prices.drop('SPY', axis=1, inplace=True)       # remove SPY

        dates = prices.index                           # update dates
        # init daily shares
        shares = prices.copy()                         # record the shares every day
        shares.loc[:, :] = np.nan
        last_share = dict.fromkeys(shares.columns, 0)  # record the total shares of each symbol
        # init daily cashes
        cashes = pd.Series({'Cash':np.nan}, index=dates) # record the daily cashes
        last_cash = self.start_val                          # record total cash

        # iterate orders and simulate the trades
        for i in range(len(orders)):
            symbol = orders.loc[i, 'Symbol']
            share = orders.loc[i, 'Shares']
            date = orders.loc[i, 'Date']
            operate = orders.loc[i, 'Order']
            price = prices.loc[date, symbol]

            # check leverage
            tmp_leverage = self.compute_leverage(prices.loc[date, :], last_share.copy(), last_cash,
                                            operate, share, symbol)
            if tmp_leverage > self.leverage:
                continue

            if operate == 'BUY':
                last_share[symbol] += share
                shares.loc[date, symbol] = last_share[symbol]
                val = last_cash - price * share
                cashes[date] = last_cash = val
            else:
                temp_share = last_share[symbol] - share
                # short check
                if not self.allow_short and temp_share < 0:
                    continue
                shares.loc[date, symbol] = last_share[symbol] = temp_share
                last_cash += price * share
                cashes[date] = last_cash

        # init the nan values of the first row of shares before invoking fillna
        for symbol in shares.columns:
            if pd.isnull(shares.loc[dates[0], symbol]):
                shares.loc[dates[0], symbol] = 0

        shares.fillna(method="ffill", inplace=True)
        # init the nan value of the first row of cashes before invoking fillna
        if pd.isnull(cashes.ix[0]):
            cashes.ix[0] = self.start_val
        cashes.fillna(method='ffill', inplace=True)
        values = (prices * shares).sum(axis=1)
        portvals = (values + cashes).to_frame()
        portvals.rename(columns={portvals.columns[0]: "Portfolio"}, inplace=True)
        return portvals


# Author: Alexander Fabisch <afabisch@informatik.uni-bremen.de>
#
# License: BSD 3 clause

import sys
from sklearn.externals.six.moves import cStringIO as StringIO
import numpy as np
import warnings
from sklearn.base import BaseEstimator
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_warns
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_false
from sklearn.datasets import make_classification

with warnings.catch_warnings():
    warnings.simplefilter('ignore')
    from sklearn.learning_curve import learning_curve, validation_curve
    from sklearn.cross_validation import KFold

from sklearn.linear_model import PassiveAggressiveClassifier

class MockImprovingEstimator(BaseEstimator):
    """Dummy classifier to test the learning curve"""
    def __init__(self, n_max_train_sizes):
        self.n_max_train_sizes = n_max_train_sizes
        self.train_sizes = 0
        self.X_subset = None

    def fit(self, X_subset, y_subset=None):
        self.X_subset = X_subset
        self.train_sizes = X_subset.shape[0]
        return self

    def predict(self, X):
        raise NotImplementedError

    def score(self, X=None, Y=None):
        # training score becomes worse (2 -> 1), test error better (0 -> 1)
        if self._is_training_data(X):
            return 2. - float(self.train_sizes) / self.n_max_train_sizes
        else:
            return float(self.train_sizes) / self.n_max_train_sizes

    def _is_training_data(self, X):
        return X is self.X_subset

class MockIncrementalImprovingEstimator(MockImprovingEstimator):
    """Dummy classifier that provides partial_fit"""
    def __init__(self, n_max_train_sizes):
        super(MockIncrementalImprovingEstimator,
              self).__init__(n_max_train_sizes)
        self.x = None

    def _is_training_data(self, X):
        return self.x in X

    def partial_fit(self, X, y=None, **params):
        self.train_sizes += X.shape[0]
        self.x = X[0]

class MockEstimatorWithParameter(BaseEstimator):
    """Dummy classifier to test the validation curve"""
    def __init__(self, param=0.5):
        self.X_subset = None
        self.param = param

    def fit(self, X_subset, y_subset):
        self.X_subset = X_subset
        self.train_sizes = X_subset.shape[0]
        return self

    def predict(self, X):
        raise NotImplementedError

    def score(self, X=None, y=None):
        return self.param if self._is_training_data(X) else 1 - self.param

    def _is_training_data(self, X):
        return X is self.X_subset

class MockEstimatorFailing(BaseEstimator):
    """Dummy classifier to test error_score in learning curve"""
    def fit(self, X_subset, y_subset):
        raise ValueError()

    def score(self, X=None, y=None):
        return None

class MockEstimatorWithSingleFitCallAllowed(MockEstimatorWithParameter):
    """Dummy classifier that disallows repeated calls of fit method"""

    def fit(self, X_subset, y_subset):
        assert_false(
            hasattr(self, 'fit_called_'),
            'fit is called the second time'
        )
        self.fit_called_ = True
        return super(type(self), self).fit(X_subset, y_subset)

def test_learning_curve():
    X, y = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockImprovingEstimator(20)
    with warnings.catch_warnings(record=True) as w:
        train_sizes, train_scores, test_scores = learning_curve(
            estimator, X, y, cv=3, train_sizes=np.linspace(0.1, 1.0, 10))
    if len(w) > 0:
        raise RuntimeError("Unexpected warning: %r" % w[0].message)
    assert_equal(train_scores.shape, (10, 3))
    assert_equal(test_scores.shape, (10, 3))
    assert_array_equal(train_sizes, np.linspace(2, 20, 10))
    assert_array_almost_equal(train_scores.mean(axis=1),
                              np.linspace(1.9, 1.0, 10))
    assert_array_almost_equal(test_scores.mean(axis=1),
                              np.linspace(0.1, 1.0, 10))

def test_learning_curve_unsupervised():
    X, _ = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockImprovingEstimator(20)
    train_sizes, train_scores, test_scores = learning_curve(
        estimator, X, y=None, cv=3, train_sizes=np.linspace(0.1, 1.0, 10))
    assert_array_equal(train_sizes, np.linspace(2, 20, 10))
    assert_array_almost_equal(train_scores.mean(axis=1),
                              np.linspace(1.9, 1.0, 10))
    assert_array_almost_equal(test_scores.mean(axis=1),
                              np.linspace(0.1, 1.0, 10))

def test_learning_curve_verbose():
    X, y = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockImprovingEstimator(20)

    old_stdout = sys.stdout
    sys.stdout = StringIO()
    try:
        train_sizes, train_scores, test_scores = \
            learning_curve(estimator, X, y, cv=3, verbose=1)
    finally:
        out = sys.stdout.getvalue()
        sys.stdout.close()
        sys.stdout = old_stdout

    assert("[learning_curve]" in out)

def test_learning_curve_error_score():
    X, y = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockEstimatorFailing()
    _, _, test_scores = learning_curve(estimator, X, y, cv=3, error_score=0)
    all_zeros = not np.any(test_scores)
    assert(all_zeros)

def test_learning_curve_error_score_default_raise():
    X, y = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockEstimatorFailing()
    assert_raises(ValueError, learning_curve, estimator, X, y, cv=3)

def test_learning_curve_incremental_learning_not_possible():
    X, y = make_classification(n_samples=2, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    # The mockup does not have partial_fit()
    estimator = MockImprovingEstimator(1)
    assert_raises(ValueError, learning_curve, estimator, X, y,
                  exploit_incremental_learning=True)

def test_learning_curve_incremental_learning():
    X, y = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockIncrementalImprovingEstimator(20)
    train_sizes, train_scores, test_scores = learning_curve(
        estimator, X, y, cv=3, exploit_incremental_learning=True,
        train_sizes=np.linspace(0.1, 1.0, 10))
    assert_array_equal(train_sizes, np.linspace(2, 20, 10))
    assert_array_almost_equal(train_scores.mean(axis=1),
                              np.linspace(1.9, 1.0, 10))
    assert_array_almost_equal(test_scores.mean(axis=1),
                              np.linspace(0.1, 1.0, 10))

def test_learning_curve_incremental_learning_unsupervised():
    X, _ = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockIncrementalImprovingEstimator(20)
    train_sizes, train_scores, test_scores = learning_curve(
        estimator, X, y=None, cv=3, exploit_incremental_learning=True,
        train_sizes=np.linspace(0.1, 1.0, 10))
    assert_array_equal(train_sizes, np.linspace(2, 20, 10))
    assert_array_almost_equal(train_scores.mean(axis=1),
                              np.linspace(1.9, 1.0, 10))
    assert_array_almost_equal(test_scores.mean(axis=1),
                              np.linspace(0.1, 1.0, 10))

def test_learning_curve_batch_and_incremental_learning_are_equal():
    X, y = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    train_sizes = np.linspace(0.2, 1.0, 5)
    estimator = PassiveAggressiveClassifier(max_iter=1, tol=None,
                                            shuffle=False)

    train_sizes_inc, train_scores_inc, test_scores_inc = \
        learning_curve(
            estimator, X, y, train_sizes=train_sizes,
            cv=3, exploit_incremental_learning=True)
    train_sizes_batch, train_scores_batch, test_scores_batch = \
        learning_curve(
            estimator, X, y, cv=3, train_sizes=train_sizes,
            exploit_incremental_learning=False)

    assert_array_equal(train_sizes_inc, train_sizes_batch)
    assert_array_almost_equal(train_scores_inc.mean(axis=1),
                              train_scores_batch.mean(axis=1))
    assert_array_almost_equal(test_scores_inc.mean(axis=1),
                              test_scores_batch.mean(axis=1))

def test_learning_curve_n_sample_range_out_of_bounds():
    X, y = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockImprovingEstimator(20)
    assert_raises(ValueError, learning_curve, estimator, X, y, cv=3,
                  train_sizes=[0, 1])
    assert_raises(ValueError, learning_curve, estimator, X, y, cv=3,
                  train_sizes=[0.0, 1.0])
    assert_raises(ValueError, learning_curve, estimator, X, y, cv=3,
                  train_sizes=[0.1, 1.1])
    assert_raises(ValueError, learning_curve, estimator, X, y, cv=3,
                  train_sizes=[0, 20])
    assert_raises(ValueError, learning_curve, estimator, X, y, cv=3,
                  train_sizes=[1, 21])

def test_learning_curve_remove_duplicate_sample_sizes():
    X, y = make_classification(n_samples=3, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockImprovingEstimator(2)
    train_sizes, _, _ = assert_warns(
        RuntimeWarning, learning_curve, estimator, X, y, cv=3,
        train_sizes=np.linspace(0.33, 1.0, 3))
    assert_array_equal(train_sizes, [1, 2])

def test_learning_curve_with_boolean_indices():
    X, y = make_classification(n_samples=30, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    estimator = MockImprovingEstimator(20)
    cv = KFold(n=30, n_folds=3)
    train_sizes, train_scores, test_scores = learning_curve(
        estimator, X, y, cv=cv, train_sizes=np.linspace(0.1, 1.0, 10))
    assert_array_equal(train_sizes, np.linspace(2, 20, 10))
    assert_array_almost_equal(train_scores.mean(axis=1),
                              np.linspace(1.9, 1.0, 10))
    assert_array_almost_equal(test_scores.mean(axis=1),
                              np.linspace(0.1, 1.0, 10))

def test_validation_curve():
    X, y = make_classification(n_samples=2, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)
    param_range = np.linspace(0, 1, 10)
    with warnings.catch_warnings(record=True) as w:
        train_scores, test_scores = validation_curve(
            MockEstimatorWithParameter(), X, y, param_name="param",
            param_range=param_range, cv=2
        )
    if len(w) > 0:
        raise RuntimeError("Unexpected warning: %r" % w[0].message)

    assert_array_almost_equal(train_scores.mean(axis=1), param_range)
    assert_array_almost_equal(test_scores.mean(axis=1), 1 - param_range)

def test_validation_curve_clone_estimator():
    X, y = make_classification(n_samples=2, n_features=1, n_informative=1,
                               n_redundant=0, n_classes=2,
                               n_clusters_per_class=1, random_state=0)

    param_range = np.linspace(1, 0, 10)
    _, _ = validation_curve(
        MockEstimatorWithSingleFitCallAllowed(), X, y,
        param_name="param", param_range=param_range, cv=2
    )


#!/usr/bin/env python
# -*- coding: utf-8 -*-
__author__ = 'Will Brennan'

# Built-in Module
import os
import time
import logging
import warnings
import cPickle as pickle
from datetime import datetime
# Standard Modules
import numpy
import sklearn
import theano
import theano.tensor as T
# Custom Modules
import Scripts
import Layers

logger = logging.getLogger('main')
warnings.simplefilter("ignore", DeprecationWarning)

class DeepConv(object):
    def __init__(self, debug=False, load=False, save=False):
        self.args_debug = debug
        self.args_load = load
        self.args_save = save

        if self.args_debug:
            theano.exception_verbosity = 'high'
        if self.args_load:
            self.load()
        else:
            self.layers = None
        self.test_model = None
        self.validate_model = None
        self.train_model = None
        self.pred_model = None
        self.index, self.x, self.y = T.lscalar(), T.matrix('x'), T.ivector('y')

    def fit(self, data, labels, test_data, test_labels, learning_rate=0.1, n_epochs=250, nkerns=[20, 50], batch_size=500):
        logger.info('Initialising the classifier')
        rng = numpy.random.RandomState()
        data, labels = Scripts.shared_dataset(data_x=data, data_y=labels)
        test_data, test_labels = Scripts.shared_dataset(data_x=test_data, data_y=test_labels)
        if batch_size < 1:
            batch_size = data.get_value(borrow=True).shape[0]
        n_train_batches = data.get_value(borrow=True).shape[0]/batch_size
        n_test_batches = test_data.get_value(borrow=True).shape[0]/batch_size
        logger.info('Constructing the classifier')
        self.layers = []
        self.layers.append(Layers.PoolingLayer(
            rng,
            input=self.x.reshape((batch_size, 1, 28, 28)),
            image_shape=(batch_size, 1, 28, 28),
            filter_shape=(nkerns[0], 1, 5, 5),
            poolsize=(2, 2)
        ))
        self.layers.append(Layers.PoolingLayer(
            rng,
            input=self.layers[-1].output,
            image_shape=(batch_size, nkerns[0], 12, 12),
            filter_shape=(nkerns[1], nkerns[0], 5, 5),
            poolsize=(2, 2)
        ))
        self.layers.append(Layers.HiddenLayer(
            rng,
            input=self.layers[-1].output.flatten(2),
            n_in=nkerns[1] * 4 * 4,
            n_out=500,
            activation=T.tanh
        ))
        self.layers.append(Layers.LogisticRegression(
            input=self.layers[-1].output,
            n_in=500,
            n_out=10
        ))
        test_givens = {self.x: test_data[self.index * batch_size: (self.index + 1) * batch_size], self.y: test_labels[self.index * batch_size: (self.index + 1) * batch_size]}
        self.test_model = theano.function([self.index], self.layers[-1].errors(self.y), givens=test_givens)
        params = self.layers[0].params + self.layers[1].params + self.layers[2].params + self.layers[3].params
        cost = self.layers[-1].negative_log_likelihood(self.y)
        grads = T.grad(cost, params)
        updates = [(param_i, param_i - learning_rate * grad_i) for param_i, grad_i in zip(params, grads)]
        train_givens = {self.x: data[self.index * batch_size: (self.index + 1) * batch_size], self.y: labels[self.index * batch_size: (self.index + 1) * batch_size]}
        self.train_model = theano.function([self.index], cost, updates=updates, givens=train_givens)
        patience, patience_increase = 10000, 2
        validation_frequency = min(n_train_batches, patience / 2)
        epoch, count = 0, 0
        start_time = time.time()
        n_iters = n_epochs*n_train_batches
        logger.info("Fitting Classifier")
        logger.debug("{0} epochs, {1} batches, {2} iterations".format(n_epochs, n_train_batches, n_iters))
        while epoch < n_epochs and patience > count:
            epoch += 1
            for minibatch_index in xrange(n_train_batches):
                count = (epoch - 1) * n_train_batches + minibatch_index
                if count % 50 == 0:
                    percentage = round(100.0*count/n_iters, 2)
                    if percentage == 0:
                        time_stamp = "Null"
                    else:
                        time_stamp = datetime.utcfromtimestamp((time.time()-start_time)*(100.0/percentage)+start_time)
                    logger.info("training is {0}% complete (Completion at {1})".format(round(percentage, 2), time_stamp))
                train_cost = self.train_model(minibatch_index)
                if (count + 1) % validation_frequency == 0:
                    testlosses = [self.test_model(i) for i in xrange(n_test_batches)]
                    test_score = numpy.mean(testlosses)
                    logger.info('Test error of {0}% achieved on Epoch {1} Iteration {2}'.format(test_score*100.0, epoch, count+1))
                logger.debug("Iteration number {0}".format(count))
        logger.debug('Optimization complete.')
        logger.debug('Conducting final model testing')
        testlosses = [self.test_model(i) for i in xrange(n_test_batches)]
        test_score = numpy.mean(testlosses)
        t_taken = int((time.time()-start_time)/60.0)
        logger.info('Training Complete')
        logger.info('Test score of {0}%, training time {1}m'.format(test_score*100.0, t_taken))
        if self.args_save:
            self.save()

    def predict(self, x_data, batch_size=500):
        assert isinstance(x_data, numpy.ndarray), "input features must be a numpy array"
        assert len(x_data.shape) == 2, "it must be an array of feature vectors"
        logger.info('classifier prediction called')
        logger.debug('x_data shape: {0}'.format(x_data.shape))
        logger.debug('forming prediction function')
        x_data = Scripts.shared_dataset(data_x=x_data)
        givens = {self.x: x_data[self.index * batch_size: (self.index + 1) * batch_size]}
        pred_model = theano.function(inputs=[self.index], outputs=self.layers[-1].y_pred, givens=givens, on_unused_input='warn', allow_input_downcast=True)
        logger.debug('input shape: {0}'.format(x_data.get_value(borrow=True).shape))
        logger.info('beginning prediction on x_data')
        n_batches = x_data.get_value(borrow=True).shape[0]/batch_size
        result = []
        for batch_index in range(n_batches):
            logger.debug('processing batch {0}'.format(batch_index))
            batch_result = pred_model(batch_index)
            logger.debug('result generated')
            result = numpy.hstack((result, batch_result))
        logger.debug('output shape: {0}'.format(len(result)))
        # batch size, rows, columns, channels.
        return result

    def score(self, test_data, test_labels, batch_size=500):
        logger.info('Generating Classification Score')
        logger.debug('creating shared datasets')
        test_data, test_labels = Scripts.shared_dataset(data_x=test_data, data_y=test_labels)
        logger.debug('producing batch information')
        n_test_batches = test_data.get_value(borrow=True).shape[0]
        n_test_batches /= batch_size
        logger.debug('generating theano functions')
        test_givens = {self.x: test_data[self.index * batch_size: (self.index + 1) * batch_size], self.y: test_labels[self.index * batch_size: (self.index + 1) * batch_size]}
        test_model = theano.function(inputs=[self.index], outputs=self.layers[-1].errors(self.y), givens=test_givens, on_unused_input='warn')
        logger.debug('producing test results')
        losses = [test_model(i) for i in range(n_test_batches)]
        return 1.0-numpy.mean(losses)

    def score_report(self, y_test, y_pred):
        scores = sklearn.metrics.classification_report(y_test, y_pred)
        logger.info("\n"+scores)

    def save(self, path="DeepConvolution.pkl"):
        path = os.path.join(os.path.split(__file__)[0], path)
        logger.info("Saving layers to {0}".format(path))
        with open(path, 'wb') as output:
            pickle.dump(self.layers, output, pickle.HIGHEST_PROTOCOL)
        logger.debug("Successfully saved")

    def load(self, path="DeepConvolution.pkl"):
        path = os.path.join(os.path.split(__file__)[0], path)
        logger.info("Loading layers from {0}".format(path))
        assert os.path.exists(path), "Specified Path is not valid"
        with open(path, "rb") as input_file:
            self.layers = pickle.load(input_file)
        logger.debug("Successfully loaded")


'''
Created on 7 Jun 2013

@author: Jamie
'''

import urllib2
import math
import re
import itertools
import argparse
from bs4 import BeautifulSoup
import matplotlib.pyplot as plt

markets = {'UK': {'url': 'gb/united%20kingdom/', 'curr': 'GBP'},
           'USA': {'url': 'us/united%20states/', 'curr': 'USD'},
           'GERMANY': {'url': 'de/germany/', 'curr': 'EUR'},
           'ITALY': {'url': 'it/italy/', 'curr': 'EUR'},
           'SPAIN': {'url': 'es/spain/', 'curr': 'EUR'},
           'AUSTRALIA': {'url': 'au/australia/', 'curr': 'AUD'},
           'ARGENTINA': {'url': 'ar/argentina/', 'curr': 'ARS'},
           'NETHERLANDS': {'url': 'nl/netherlands/', 'curr': 'EUR'},
           'BRAZIL': {'url': 'br/brazil/', 'curr': 'BRL'},
           'FRANCE': {'url': 'fr/france/', 'curr': 'EUR'},
           'GBP': {'url': 'gbp/', 'curr': 'GBP'},
           'USD': {'url': 'usd/', 'curr': 'USD'},
           'EUR': {'url': 'eur/', 'curr': 'EUR'},
           }

methods = {'NATIONAL_BANK_TRANSFER': 'national-bank-transfer/'}
method = ''

buy_url = 'https://localbitcoins.com/buy-bitcoins-online/'
sell_url = 'https://localbitcoins.com/sell-bitcoins-online/'

def get_ads_dict(soup, buy_sell):
    prices = get_prices(soup)
    users = get_users(soup)
    amounts = get_amounts(soup)
    amounts = [a/p for a,p in zip(amounts, prices)] # To give amount in BTC
    currency = get_currency(soup)
    methods = get_methods(soup)
    lists = set(zip(prices, users, amounts, currency))
    if buy_sell == 'buy':
        sorted_ads = sorted(lists)
    elif buy_sell == 'sell':
        sorted_ads = sorted(lists)[::-1]
    prices = [item[0] for item in sorted_ads]
    users = [item[1] for item in sorted_ads]
    amounts = [item[2] for item in sorted_ads]
    currency = [item[3] for item in sorted_ads]
    depth = get_depth(amounts)
    ads_dict = {'users': users, 'prices': prices, 'amounts': amounts,
                'depth': depth, 'currency': currency, 'methods': methods}
    return ads_dict

def get_prices(soup):
    ''' Returns a list of prices '''
    prices = soup.find_all('td', attrs={'class':"column-price"})
    prices = [float(re.findall("\d+.\d+", price.get_text())[0]) for price in prices]
    return prices

def get_currency(soup):
    ''' Returns a list of currencies '''
    prices = soup.find_all('td', attrs={'class':"column-price"})
    currencies = [price.get_text().split()[-1] for price in prices]
    return currencies

def get_methods(soup):
    ''' Returns a list of payment methods '''
    methods = soup.find_all('tr', attrs={'class':"clickable"})
    methods = [method.get_text().split('\n')[-7].strip() for method in methods]
    return methods

def get_users(soup):
    ''' Returns a list of users '''
    users = soup.find_all('td', attrs={'class':"column-user"})
    users = [user.get_text().split()[0] for user in users]
    return users

def get_amounts(soup):
    ''' Returns a list of amounts '''
    raw_amounts = soup.find_all('td', attrs={'class':"column-limit"})
    amounts = []
    for amount in raw_amounts:
        try:
            amounts += [float(amount.get_text().split()[2])]
        except:
            amounts += [0.0]
    return amounts

def get_depth(amounts):
    ''' Generates the cumulative amount for each point on the curve '''
    cum_amounts = []
    cum_amount = 0
    for amount in amounts:
        cum_amount += amount
        cum_amounts += [cum_amount]
    return cum_amounts

def get_buy_curve(market):
    response = urllib2.urlopen(buy_url + market['url'] + method)
    soup = BeautifulSoup(response)
    buy_ads = get_ads_dict(soup, 'buy')
    buy_prices = [i for i,j in zip(buy_ads['prices'], buy_ads['currency']) if j == market['curr']]
    buy_depth = [i for i,j in zip(buy_ads['depth'], buy_ads['currency']) if j == market['curr']]
    buy_prices = double_list(buy_prices)[1:]
    buy_depth = double_list(buy_depth)[:-1]
    return buy_prices[:-2], buy_depth[:-2]

def get_sell_curve(market):
    response = urllib2.urlopen(sell_url + market['url'] + method)
    soup = BeautifulSoup(response)
    sell_ads = get_ads_dict(soup, 'sell')
    sell_prices = [i for i,j in zip(sell_ads['prices'], sell_ads['currency']) if j == market['curr']][::-1]
    sell_depth = [i for i,j in zip(sell_ads['depth'], sell_ads['currency']) if j == market['curr']][::-1]
    sell_prices = double_list(sell_prices)[1:]
    sell_depth = double_list(sell_depth)[:-1]
    return sell_prices, sell_depth

def plot_chart(ax, buy, sell):
    ax.plot(buy[0], buy[1], color='r')
    ax.plot(sell[0], sell[1], color='g')

def double_list(list_in):
    iters = [iter(list_in), iter(list_in)]
    return list(it.next() for it in itertools.cycle(iters))

def get_bid(country):
    market = markets[country]
    response = urllib2.urlopen(buy_url + market['url'] + method)
    soup = BeautifulSoup(response)
    buy_ads = get_ads_dict(soup, 'buy')
    bid = buy_ads['prices'][0]
    return bid

def get_ask(country):
    market = markets[country]
    response = urllib2.urlopen(sell_url + market['url'] + method)
    soup = BeautifulSoup(response)
    sell_ads = get_ads_dict(soup, 'sell')
    ask = sell_ads['prices'][0]
    return ask

def make_charts(*args):
    if len(args[0].countries) == 0:
        selection = ['UK','USA','SPAIN','FRANCE','GERMANY','BRAZIL']
    else:
        selection = args[0].countries
    fig = plt.figure()
    dim = math.ceil(len(selection)**0.5)
    for x, s in enumerate(selection):
        market = markets[s]
        # method = methods['NATIONAL_BANK_TRANSFER']
        ax = fig.add_subplot(dim, dim, x+1)
        ax.set_xlabel(market['curr'])
        ax.set_ylabel('BTC')
        ax.set_title('Local Bitcoins online: %s' % s)
        buy_curve = get_buy_curve(market)
        sell_curve = get_sell_curve(market)
        plot_chart(ax, buy_curve, sell_curve)
    plt.tight_layout()
    plt.show()

def main():
    parser = argparse.ArgumentParser(description='Display charts of the Local Bitcoin market depth.')
    parser.add_argument('countries', type=str, nargs='*',
                        help='optionally specify any number of country names')
    args = parser.parse_args()
    make_charts(args)

if __name__ == '__main__':
    main()


from matplotlib.colors import LinearSegmentedColormap
from numpy import nan, inf
cm_data = [[1., 1., 1.],
[1., 1., 1.],
[0.498039, 0.498039, 0.498039],
[0., 0., 0.513725],
[0., 0., 0.533333],
[0., 0., 0.54902],
[0., 0., 0.564706],
[0., 0., 0.580392],
[0., 0., 0.6],
[0., 0., 0.615686],
[0., 0., 0.568627],
[0., 0., 0.584314],
[0., 0., 0.666667],
[0., 0., 0.682353],
[0., 0., 0.698039],
[0., 0., 0.713725],
[0., 0., 0.733333],
[0., 0., 0.74902],
[0., 0., 0.764706],
[0., 0., 0.780392],
[0., 0., 0.717647],
[0., 0., 0.733333],
[0., 0., 0.831373],
[0., 0., 0.847059],
[0., 0., 0.866667],
[0., 0., 0.882353],
[0., 0., 0.898039],
[0., 0., 0.913725],
[0., 0., 0.933333],
[0., 0., 0.94902],
[0., 0., 0.866667],
[0., 0., 0.882353],
[0., 0., 1.],
[0., 0.027451, 0.968627],
[0., 0.0588235, 0.937255],
[0., 0.0901961, 0.905882],
[0., 0.121569, 0.87451],
[0., 0.152941, 0.843137],
[0., 0.184314, 0.811765],
[0., 0.215686, 0.780392],
[0., 0.223529, 0.67451],
[0., 0.25098, 0.643137],
[0., 0.309804, 0.686275],
[0., 0.341176, 0.654902],
[0., 0.372549, 0.623529],
[0., 0.403922, 0.592157],
[0., 0.435294, 0.560784],
[0., 0.466667, 0.529412],
[0., 0.498039, 0.498039],
[0., 0.529412, 0.466667],
[0., 0.505882, 0.392157],
[0., 0.533333, 0.364706],
[0., 0.623529, 0.372549],
[0., 0.654902, 0.341176],
[0., 0.686275, 0.309804],
[0., 0.717647, 0.278431],
[0., 0.74902, 0.247059],
[0., 0.780392, 0.215686],
[0., 0.811765, 0.184314],
[0., 0.843137, 0.152941],
[0., 0.784314, 0.109804],
[0., 0.811765, 0.0823529],
[0., 0.937255, 0.0588235],
[0., 0.968627, 0.027451],
[0., 1., 0.],
[0.0352941, 1., 0.],
[0.0705882, 1., 0.],
[0.105882, 1., 0.],
[0.141176, 1., 0.],
[0.176471, 1., 0.],
[0.192157, 0.898039, 0.],
[0.223529, 0.898039, 0.],
[0.282353, 1., 0.],
[0.317647, 1., 0.],
[0.356863, 1., 0.],
[0.392157, 1., 0.],
[0.427451, 1., 0.],
[0.462745, 1., 0.],
[0.498039, 1., 0.],
[0.533333, 1., 0.],
[0.513725, 0.898039, 0.],
[0.545098, 0.898039, 0.],
[0.639216, 1., 0.],
[0.678431, 1., 0.],
[0.713725, 1., 0.],
[0.74902, 1., 0.],
[0.784314, 1., 0.],
[0.819608, 1., 0.],
[0.854902, 1., 0.],
[0.890196, 1., 0.],
[0.835294, 0.898039, 0.],
[0.866667, 0.898039, 0.],
[1., 1., 0.],
[1., 0.980392, 0.],
[1., 0.964706, 0.],
[1., 0.94902, 0.],
[1., 0.933333, 0.],
[1., 0.913725, 0.],
[1., 0.898039, 0.],
[1., 0.882353, 0.],
[0.898039, 0.776471, 0.],
[0.898039, 0.764706, 0.],
[1., 0.831373, 0.],
[1., 0.815686, 0.],
[1., 0.8, 0.],
[1., 0.780392, 0.],
[1., 0.764706, 0.],
[1., 0.74902, 0.],
[1., 0.733333, 0.],
[1., 0.713725, 0.],
[0.898039, 0.627451, 0.],
[0.898039, 0.611765, 0.],
[1., 0.662745, 0.],
[1., 0.647059, 0.],
[1., 0.631373, 0.],
[1., 0.615686, 0.],
[1., 0.6, 0.],
[1., 0.580392, 0.],
[1., 0.564706, 0.],
[1., 0.54902, 0.],
[0.898039, 0.478431, 0.],
[0.898039, 0.462745, 0.],
[1., 0.498039, 0.],
[1., 0.490196, 0.],
[1., 0.482353, 0.],
[1., 0.47451, 0.],
[1., 0.466667, 0.],
[1., 0.454902, 0.],
[1., 0.447059, 0.],
[1., 0.439216, 0.],
[0.898039, 0.388235, 0.],
[0.898039, 0.380392, 0.],
[1., 0.415686, 0.],
[1., 0.407843, 0.],
[1., 0.4, 0.],
[1., 0.388235, 0.],
[1., 0.380392, 0.],
[1., 0.372549, 0.],
[1., 0.364706, 0.],
[1., 0.356863, 0.],
[0.898039, 0.313725, 0.],
[0.898039, 0.305882, 0.],
[1., 0.329412, 0.],
[1., 0.321569, 0.],
[1., 0.313725, 0.],
[1., 0.305882, 0.],
[1., 0.298039, 0.],
[1., 0.290196, 0.],
[1., 0.282353, 0.],
[1., 0.27451, 0.],
[0.898039, 0.239216, 0.],
[0.898039, 0.231373, 0.],
[1., 0.247059, 0.],
[1., 0.239216, 0.],
[1., 0.231373, 0.],
[1., 0.223529, 0.],
[1., 0.215686, 0.],
[1., 0.207843, 0.],
[1., 0.196078, 0.],
[1., 0.188235, 0.],
[0.898039, 0.164706, 0.],
[0.898039, 0.156863, 0.],
[1., 0.164706, 0.],
[1., 0.156863, 0.],
[1., 0.14902, 0.],
[1., 0.141176, 0.],
[1., 0.129412, 0.],
[1., 0.121569, 0.],
[1., 0.113725, 0.],
[1., 0.105882, 0.],
[0.898039, 0.0862745, 0.],
[0.898039, 0.0823529, 0.],
[1., 0.0823529, 0.],
[1., 0.0745098, 0.],
[1., 0.0627451, 0.],
[1., 0.054902, 0.],
[1., 0.0470588, 0.],
[1., 0.0509804, 0.],
[1., 0.0313725, 0.],
[1., 0.0235294, 0.],
[0.898039, 0.0117647, 0.],
[0.898039, 0.00392157, 0.],
[1., 0., 0.],
[0.992157, 0., 0.],
[0.984314, 0., 0.],
[0.976471, 0., 0.],
[0.968627, 0., 0.],
[0.960784, 0., 0.],
[0.952941, 0., 0.],
[0.945098, 0., 0.],
[0.843137, 0., 0.],
[0.839216, 0., 0.],
[0.921569, 0., 0.],
[0.917647, 0., 0.],
[0.909804, 0., 0.],
[0.901961, 0., 0.],
[0.894118, 0., 0.],
[0.886275, 0., 0.],
[0.878431, 0., 0.],
[0.870588, 0., 0.],
[0.776471, 0., 0.],
[0.768627, 0., 0.],
[0.847059, 0., 0.],
[0.843137, 0., 0.],
[0.835294, 0., 0.],
[0.827451, 0., 0.],
[0.819608, 0., 0.],
[0.811765, 0., 0.],
[0.803922, 0., 0.],
[0.796078, 0., 0.],
[0.709804, 0., 0.],
[0.701961, 0., 0.],
[0.772549, 0., 0.],
[0.768627, 0., 0.],
[0.760784, 0., 0.],
[0.752941, 0., 0.],
[0.745098, 0., 0.],
[0.737255, 0., 0.],
[0.729412, 0., 0.],
[0.721569, 0., 0.],
[0.643137, 0., 0.],
[0.635294, 0., 0.],
[0.698039, 0., 0.],
[0.690196, 0., 0.],
[0.686275, 0., 0.],
[0.678431, 0., 0.],
[0.670588, 0., 0.],
[0.662745, 0., 0.],
[0.654902, 0., 0.],
[0.647059, 0., 0.],
[0.576471, 0., 0.],
[0.568627, 0., 0.],
[0.623529, 0., 0.],
[0.615686, 0., 0.],
[0.611765, 0., 0.],
[0.603922, 0., 0.],
[0.596078, 0., 0.],
[0.588235, 0., 0.],
[0.580392, 0., 0.],
[0.572549, 0., 0.],
[0.509804, 0., 0.],
[0.501961, 0., 0.],
[0.54902, 0., 0.],
[0.541176, 0., 0.],
[0.537255, 0., 0.],
[0.529412, 0., 0.],
[0.521569, 0., 0.],
[0.513725, 0., 0.],
[0.505882, 0., 0.],
[0.498039, 0., 0.],
[0.443137, 0., 0.],
[0.435294, 0., 0.],
[0.47451, 0., 0.],
[0.466667, 0., 0.],
[0.458824, 0., 0.],
[0.458824, 0., 0.]]

test_cm = LinearSegmentedColormap.from_list(__file__, cm_data)

if __name__ == "__main__":
    import matplotlib.pyplot as plt
    import numpy as np

    try:
        from pycam02ucs.cm.viscm import viscm
        viscm(test_cm)
    except ImportError:
        print("pycam02ucs not found, falling back on simple display")
        plt.imshow(np.linspace(0, 100, 256)[None, :], aspect='auto',
                   cmap=test_cm)
    plt.show()


import matplotlib.pyplot as plt
import numpy as np
import scipy as sp

def trendline(xd, yd, order=1, c='r', alpha=1, plot_r=False, text_pos=None):
    """Make a line of best fit"""

    #Calculate trendline
    coeffs = np.polyfit(xd, yd, order)

    intercept = coeffs[-1]
    slope = coeffs[-2]
    if order == 2: power = coeffs[0]
    else: power = 0

    minxd = np.min(xd)
    maxxd = np.max(xd)

    xl = np.array([minxd, maxxd])
    yl = power * xl ** 2 + slope * xl + intercept

    #Plot trendline
    plt.plot(xl, yl, color=c, alpha=alpha)

    #Calculate R Squared
    r = sp.stats.pearsonr(xd, yd)[0]

    if plot_r == False:
        #Plot R^2 value
        if text_pos == None:
            text_pos = (0.9 * maxxd + 0.1 * minxd, 0.9 * np.max(yd) + 0.1 * np.min(yd),)
        plt.text(text_pos[0], text_pos[1], '$R = %0.2f$' % r)
    else:
        #Return the R^2 value:
        return r

def plot_nice_err(x, y, y_err, color='blue', ls='-', lw=1):
   plt.plot(x, y, color=color, ls=ls, lw=lw)
   plt.fill_between(x, y-y_err, y+y_err, alpha=0.1, color=color)

def plot_word_dist(info, words, start_year, end_year, one_minus=False, legend_loc='upper left'):
    colors = ['r', 'g', 'b', 'c', 'm', 'y', 'k']
    plot_info = {}
    for word in words:
        plot_info[word] = info[word]
    for title, data_dict in plot_info.iteritems():
        x = []; y = []
        for year, val in data_dict.iteritems():
            if year >= start_year and year <= end_year:
                x.append(year)
                if one_minus:
                    val = 1 - val
                y.append(val)
        color = colors.pop()
        plt.plot(x, smooth(np.array(y)), color=color)
        plt.scatter(x, y, marker='.', color=color)
    plt.legend(plot_info.keys(), loc=legend_loc)
    return plt

def get_ccdf(deg_hist, x_min=1):
    cum_counts = [0]
    degs = range(x_min, np.max(deg_hist.keys()))
    total_sum = 0
    for deg in degs:
        if deg in deg_hist:
            deg_count = deg_hist[deg]
        else:
            deg_count = 0
        total_sum += deg_count
        cum_counts.append((cum_counts[-1] + deg_count))
    return np.array(degs), 1 - np.array(cum_counts[1:]) / float(total_sum)

def plot_word_basic(info, words, start_year, end_year, datatype):
    colors = ['r', 'g', 'b', 'c', 'm', 'y', 'k']
    plot_info = {}
    for word in words:
        plot_info[word] = info[word]
    for title, data_dict in plot_info.iteritems():
        x = []; y = []
        for year, val in data_dict[datatype].iteritems():
            if year >= start_year and year <= end_year:
                x.append(year)
                y.append(val)
        color = colors.pop()
        plt.plot(x, smooth(np.array(y)), color=color)
        plt.scatter(x, y, marker='.', color=color)
    plt.legend(plot_info.keys())
    plt.show()

def plot_basic(plot_info, start_year, end_year):
    for title, data_dict in plot_info.iteritems():
        x = []; y = []
        for year, val in data_dict.iteritems():
            if year >= start_year and year <= end_year:
                x.append(year)
                y.append(val)
        plt.plot(x, y)
    plt.legend(plot_info.keys())
    plt.show()

def plot_smooth(x, y, color='blue', window_len=7, window='hanning', ax=None, lw=1.0, ls="-", **kwargs):
    if ax == None:
        _, ax = plt.subplots(1,1)
    ax.plot(x, smooth(np.array(y), window_len=window_len), color=color, lw=lw, ls=ls)
    ax.scatter(x, y, color=color, **kwargs)
    return ax

def smooth(x, window_len=7, window='hanning'):
    """smooth the data using a window with requested size.

    This method is based on the convolution of a scaled window with the signal.
    The signal is prepared by introducing reflected copies of the signal
    (with the window size) in both ends so that transient parts are minimized
    in the begining and end part of the output signal.

    input:
        x: the input signal
        window_len: the dimension of the smoothing window; should be an odd integer
        window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
            flat window will produce a moving average smoothing.

    output:
        the smoothed signal

    example:

    t=linspace(-2,2,0.1)
    x=sin(t)+randn(len(t))*0.1
    y=smooth(x)

    see also:

    numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
    scipy.signal.lfilter

    TODO: the window parameter could be the window itself if an array instead of a string
    NOTE: length(output) != length(input), to correct this: return y[(window_len/2-1):-(window_len/2)] instead of just y.
    """

    if x.ndim != 1:
        raise ValueError, "smooth only accepts 1 dimension arrays."

    if x.size < window_len:
        raise ValueError, "Input vector needs to be bigger than window size."

    if window_len<3:
        return x

    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
        raise ValueError, "Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'"

    s=np.r_[x[window_len-1:0:-1],x,x[-1:-window_len:-1]]
    #print(len(s))
    if window == 'flat': #moving average
        w=np.ones(window_len,'d')
    else:
        w=eval('np.'+window+'(window_len)')

    y=np.convolve(w/w.sum(),s,mode='valid')
    y = y[(window_len/2 - 1):-(window_len/2 + 1)]
    return y


#! /usr/bin/env python
'''
Testsuite for the UniformPDF class
'''
from __future__ import absolute_import, division, print_function

import numpy as np
from puq import *
import scipy.stats as stats

def _hisplot(y, nbins):
    n, bins = np.histogram(y, nbins, normed=True)
    mids = bins[:-1] + np.diff(bins) / 2.0
    return mids, n

def compare_curves(x1, y1, x2, y2, **args):
    ay = np.interp(x2, x1, y1)
    rmse = np.sqrt(np.sum((ay - y2)**2))
    print("maximum difference is", np.max(np.abs(ay - y2)))
    print("RMSE=%s" % rmse)
    # assert rmse < .002
    assert np.allclose(ay, y2, **args)

def _test_updf(min, max):
    options['pdf']['samples'] = 1000

    c = UniformPDF(min=min, max=max)
    assert isinstance(c, PDF)

    x = c.x
    y = stats.uniform(min, max-min).pdf(x)
    rmse = np.sqrt(np.sum((c.y - y)**2))
    print("RMSE=%s" % rmse)
    print("MaxError=", np.max(abs(c.y - y)))
    assert rmse < 1e-11

def _test_ucdf(min, max):
    options['pdf']['samples'] = 1000

    c = UniformPDF(min=min, max=max)

    cdfy = stats.uniform(min, max-min).cdf(c.x)
    rmse = np.sqrt(np.sum((c.cdfy - cdfy)**2))
    print("RMSE=%s" % rmse)
    print("MaxError=", np.max(abs(c.cdfy - cdfy)))
    assert rmse < 1e-11

    """
    import matplotlib.pyplot as plt
    plt.plot(c.x, c.cdfy, color='green')
    plt.plot(c.x, cdfy, color='red')
    plt.show()
    """

# test mean, min, max and deviation
def _test_uniform_minmeanmax(min, mean, max):
    c = UniformPDF(min=min, mean=mean, max=max)
    cmin, cmax = c.range
    print("min=%s mean=%s  max=%s" % (cmin, c.mean, cmax))

    if min is not None:
        assert min == cmin
    else:
        assert cmin == mean - (max - mean)
    if max is not None:
        assert max == cmax
    else:
        assert cmax == mean + (mean - min)
    if mean is not None:
        assert np.allclose(mean, c.mean)
    else:
        assert np.allclose(c.mean, (min + max) / 2.0)

# test lhs()
def _test_uniform_lhs(min, max):
    c = UniformPDF(min=min, max=max)

    # test the lhs() function to see if the curve it generates is
    # close enough
    data = c.ds(10000)
    assert len(data) == 10000
    assert np.min(data) >= min
    assert np.max(data) <= max
    dx, dy = _hisplot(data, 20)

    x = dx
    y = stats.uniform(min, max-min).pdf(x)
    compare_curves(x, y, dx, dy, atol=.0001)

    """
    import matplotlib.pyplot as plt
    plt.plot(x, y, color='red')
    plt.plot(dx, dy, color='blue')
    plt.show()
    """

    assert np.allclose(c.mean, np.mean(data), rtol=.001), 'mean=%s' % np.mean(data)

# test lhs1()
def _test_uniform_lhs1(min, max):
    c = UniformPDF(min=min, max=max)

    data = c.ds1(1000)
    xs = data
    assert len(xs) == 1000

    assert min, max == c.range
    # scale [-1,1] back to original size

    mean = (min + max)/2.0
    xs *= max - mean
    xs += mean
    dx, dy = _hisplot(xs, 20)

    x = dx
    y = stats.uniform(min, max-min).pdf(x)
    compare_curves(x, y, dx, dy, atol=.001)

    """
    import matplotlib.pyplot as plt
    plt.plot(x, y, color='green')
    plt.plot(dx, dy, color='blue')
    plt.show()
    """
    assert np.allclose(c.mean, np.mean(data), rtol=.001), 'mean=%s' % np.mean(data)

def _test_uniform_random(min, max):
    c = UniformPDF(min=min, max=max)

    data = c.random(1000000)
    assert len(data) == 1000000
    dx, dy = _hisplot(data, 20)

    x = dx
    y = stats.uniform(min, max-min).pdf(x)
    compare_curves(x, y, dx, dy, atol=.02)

    assert np.min(data) >= min
    assert np.max(data) <= max

    """
    import matplotlib.pyplot as plt
    plt.plot(dx, dy, color='red')
    plt.plot(x, y, color='blue')
    plt.show()
    """
    assert np.allclose(c.mean, np.mean(data), rtol=.001), 'mean=%s' % np.mean(data)

def test_updf():
    _test_updf(10,20)
    _test_updf(-20,-10)

def test_ucdf():
    _test_ucdf(100,105)
    _test_ucdf(-1,2)

def test_uniform_minmeanmax():
    _test_uniform_minmeanmax(0,None,20)
    _test_uniform_minmeanmax(None,0.5,2)
    _test_uniform_minmeanmax(5,10,15)
    _test_uniform_minmeanmax(5,10,None)

def test_uniform_lhs():
    _test_uniform_lhs(10,20)
    _test_uniform_lhs(-100, -50)

def test_uniform_lhs1():
    _test_uniform_lhs1(10,20)
    _test_uniform_lhs1(-100, -50)

def test_uniform_random():
    _test_uniform_random(10,20)

if __name__ == "__main__":
    test_updf()
    test_ucdf()
    test_uniform_minmeanmax()
    test_uniform_lhs()
    test_uniform_lhs1()
    test_uniform_random()