"""Benchmarks of Lasso regularization path computation using Lars and CD

The input data is mostly low rank but is a fat infinite tail.
"""
from __future__ import print_function

from collections import defaultdict
import gc
import sys
from time import time

import numpy as np

from sklearn.linear_model import lars_path
from sklearn.linear_model import lasso_path
from sklearn.datasets.samples_generator import make_regression

def compute_bench(samples_range, features_range):

    it = 0

    results = defaultdict(lambda: [])

    max_it = len(samples_range) * len(features_range)
    for n_samples in samples_range:
        for n_features in features_range:
            it += 1
            print('====================')
            print('Iteration %03d of %03d' % (it, max_it))
            print('====================')
            dataset_kwargs = {
                'n_samples': n_samples,
                'n_features': n_features,
                'n_informative': n_features / 10,
                'effective_rank': min(n_samples, n_features) / 10,
                #'effective_rank': None,
                'bias': 0.0,
            }
            print("n_samples: %d" % n_samples)
            print("n_features: %d" % n_features)
            X, y = make_regression(**dataset_kwargs)

            gc.collect()
            print("benchmarking lars_path (with Gram):", end='')
            sys.stdout.flush()
            tstart = time()
            G = np.dot(X.T, X)  # precomputed Gram matrix
            Xy = np.dot(X.T, y)
            lars_path(X, y, Xy=Xy, Gram=G, method='lasso')
            delta = time() - tstart
            print("%0.3fs" % delta)
            results['lars_path (with Gram)'].append(delta)

            gc.collect()
            print("benchmarking lars_path (without Gram):", end='')
            sys.stdout.flush()
            tstart = time()
            lars_path(X, y, method='lasso')
            delta = time() - tstart
            print("%0.3fs" % delta)
            results['lars_path (without Gram)'].append(delta)

            gc.collect()
            print("benchmarking lasso_path (with Gram):", end='')
            sys.stdout.flush()
            tstart = time()
            lasso_path(X, y, precompute=True)
            delta = time() - tstart
            print("%0.3fs" % delta)
            results['lasso_path (with Gram)'].append(delta)

            gc.collect()
            print("benchmarking lasso_path (without Gram):", end='')
            sys.stdout.flush()
            tstart = time()
            lasso_path(X, y, precompute=False)
            delta = time() - tstart
            print("%0.3fs" % delta)
            results['lasso_path (without Gram)'].append(delta)

    return results

if __name__ == '__main__':
    from mpl_toolkits.mplot3d import axes3d  # register the 3d projection
    import matplotlib.pyplot as plt

    samples_range = np.linspace(10, 2000, 5).astype(np.int)
    features_range = np.linspace(10, 2000, 5).astype(np.int)
    results = compute_bench(samples_range, features_range)

    max_time = max(max(t) for t in results.values())

    fig = plt.figure('scikit-learn Lasso path benchmark results')
    i = 1
    for c, (label, timings) in zip('bcry', sorted(results.items())):
        ax = fig.add_subplot(2, 2, i, projection='3d')
        X, Y = np.meshgrid(samples_range, features_range)
        Z = np.asarray(timings).reshape(samples_range.shape[0],
                                        features_range.shape[0])

        # plot the actual surface
        ax.plot_surface(X, Y, Z.T, cstride=1, rstride=1, color=c, alpha=0.8)

        # dummy point plot to stick the legend to since surface plot do not
        # support legends (yet?)
        # ax.plot([1], [1], [1], color=c, label=label)

        ax.set_xlabel('n_samples')
        ax.set_ylabel('n_features')
        ax.set_zlabel('Time (s)')
        ax.set_zlim3d(0.0, max_time * 1.1)
        ax.set_title(label)
        # ax.legend()
        i += 1
    plt.show()


from datetime import datetime as dt, timedelta

import bson
import numpy as np
from mock import patch

from arctic._util import mongo_count
from arctic.arctic import Arctic

def test_save_read_bson(library):
    blob = {'foo': dt(2015, 1, 1), 'bar': ['a', 'b', ['x', 'y', 'z']]}
    library.write('BLOB', blob)
    saved_blob = library.read('BLOB').data
    assert blob == saved_blob

'''
Run test at your own discretion. Takes > 60 secs
def test_save_read_MASSIVE(library):
    import pandas as pd
    df = pd.DataFrame(data={'data': [1] * 150000000})
    data = (df, df)
    library.write('BLOB', data)
    saved_blob = library.read('BLOB').data
    assert(saved_blob[0].equals(df))
    assert(saved_blob[1].equals(df))
'''

def test_save_read_big_encodable(library):
    blob = {'foo': 'a' * 1024 * 1024 * 20}
    library.write('BLOB', blob)
    saved_blob = library.read('BLOB').data
    assert blob == saved_blob

def test_save_read_bson_object(library):
    blob = {'foo': dt(2015, 1, 1), 'object': Arctic}
    library.write('BLOB', blob)
    saved_blob = library.read('BLOB').data
    assert blob == saved_blob

def test_get_info_bson_object(library):
    blob = {'foo': dt(2015, 1, 1), 'object': Arctic}
    library.write('BLOB', blob)
    assert library.get_info('BLOB')['handler'] == 'PickleStore'

def test_bson_large_object(library):
    blob = {'foo': dt(2015, 1, 1), 'object': Arctic,
            'large_thing': np.random.rand(int(2.1 * 1024 * 1024)).tostring()}
    assert len(blob['large_thing']) > 16 * 1024 * 1024
    library.write('BLOB', blob)
    saved_blob = library.read('BLOB').data
    assert blob == saved_blob

def test_bson_leak_objects_delete(library):
    blob = {'foo': dt(2015, 1, 1), 'object': Arctic}
    library.write('BLOB', blob)
    assert mongo_count(library._collection) == 1
    assert mongo_count(library._collection.versions) == 1
    library.delete('BLOB')
    assert mongo_count(library._collection) == 0
    assert mongo_count(library._collection.versions) == 0

def test_bson_leak_objects_prune_previous(library):
    blob = {'foo': dt(2015, 1, 1), 'object': Arctic}

    yesterday = dt.utcnow() - timedelta(days=1, seconds=1)
    _id = bson.ObjectId.from_datetime(yesterday)
    with patch("bson.ObjectId", return_value=_id):
        library.write('BLOB', blob)
    assert mongo_count(library._collection) == 1
    assert mongo_count(library._collection.versions) == 1

    _id = bson.ObjectId.from_datetime(dt.utcnow() - timedelta(minutes=130))
    with patch("bson.ObjectId", return_value=_id):
        library.write('BLOB', {}, prune_previous_version=False)
    assert mongo_count(library._collection) == 1
    assert mongo_count(library._collection.versions) == 2

    # This write should pruned the oldest version in the chunk collection
    library.write('BLOB', {})
    assert mongo_count(library._collection) == 0
    assert mongo_count(library._collection.versions) == 2

def test_prune_previous_doesnt_kill_other_objects(library):
    blob = {'foo': dt(2015, 1, 1), 'object': Arctic}

    yesterday = dt.utcnow() - timedelta(days=1, seconds=1)
    _id = bson.ObjectId.from_datetime(yesterday)
    with patch("bson.ObjectId", return_value=_id):
        library.write('BLOB', blob, prune_previous_version=False)
    assert mongo_count(library._collection) == 1
    assert mongo_count(library._collection.versions) == 1

    _id = bson.ObjectId.from_datetime(dt.utcnow() - timedelta(hours=10))
    with patch("bson.ObjectId", return_value=_id):
        library.write('BLOB', blob, prune_previous_version=False)
    assert mongo_count(library._collection) == 1
    assert mongo_count(library._collection.versions) == 2

    # This write should pruned the oldest version in the chunk collection
    library.write('BLOB', {})
    assert mongo_count(library._collection) == 1
    assert mongo_count(library._collection.versions) == 2

    library._delete_version('BLOB', 2)
    assert mongo_count(library._collection) == 0
    assert mongo_count(library._collection.versions) == 1

def test_write_metadata(library):
    blob = {'foo': dt(2015, 1, 1), 'object': Arctic}
    library.write(symbol='symX', data=blob, metadata={'key1': 'value1'})
    library.write_metadata(symbol='symX', metadata={'key2': 'value2'})
    v = library.read('symX')
    assert v.data == blob
    assert v.metadata == {'key2': 'value2'}


"""Unsupervised nearest neighbors learner"""

from .base import NeighborsBase
from .base import KNeighborsMixin
from .base import RadiusNeighborsMixin
from .base import UnsupervisedMixin

class NearestNeighbors(NeighborsBase, KNeighborsMixin,
                       RadiusNeighborsMixin, UnsupervisedMixin):
    """Unsupervised learner for implementing neighbor searches.

    Read more in the :ref:`User Guide <unsupervised_neighbors>`.

    Parameters
    ----------
    n_neighbors : int, optional (default = 5)
        Number of neighbors to use by default for :meth:`k_neighbors` queries.

    radius : float, optional (default = 1.0)
        Range of parameter space to use by default for :meth`radius_neighbors`
        queries.

    algorithm : {'auto', 'ball_tree', 'kd_tree', 'brute'}, optional
        Algorithm used to compute the nearest neighbors:

        - 'ball_tree' will use :class:`BallTree`
        - 'kd_tree' will use :class:`KDtree`
        - 'brute' will use a brute-force search.
        - 'auto' will attempt to decide the most appropriate algorithm
          based on the values passed to :meth:`fit` method.

        Note: fitting on sparse input will override the setting of
        this parameter, using brute force.

    leaf_size : int, optional (default = 30)
        Leaf size passed to BallTree or KDTree.  This can affect the
        speed of the construction and query, as well as the memory
        required to store the tree.  The optimal value depends on the
        nature of the problem.

    p: integer, optional (default = 2)
        Parameter for the Minkowski metric from
        sklearn.metrics.pairwise.pairwise_distances. When p = 1, this is
        equivalent to using manhattan_distance (l1), and euclidean_distance
        (l2) for p = 2. For arbitrary p, minkowski_distance (l_p) is used.

    metric : string or callable, default 'minkowski'
        metric to use for distance computation. Any metric from scikit-learn
        or scipy.spatial.distance can be used.

        If metric is a callable function, it is called on each
        pair of instances (rows) and the resulting value recorded. The callable
        should take two arrays as input and return one value indicating the
        distance between them. This works for Scipy's metrics, but is less
        efficient than passing the metric name as a string.

        Distance matrices are not supported.

        Valid values for metric are:

        - from scikit-learn: ['cityblock', 'cosine', 'euclidean', 'l1', 'l2',
          'manhattan']

        - from scipy.spatial.distance: ['braycurtis', 'canberra', 'chebyshev',
          'correlation', 'dice', 'hamming', 'jaccard', 'kulsinski',
          'mahalanobis', 'matching', 'minkowski', 'rogerstanimoto',
          'russellrao', 'seuclidean', 'sokalmichener', 'sokalsneath',
          'sqeuclidean', 'yule']

        See the documentation for scipy.spatial.distance for details on these
        metrics.

    metric_params: dict, optional (default = None)
        additional keyword arguments for the metric function.

    Examples
    --------
      >>> import numpy as np
      >>> from sklearn.neighbors import NearestNeighbors
      >>> samples = [[0, 0, 2], [1, 0, 0], [0, 0, 1]]

      >>> neigh = NearestNeighbors(2, 0.4)
      >>> neigh.fit(samples)  #doctest: +ELLIPSIS
      NearestNeighbors(...)

      >>> neigh.kneighbors([[0, 0, 1.3]], 2, return_distance=False)
      ... #doctest: +ELLIPSIS
      array([[2, 0]]...)

      >>> rng = neigh.radius_neighbors([0, 0, 1.3], 0.4, return_distance=False)
      >>> np.asarray(rng[0][0])
      array(2)

    See also
    --------
    KNeighborsClassifier
    RadiusNeighborsClassifier
    KNeighborsRegressor
    RadiusNeighborsRegressor
    BallTree

    Notes
    -----
    See :ref:`Nearest Neighbors <neighbors>` in the online documentation
    for a discussion of the choice of ``algorithm`` and ``leaf_size``.

    http://en.wikipedia.org/wiki/K-nearest_neighbor_algorithm
    """

    def __init__(self, n_neighbors=5, radius=1.0,
                 algorithm='auto', leaf_size=30, metric='minkowski',
                 p=2, metric_params=None, **kwargs):
        self._init_params(n_neighbors=n_neighbors,
                          radius=radius,
                          algorithm=algorithm,
                          leaf_size=leaf_size, metric=metric, p=p,
                          metric_params=metric_params, **kwargs)


import re, random, math
import h2o_args
import h2o_nodes
import h2o_cmd
from h2o_test import verboseprint, dump_json, check_sandbox_for_errors

def plotLists(xList, xLabel=None, eListTitle=None, eList=None, eLabel=None, fListTitle=None, fList=None, fLabel=None, server=False):
    if h2o_args.python_username!='kevin':
        return

    # Force matplotlib to not use any Xwindows backend.
    if server:
        import matplotlib
        matplotlib.use('Agg')

    import pylab as plt
    print "xList", xList
    print "eList", eList
    print "fList", fList

    font = {'family' : 'normal',
            'weight' : 'normal',
            'size'   : 26}
    ### plt.rc('font', **font)
    plt.rcdefaults()

    if eList:
        if eListTitle:
            plt.title(eListTitle)
        plt.figure()
        plt.plot (xList, eList)
        plt.xlabel(xLabel)
        plt.ylabel(eLabel)
        plt.draw()
        plt.savefig('eplot.jpg',format='jpg')
        # Image.open('testplot.jpg').save('eplot.jpg','JPEG')

    if fList:
        if fListTitle:
            plt.title(fListTitle)
        plt.figure()
        plt.plot (xList, fList)
        plt.xlabel(xLabel)
        plt.ylabel(fLabel)
        plt.draw()
        plt.savefig('fplot.jpg',format='jpg')
        # Image.open('fplot.jpg').save('fplot.jpg','JPEG')

    if eList or fList:
        plt.show()

# pretty print a cm that the C
def pp_cm(jcm, header=None):
    # header = jcm['header']
    # hack col index header for now..where do we get it?
    header = ['"%s"'%i for i in range(len(jcm[0]))]
    # cm = '   '.join(header)
    cm = '{0:<8}'.format('')
    for h in header:
        cm = '{0}|{1:<8}'.format(cm, h)

    cm = '{0}|{1:<8}'.format(cm, 'error')
    c = 0
    for line in jcm:
        lineSum  = sum(line)
        if c < 0 or c >= len(line):
            raise Exception("Error in h2o_gbm.pp_cm. c: %s line: %s len(line): %s jcm: %s" % (c, line, len(line), dump_json(jcm)))
        print "c:", c, "line:", line
        errorSum = lineSum - line[c]
        if (lineSum>0):
            err = float(errorSum) / lineSum
        else:
            err = 0.0
        fl = '{0:<8}'.format(header[c])
        for num in line: fl = '{0}|{1:<8}'.format(fl, num)
        fl = '{0}|{1:<8.2f}'.format(fl, err)
        cm = "{0}\n{1}".format(cm, fl)
        c += 1
    return cm

def pp_cm_summary(cm):
    # hack cut and past for now (should be in h2o_gbm.py?
    scoresList = cm
    totalScores = 0
    totalRight = 0
    # individual scores can be all 0 if nothing for that output class
    # due to sampling
    classErrorPctList = []
    predictedClassDict = {} # may be missing some? so need a dict?
    for classIndex,s in enumerate(scoresList):
        classSum = sum(s)
        if classSum == 0 :
            # why would the number of scores for a class be 0?
            # in any case, tolerate. (it shows up in test.py on poker100)
            print "classIndex:", classIndex, "classSum", classSum, "<- why 0?"
        else:
            if classIndex >= len(s):
                print "Why is classindex:", classIndex, 'for s:"', s
            else:
            # H2O should really give me this since it's in the browser, but it doesn't
                classRightPct = ((s[classIndex] + 0.0)/classSum) * 100
                totalRight += s[classIndex]
                classErrorPct = 100 - classRightPct
                classErrorPctList.append(classErrorPct)
                ### print "s:", s, "classIndex:", classIndex
                print "class:", classIndex, "classSum", classSum, "classErrorPct:", "%4.2f" % classErrorPct

                # gather info for prediction summary
                for pIndex,p in enumerate(s):
                    if pIndex not in predictedClassDict:
                        predictedClassDict[pIndex] = p
                    else:
                        predictedClassDict[pIndex] += p

        totalScores += classSum

    print "Predicted summary:"
    # FIX! Not sure why we weren't working with a list..hack with dict for now
    for predictedClass,p in predictedClassDict.items():
        print str(predictedClass)+":", p

    # this should equal the num rows in the dataset if full scoring? (minus any NAs)
    print "totalScores:", totalScores
    print "totalRight:", totalRight
    if totalScores != 0:  pctRight = 100.0 * totalRight/totalScores
    else: pctRight = 0.0
    print "pctRight:", "%5.2f" % pctRight
    pctWrong = 100 - pctRight
    print "pctWrong:", "%5.2f" % pctWrong

    return pctWrong

# I just copied and changed GBM to GBM. Have to update to match GBM params and responses

def pickRandGbmParams(paramDict, params):
    colX = 0
    randomGroupSize = random.randint(1,len(paramDict))
    for i in range(randomGroupSize):
        randomKey = random.choice(paramDict.keys())
        randomV = paramDict[randomKey]
        randomValue = random.choice(randomV)
        params[randomKey] = randomValue

# compare this glm to last one. since the files are concatenations,
# the results should be similar? 10% of first is allowed delta
def compareToFirstGbm(self, key, glm, firstglm):
    # if isinstance(firstglm[key], list):
    # in case it's not a list allready (err is a list)
    verboseprint("compareToFirstGbm key:", key)
    verboseprint("compareToFirstGbm glm[key]:", glm[key])
    # key could be a list or not. if a list, don't want to create list of that list
    # so use extend on an empty list. covers all cases?
    if type(glm[key]) is list:
        kList  = glm[key]
        firstkList = firstglm[key]
    elif type(glm[key]) is dict:
        raise Exception("compareToFirstGLm: Not expecting dict for " + key)
    else:
        kList  = [glm[key]]
        firstkList = [firstglm[key]]

    for k, firstk in zip(kList, firstkList):
        # delta must be a positive number ?
        delta = .1 * abs(float(firstk))
        msg = "Too large a delta (" + str(delta) + ") comparing current and first for: " + key
        self.assertAlmostEqual(float(k), float(firstk), delta=delta, msg=msg)
        self.assertGreaterEqual(abs(float(k)), 0.0, str(k) + " abs not >= 0.0 in current")

def goodXFromColumnInfo(y,
    num_cols=None, missingValuesDict=None, constantValuesDict=None, enumSizeDict=None,
    colTypeDict=None, colNameDict=None, keepPattern=None, key=None,
    timeoutSecs=120, forRF=False, noPrint=False):

    y = str(y)

    # if we pass a key, means we want to get the info ourselves here
    if key is not None:
        (missingValuesDict, constantValuesDict, enumSizeDict, colTypeDict, colNameDict) = \
            h2o_cmd.columnInfoFromInspect(key, exceptionOnMissingValues=False,
            max_column_display=99999999, timeoutSecs=timeoutSecs)
        num_cols = len(colNameDict)

    # now remove any whose names don't match the required keepPattern
    if keepPattern is not None:
        keepX = re.compile(keepPattern)
    else:
        keepX = None

    x = range(num_cols)
    # need to walk over a copy, cause we change x
    xOrig = x[:]
    ignore_x = [] # for use by RF
    for k in xOrig:
        name = colNameDict[k]
        # remove it if it has the same name as the y output
        if str(k)== y: # if they pass the col index as y
            if not noPrint:
                print "Removing %d because name: %s matches output %s" % (k, str(k), y)
            x.remove(k)
            # rf doesn't want it in ignore list
            # ignore_x.append(k)
        elif name == y: # if they pass the name as y
            if not noPrint:
                print "Removing %d because name: %s matches output %s" % (k, name, y)
            x.remove(k)
            # rf doesn't want it in ignore list
            # ignore_x.append(k)

        elif keepX is not None and not keepX.match(name):
            if not noPrint:
                print "Removing %d because name: %s doesn't match desired keepPattern %s" % (k, name, keepPattern)
            x.remove(k)
            ignore_x.append(k)

        # missing values reports as constant also. so do missing first.
        # remove all cols with missing values
        # could change it against num_rows for a ratio
        elif k in missingValuesDict:
            value = missingValuesDict[k]
            if not noPrint:
                print "Removing %d with name: %s because it has %d missing values" % (k, name, value)
            x.remove(k)
            ignore_x.append(k)

        elif k in constantValuesDict:
            value = constantValuesDict[k]
            if not noPrint:
                print "Removing %d with name: %s because it has constant value: %s " % (k, name, str(value))
            x.remove(k)
            ignore_x.append(k)

        # this is extra pruning..
        # remove all cols with enums, if not already removed
        elif k in enumSizeDict:
            value = enumSizeDict[k]
            if not noPrint:
                print "Removing %d %s because it has enums of size: %d" % (k, name, value)
            x.remove(k)
            ignore_x.append(k)

    if not noPrint:
        print "x has", len(x), "cols"
        print "ignore_x has", len(ignore_x), "cols"
    x = ",".join(map(str,x))
    ignore_x = ",".join(map(str,ignore_x))

    if not noPrint:
        print "\nx:", x
        print "\nignore_x:", ignore_x

    if forRF:
        return ignore_x
    else:
        return x

def showGBMGridResults(GBMResult, expectedErrorMax, classification=True):
    # print "GBMResult:", dump_json(GBMResult)
    jobs = GBMResult['jobs']
    print "GBM jobs:", jobs
    for jobnum, j in enumerate(jobs):
        _distribution = j['_distribution']
        model_key = j['destination_key']
        job_key = j['job_key']
        # inspect = h2o_cmd.runInspect(key=model_key)
        # print "jobnum:", jobnum, dump_json(inspect)
        gbmTrainView = h2o_cmd.runGBMView(model_key=model_key)
        print "jobnum:", jobnum, dump_json(gbmTrainView)

        if classification:
            cms = gbmTrainView['gbm_model']['cms']
            cm = cms[-1]['_arr'] # take the last one
            print "GBM cms[-1]['_predErr']:", cms[-1]['_predErr']
            print "GBM cms[-1]['_classErr']:", cms[-1]['_classErr']

            pctWrongTrain = pp_cm_summary(cm);
            if pctWrongTrain > expectedErrorMax:
                raise Exception("Should have < %s error here. pctWrongTrain: %s" % (expectedErrorMax, pctWrongTrain))

            errsLast = gbmTrainView['gbm_model']['errs'][-1]
            print "\nTrain", jobnum, job_key, "\n==========\n", "pctWrongTrain:", pctWrongTrain, "errsLast:", errsLast
            print "GBM 'errsLast'", errsLast
            print pp_cm(cm)
        else:
            print "\nTrain", jobnum, job_key, "\n==========\n", "errsLast:", errsLast
            print "GBMTrainView errs:", gbmTrainView['gbm_model']['errs']

def simpleCheckGBMView(node=None, gbmv=None, noPrint=False, **kwargs):
    if not node:
        node = h2o_nodes.nodes[0]

    if 'warnings' in gbmv:
        warnings = gbmv['warnings']
        # catch the 'Failed to converge" for now
        for w in warnings:
            if not noPrint: print "\nwarning:", w
            if ('Failed' in w) or ('failed' in w):
                raise Exception(w)

    if 'cm' in gbmv:
        cm = gbmv['cm'] # only one
    else:
        if 'gbm_model' in gbmv:
            gbm_model = gbmv['gbm_model']
        else:
            raise Exception("no gbm_model in gbmv? %s" % dump_json(gbmv))

        cms = gbm_model['cms']
        print "number of cms:", len(cms)
        print "FIX! need to add reporting of h2o's _perr per class error"
        # FIX! what if regression. is rf only classification?
        print "cms[-1]['_arr']:", cms[-1]['_arr']
        print "cms[-1]['_predErr']:", cms[-1]['_predErr']
        print "cms[-1]['_classErr']:", cms[-1]['_classErr']

        ## print "cms[-1]:", dump_json(cms[-1])
        ## for i,c in enumerate(cms):
        ##    print "cm %s: %s" % (i, c['_arr'])

        cm = cms[-1]['_arr'] # take the last one

    scoresList = cm

    used_trees = gbm_model['N']
    errs = gbm_model['errs']
    print "errs[0]:", errs[0]
    print "errs[-1]:", errs[-1]
    print "errs:", errs
    # if we got the ntree for comparison. Not always there in kwargs though!
    param_ntrees = kwargs.get('ntrees',None)
    if (param_ntrees is not None and used_trees != param_ntrees):
        raise Exception("used_trees should == param_ntree. used_trees: %s"  % used_trees)
    if (used_trees+1)!=len(cms) or (used_trees+1)!=len(errs):
        raise Exception("len(cms): %s and len(errs): %s should be one more than N %s trees" % (len(cms), len(errs), used_trees))

    totalScores = 0
    totalRight = 0
    # individual scores can be all 0 if nothing for that output class
    # due to sampling
    classErrorPctList = []
    predictedClassDict = {} # may be missing some? so need a dict?
    for classIndex,s in enumerate(scoresList):
        classSum = sum(s)
        if classSum == 0 :
            # why would the number of scores for a class be 0? does GBM CM have entries for non-existent classes
            # in a range??..in any case, tolerate. (it shows up in test.py on poker100)
            if not noPrint: print "class:", classIndex, "classSum", classSum, "<- why 0?"
        else:
            # H2O should really give me this since it's in the browser, but it doesn't
            classRightPct = ((s[classIndex] + 0.0)/classSum) * 100
            totalRight += s[classIndex]
            classErrorPct = round(100 - classRightPct, 2)
            classErrorPctList.append(classErrorPct)
            ### print "s:", s, "classIndex:", classIndex
            if not noPrint: print "class:", classIndex, "classSum", classSum, "classErrorPct:", "%4.2f" % classErrorPct

            # gather info for prediction summary
            for pIndex,p in enumerate(s):
                if pIndex not in predictedClassDict:
                    predictedClassDict[pIndex] = p
                else:
                    predictedClassDict[pIndex] += p

        totalScores += classSum

    #****************************
    if not noPrint:
        print "Predicted summary:"
        # FIX! Not sure why we weren't working with a list..hack with dict for now
        for predictedClass,p in predictedClassDict.items():
            print str(predictedClass)+":", p

        # this should equal the num rows in the dataset if full scoring? (minus any NAs)
        print "totalScores:", totalScores
        print "totalRight:", totalRight
        if totalScores != 0:
            pctRight = 100.0 * totalRight/totalScores
        else:
            pctRight = 0.0
        pctWrong = 100 - pctRight
        print "pctRight:", "%5.2f" % pctRight
        print "pctWrong:", "%5.2f" % pctWrong

    #****************************
    # more testing for GBMView
    # it's legal to get 0's for oobe error # if sample_rate = 1

    sample_rate = kwargs.get('sample_rate', None)
    validation = kwargs.get('validation', None)
    if (sample_rate==1 and not validation):
        pass
    elif (totalScores<=0 or totalScores>5e9):
        raise Exception("scores in GBMView seems wrong. scores:", scoresList)

    varimp = gbm_model['varimp']
    treeStats = gbm_model['treeStats']
    if not treeStats:
        raise Exception("treeStats not right?: %s" % dump_json(treeStats))
    # print "json:", dump_json(gbmv)
    data_key = gbm_model['_dataKey']
    model_key = gbm_model['_key']
    classification_error = pctWrong

    if not noPrint:
        if 'minLeaves' not in treeStats or not treeStats['minLeaves']:
            raise Exception("treeStats seems to be missing minLeaves %s" % dump_json(treeStats))
        print """
         Leaves: {0} / {1} / {2}
          Depth: {3} / {4} / {5}
            Err: {6:0.2f} %
        """.format(
                treeStats['minLeaves'],
                treeStats['meanLeaves'],
                treeStats['maxLeaves'],
                treeStats['minDepth'],
                treeStats['meanDepth'],
                treeStats['maxDepth'],
                classification_error,
                )

    ### modelInspect = node.inspect(model_key)
    dataInspect = h2o_cmd.runInspect(key=data_key)
    check_sandbox_for_errors()
    return (round(classification_error,2), classErrorPctList, totalScores)


__author__ = 'robert'

from pypet import Environment, Trajectory
from pypet.tests.testutils.ioutils import make_temp_dir, get_log_config
import os
import matplotlib.pyplot as plt
import numpy as np
import time

import numpy as np
import scipy.sparse as spsp
from pycallgraph import PyCallGraph, Config, GlobbingFilter
from pycallgraph.output import GraphvizOutput
from pycallgraph.color import Color

class CustomOutput(GraphvizOutput):
    def node_color(self, node):
        value = float(node.time.fraction)
        return Color.hsv(value / 2 + .5, value, 0.9)

    def edge_color(self, edge):
        value = float(edge.time.fraction)
        return Color.hsv(value / 2 + .5, value, 0.7)

def job(traj):
    traj.f_ares('$set.$', 42, comment='A result')

def get_runtime(length):
    filename = os.path.join('tmp', 'hdf5', 'many_runs.hdf5')

    with Environment(filename = filename,
                      log_levels=20, report_progress=(0.0000002, 'progress', 50),
                      overwrite_file=True, purge_duplicate_comments=False,
                      log_stdout=False,
                      summary_tables=False, small_overview_tables=False) as env:

        traj = env.v_traj

        traj.par.f_apar('x', 0, 'parameter')

        traj.f_explore({'x': range(length)})

        max_run = 100

        for idx in range(len(traj)):
            if idx > max_run:
                traj.f_get_run_information(idx, copy=False)['completed'] = 1
        traj.f_store()

        if not os.path.isdir('./tmp'):
            os.mkdir('tmp')
        graphviz = CustomOutput()
        graphviz.output_file = './tmp/run_profile_storage_%d.png' % len(traj)
        service_filter = GlobbingFilter(include=['*storageservice.*'])

        config = Config(groups=True, verbose=True)
        config.trace_filter = service_filter

        print('RUN PROFILE')
        with PyCallGraph(config=config, output=graphviz):
            # start = time.time()
            # env.f_run(job)
            # end = time.time()
            for irun in range(100):
                traj._make_single_run(irun+len(traj)/2)
                # Measure start time
                traj._set_start()
                traj.f_ares('$set.$', 42, comment='A result')
                traj._set_finish()
                traj._store_final(store_data=2)
                traj._finalize_run()
            print('STARTING_to_PLOT')
        print('DONE RUN PROFILE')

        # dicts = [traj.f_get_run_information(x) for x in range(min(len(traj), max_run))]
    # total = end - start
    # return total/float(min(len(traj), max_run)), total/float(min(len(traj), max_run)) * len(traj)

def main():
    lengths = [1000, 1000000]
    runtimes = [get_runtime(x) for x in lengths]
    # avg_runtimes = [x[0] for x in runtimes]
    # summed_runtime = [x[1] for x in runtimes]

    # plt.subplot(2, 1, 1)
    # plt.semilogx(list(reversed(lengths)), list(reversed(avg_runtimes)), linewidth=2)
    # plt.xlabel('Runs')
    # plt.ylabel('t[s]')
    # plt.title('Average Runtime per single run')
    # plt.grid()
    # plt.subplot(2, 1, 2)
    # plt.loglog(lengths, summed_runtime, linewidth=2)
    # plt.grid()
    # plt.xlabel('Runs')
    # plt.ylabel('t[s]')
    # plt.title('Total runtime of experiment')
    # plt.savefig('avg_runtime_as_func_of_lenght_100')
    # plt.show()

if __name__ == '__main__':
    main()


"""
Testing for the nearest centroid module.
"""

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

from sklearn.neighbors import NearestCentroid
from sklearn import datasets
from sklearn.metrics.pairwise import pairwise_distances

# toy sample
X = [[-2, -1], [-1, -1], [-1, -2], [1, 1], [1, 2], [2, 1]]
X_csr = sp.csr_matrix(X)  # Sparse matrix
y = [-1, -1, -1, 1, 1, 1]
T = [[-1, -1], [2, 2], [3, 2]]
T_csr = sp.csr_matrix(T)
true_result = [-1, 1, 1]

# also load the iris dataset
# and randomly permute it
iris = datasets.load_iris()
rng = np.random.RandomState(1)
perm = rng.permutation(iris.target.size)
iris.data = iris.data[perm]
iris.target = iris.target[perm]

def test_classification_toy():
    # Check classification on a toy dataset, including sparse versions.
    clf = NearestCentroid()
    clf.fit(X, y)
    assert_array_equal(clf.predict(T), true_result)

    # Same test, but with a sparse matrix to fit and test.
    clf = NearestCentroid()
    clf.fit(X_csr, y)
    assert_array_equal(clf.predict(T_csr), true_result)

    # Fit with sparse, test with non-sparse
    clf = NearestCentroid()
    clf.fit(X_csr, y)
    assert_array_equal(clf.predict(T), true_result)

    # Fit with non-sparse, test with sparse
    clf = NearestCentroid()
    clf.fit(X, y)
    assert_array_equal(clf.predict(T_csr), true_result)

    # Fit and predict with non-CSR sparse matrices
    clf = NearestCentroid()
    clf.fit(X_csr.tocoo(), y)
    assert_array_equal(clf.predict(T_csr.tolil()), true_result)

def test_precomputed():
    clf = NearestCentroid(metric="precomputed")
    clf.fit(X, y)
    S = pairwise_distances(T, clf.centroids_)
    assert_array_equal(clf.predict(S), true_result)

def test_iris():
    # Check consistency on dataset iris.
    for metric in ('euclidean', 'cosine'):
        clf = NearestCentroid(metric=metric).fit(iris.data, iris.target)
        score = np.mean(clf.predict(iris.data) == iris.target)
        assert score > 0.9, "Failed with score = " + str(score)

def test_iris_shrinkage():
    # Check consistency on dataset iris, when using shrinkage.
    for metric in ('euclidean', 'cosine'):
        for shrink_threshold in [None, 0.1, 0.5]:
            clf = NearestCentroid(metric=metric,
                                  shrink_threshold=shrink_threshold)
            clf = clf.fit(iris.data, iris.target)
            score = np.mean(clf.predict(iris.data) == iris.target)
            assert score > 0.8, "Failed with score = " + str(score)

def test_pickle():
    import pickle

    # classification
    obj = NearestCentroid()
    obj.fit(iris.data, iris.target)
    score = obj.score(iris.data, iris.target)
    s = pickle.dumps(obj)

    obj2 = pickle.loads(s)
    assert_equal(type(obj2), obj.__class__)
    score2 = obj2.score(iris.data, iris.target)
    assert_array_equal(score, score2,
                       "Failed to generate same score"
                       " after pickling (classification).")

def test_shrinkage_threshold_decoded_y():
    clf = NearestCentroid(shrink_threshold=0.01)
    y_ind = np.asarray(y)
    y_ind[y_ind == -1] = 0
    clf.fit(X, y_ind)
    centroid_encoded = clf.centroids_
    clf.fit(X, y)
    assert_array_equal(centroid_encoded, clf.centroids_)

def test_predict_translated_data():
    # Test that NearestCentroid gives same results on translated data

    rng = np.random.RandomState(0)
    X = rng.rand(50, 50)
    y = rng.randint(0, 3, 50)
    noise = rng.rand(50)
    clf = NearestCentroid(shrink_threshold=0.1)
    clf.fit(X, y)
    y_init = clf.predict(X)
    clf = NearestCentroid(shrink_threshold=0.1)
    X_noise = X + noise
    clf.fit(X_noise, y)
    y_translate = clf.predict(X_noise)
    assert_array_equal(y_init, y_translate)

def test_manhattan_metric():
    # Test the manhattan metric.

    clf = NearestCentroid(metric='manhattan')
    clf.fit(X, y)
    dense_centroid = clf.centroids_
    clf.fit(X_csr, y)
    assert_array_equal(clf.centroids_, dense_centroid)
    assert_array_equal(dense_centroid, [[-1, -1], [1, 1]])


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Sources for numpy arrays and pandas DataFrames."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.contrib.learn.python.learn.dataframe import transform
from tensorflow.contrib.learn.python.learn.dataframe.queues import feeding_functions

class BaseInMemorySource(transform.TensorFlowTransform):
  """Abstract parent class for NumpySource and PandasSource."""

  def __init__(self,
               data,
               num_threads=None,
               enqueue_size=None,
               batch_size=None,
               queue_capacity=None,
               shuffle=False,
               min_after_dequeue=None,
               seed=None,
               data_name="in_memory_data"):
    super(BaseInMemorySource, self).__init__()
    self._data = data
    self._num_threads = 1 if num_threads is None else num_threads
    self._batch_size = (32 if batch_size is None else batch_size)
    self._enqueue_size = max(1, int(self._batch_size / self._num_threads)
                            ) if enqueue_size is None else enqueue_size
    self._queue_capacity = (self._batch_size * 10 if queue_capacity is None else
                            queue_capacity)
    self._shuffle = shuffle
    self._min_after_dequeue = (batch_size if min_after_dequeue is None else
                               min_after_dequeue)
    self._seed = seed
    self._data_name = data_name

  @transform.parameter
  def data(self):
    return self._data

  @transform.parameter
  def num_threads(self):
    return self._num_threads

  @transform.parameter
  def enqueue_size(self):
    return self._enqueue_size

  @transform.parameter
  def batch_size(self):
    return self._batch_size

  @transform.parameter
  def queue_capacity(self):
    return self._queue_capacity

  @transform.parameter
  def shuffle(self):
    return self._shuffle

  @transform.parameter
  def min_after_dequeue(self):
    return self._min_after_dequeue

  @transform.parameter
  def seed(self):
    return self._seed

  @transform.parameter
  def data_name(self):
    return self._data_name

  @property
  def input_valency(self):
    return 0

  def _apply_transform(self, transform_input, **kwargs):
    queue = feeding_functions.enqueue_data(self.data,
                                           self.queue_capacity,
                                           self.shuffle,
                                           self.min_after_dequeue,
                                           num_threads=self.num_threads,
                                           seed=self.seed,
                                           name=self.data_name,
                                           enqueue_size=self.enqueue_size,
                                           num_epochs=kwargs.get("num_epochs"))

    dequeued = queue.dequeue_many(self.batch_size)

    # TODO(jamieas): dequeue and dequeue_many will soon return a list regardless
    # of the number of enqueued tensors. Remove the following once that change
    # is in place.
    if not isinstance(dequeued, (tuple, list)):
      dequeued = (dequeued,)
    # pylint: disable=not-callable
    return self.return_type(*dequeued)

class NumpySource(BaseInMemorySource):
  """A zero-input Transform that produces a single column from a numpy array."""

  @property
  def name(self):
    return "NumpySource"

  @property
  def _output_names(self):
    return ("index", "value")

class OrderedDictNumpySource(BaseInMemorySource):
  """A zero-input Transform that produces Series from a dict of numpy arrays."""

  def __init__(self,
               ordered_dict_of_arrays,
               num_threads=None,
               enqueue_size=None,
               batch_size=None,
               queue_capacity=None,
               shuffle=False,
               min_after_dequeue=None,
               seed=None,
               data_name="pandas_data"):
    if "index" in ordered_dict_of_arrays.keys():
      raise ValueError("Column name `index` is reserved.")
    super(OrderedDictNumpySource, self).__init__(ordered_dict_of_arrays,
                                                 num_threads, enqueue_size,
                                                 batch_size, queue_capacity,
                                                 shuffle, min_after_dequeue,
                                                 seed, data_name)

  @property
  def name(self):
    return "OrderedDictNumpySource"

  @property
  def _output_names(self):
    return tuple(["index"] + list(self._data.keys()))

class PandasSource(BaseInMemorySource):
  """A zero-input Transform that produces Series from a DataFrame."""

  def __init__(self,
               dataframe,
               num_threads=None,
               enqueue_size=None,
               batch_size=None,
               queue_capacity=None,
               shuffle=False,
               min_after_dequeue=None,
               seed=None,
               data_name="pandas_data"):
    if "index" in dataframe.columns:
      raise ValueError("Column name `index` is reserved.")
    super(PandasSource, self).__init__(dataframe, num_threads, enqueue_size,
                                       batch_size, queue_capacity, shuffle,
                                       min_after_dequeue, seed, data_name)

  @property
  def name(self):
    return "PandasSource"

  @property
  def _output_names(self):
    return tuple(["index"] + self._data.columns.tolist())


import pickle

import numpy as np

from sklearn.utils.testing import assert_almost_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_raises_regexp
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import ignore_warnings
from sklearn.utils.testing import assert_not_equal

from sklearn.base import BaseEstimator
from sklearn.metrics import (f1_score, r2_score, roc_auc_score, fbeta_score,
                             log_loss, precision_score, recall_score)
from sklearn.metrics.cluster import adjusted_rand_score
from sklearn.metrics.scorer import (check_scoring, _PredictScorer,
                                    _passthrough_scorer)
from sklearn.metrics import make_scorer, get_scorer, SCORERS
from sklearn.svm import LinearSVC
from sklearn.pipeline import make_pipeline
from sklearn.cluster import KMeans
from sklearn.dummy import DummyRegressor
from sklearn.linear_model import Ridge, LogisticRegression
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.datasets import make_blobs
from sklearn.datasets import make_classification
from sklearn.datasets import make_multilabel_classification
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.model_selection import GridSearchCV
from sklearn.multiclass import OneVsRestClassifier

REGRESSION_SCORERS = ['r2', 'mean_absolute_error', 'mean_squared_error',
                      'median_absolute_error']

CLF_SCORERS = ['accuracy', 'f1', 'f1_weighted', 'f1_macro', 'f1_micro',
               'roc_auc', 'average_precision', 'precision',
               'precision_weighted', 'precision_macro', 'precision_micro',
               'recall', 'recall_weighted', 'recall_macro', 'recall_micro',
               'log_loss',
               'adjusted_rand_score'  # not really, but works
               ]

MULTILABEL_ONLY_SCORERS = ['precision_samples', 'recall_samples', 'f1_samples']

class EstimatorWithoutFit(object):
    """Dummy estimator to test check_scoring"""
    pass

class EstimatorWithFit(BaseEstimator):
    """Dummy estimator to test check_scoring"""
    def fit(self, X, y):
        return self

class EstimatorWithFitAndScore(object):
    """Dummy estimator to test check_scoring"""
    def fit(self, X, y):
        return self

    def score(self, X, y):
        return 1.0

class EstimatorWithFitAndPredict(object):
    """Dummy estimator to test check_scoring"""
    def fit(self, X, y):
        self.y = y
        return self

    def predict(self, X):
        return self.y

class DummyScorer(object):
    """Dummy scorer that always returns 1."""
    def __call__(self, est, X, y):
        return 1

def test_check_scoring():
    # Test all branches of check_scoring
    estimator = EstimatorWithoutFit()
    pattern = (r"estimator should a be an estimator implementing 'fit' method,"
               r" .* was passed")
    assert_raises_regexp(TypeError, pattern, check_scoring, estimator)

    estimator = EstimatorWithFitAndScore()
    estimator.fit([[1]], [1])
    scorer = check_scoring(estimator)
    assert_true(scorer is _passthrough_scorer)
    assert_almost_equal(scorer(estimator, [[1]], [1]), 1.0)

    estimator = EstimatorWithFitAndPredict()
    estimator.fit([[1]], [1])
    pattern = (r"If no scoring is specified, the estimator passed should have"
               r" a 'score' method\. The estimator .* does not\.")
    assert_raises_regexp(TypeError, pattern, check_scoring, estimator)

    scorer = check_scoring(estimator, "accuracy")
    assert_almost_equal(scorer(estimator, [[1]], [1]), 1.0)

    estimator = EstimatorWithFit()
    scorer = check_scoring(estimator, "accuracy")
    assert_true(isinstance(scorer, _PredictScorer))

    estimator = EstimatorWithFit()
    scorer = check_scoring(estimator, allow_none=True)
    assert_true(scorer is None)

def test_check_scoring_gridsearchcv():
    # test that check_scoring works on GridSearchCV and pipeline.
    # slightly redundant non-regression test.

    grid = GridSearchCV(LinearSVC(), param_grid={'C': [.1, 1]})
    scorer = check_scoring(grid, "f1")
    assert_true(isinstance(scorer, _PredictScorer))

    pipe = make_pipeline(LinearSVC())
    scorer = check_scoring(pipe, "f1")
    assert_true(isinstance(scorer, _PredictScorer))

    # check that cross_val_score definitely calls the scorer
    # and doesn't make any assumptions about the estimator apart from having a
    # fit.
    scores = cross_val_score(EstimatorWithFit(), [[1], [2], [3]], [1, 0, 1],
                             scoring=DummyScorer())
    assert_array_equal(scores, 1)

def test_make_scorer():
    # Sanity check on the make_scorer factory function.
    f = lambda *args: 0
    assert_raises(ValueError, make_scorer, f, needs_threshold=True,
                  needs_proba=True)

def test_classification_scores():
    # Test classification scorers.
    X, y = make_blobs(random_state=0, centers=2)
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
    clf = LinearSVC(random_state=0)
    clf.fit(X_train, y_train)

    for prefix, metric in [('f1', f1_score), ('precision', precision_score),
                           ('recall', recall_score)]:

        score1 = get_scorer('%s_weighted' % prefix)(clf, X_test, y_test)
        score2 = metric(y_test, clf.predict(X_test), pos_label=None,
                        average='weighted')
        assert_almost_equal(score1, score2)

        score1 = get_scorer('%s_macro' % prefix)(clf, X_test, y_test)
        score2 = metric(y_test, clf.predict(X_test), pos_label=None,
                        average='macro')
        assert_almost_equal(score1, score2)

        score1 = get_scorer('%s_micro' % prefix)(clf, X_test, y_test)
        score2 = metric(y_test, clf.predict(X_test), pos_label=None,
                        average='micro')
        assert_almost_equal(score1, score2)

        score1 = get_scorer('%s' % prefix)(clf, X_test, y_test)
        score2 = metric(y_test, clf.predict(X_test), pos_label=1)
        assert_almost_equal(score1, score2)

    # test fbeta score that takes an argument
    scorer = make_scorer(fbeta_score, beta=2)
    score1 = scorer(clf, X_test, y_test)
    score2 = fbeta_score(y_test, clf.predict(X_test), beta=2)
    assert_almost_equal(score1, score2)

    # test that custom scorer can be pickled
    unpickled_scorer = pickle.loads(pickle.dumps(scorer))
    score3 = unpickled_scorer(clf, X_test, y_test)
    assert_almost_equal(score1, score3)

    # smoke test the repr:
    repr(fbeta_score)

def test_regression_scorers():
    # Test regression scorers.
    diabetes = load_diabetes()
    X, y = diabetes.data, diabetes.target
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
    clf = Ridge()
    clf.fit(X_train, y_train)
    score1 = get_scorer('r2')(clf, X_test, y_test)
    score2 = r2_score(y_test, clf.predict(X_test))
    assert_almost_equal(score1, score2)

def test_thresholded_scorers():
    # Test scorers that take thresholds.
    X, y = make_blobs(random_state=0, centers=2)
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
    clf = LogisticRegression(random_state=0)
    clf.fit(X_train, y_train)
    score1 = get_scorer('roc_auc')(clf, X_test, y_test)
    score2 = roc_auc_score(y_test, clf.decision_function(X_test))
    score3 = roc_auc_score(y_test, clf.predict_proba(X_test)[:, 1])
    assert_almost_equal(score1, score2)
    assert_almost_equal(score1, score3)

    logscore = get_scorer('log_loss')(clf, X_test, y_test)
    logloss = log_loss(y_test, clf.predict_proba(X_test))
    assert_almost_equal(-logscore, logloss)

    # same for an estimator without decision_function
    clf = DecisionTreeClassifier()
    clf.fit(X_train, y_train)
    score1 = get_scorer('roc_auc')(clf, X_test, y_test)
    score2 = roc_auc_score(y_test, clf.predict_proba(X_test)[:, 1])
    assert_almost_equal(score1, score2)

    # test with a regressor (no decision_function)
    reg = DecisionTreeRegressor()
    reg.fit(X_train, y_train)
    score1 = get_scorer('roc_auc')(reg, X_test, y_test)
    score2 = roc_auc_score(y_test, reg.predict(X_test))
    assert_almost_equal(score1, score2)

    # Test that an exception is raised on more than two classes
    X, y = make_blobs(random_state=0, centers=3)
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
    clf.fit(X_train, y_train)
    assert_raises(ValueError, get_scorer('roc_auc'), clf, X_test, y_test)

def test_thresholded_scorers_multilabel_indicator_data():
    # Test that the scorer work with multilabel-indicator format
    # for multilabel and multi-output multi-class classifier
    X, y = make_multilabel_classification(allow_unlabeled=False,
                                          random_state=0)
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

    # Multi-output multi-class predict_proba
    clf = DecisionTreeClassifier()
    clf.fit(X_train, y_train)
    y_proba = clf.predict_proba(X_test)
    score1 = get_scorer('roc_auc')(clf, X_test, y_test)
    score2 = roc_auc_score(y_test, np.vstack(p[:, -1] for p in y_proba).T)
    assert_almost_equal(score1, score2)

    # Multi-output multi-class decision_function
    # TODO Is there any yet?
    clf = DecisionTreeClassifier()
    clf.fit(X_train, y_train)
    clf._predict_proba = clf.predict_proba
    clf.predict_proba = None
    clf.decision_function = lambda X: [p[:, 1] for p in clf._predict_proba(X)]

    y_proba = clf.decision_function(X_test)
    score1 = get_scorer('roc_auc')(clf, X_test, y_test)
    score2 = roc_auc_score(y_test, np.vstack(p for p in y_proba).T)
    assert_almost_equal(score1, score2)

    # Multilabel predict_proba
    clf = OneVsRestClassifier(DecisionTreeClassifier())
    clf.fit(X_train, y_train)
    score1 = get_scorer('roc_auc')(clf, X_test, y_test)
    score2 = roc_auc_score(y_test, clf.predict_proba(X_test))
    assert_almost_equal(score1, score2)

    # Multilabel decision function
    clf = OneVsRestClassifier(LinearSVC(random_state=0))
    clf.fit(X_train, y_train)
    score1 = get_scorer('roc_auc')(clf, X_test, y_test)
    score2 = roc_auc_score(y_test, clf.decision_function(X_test))
    assert_almost_equal(score1, score2)

def test_unsupervised_scorers():
    # Test clustering scorers against gold standard labeling.
    # We don't have any real unsupervised Scorers yet.
    X, y = make_blobs(random_state=0, centers=2)
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
    km = KMeans(n_clusters=3)
    km.fit(X_train)
    score1 = get_scorer('adjusted_rand_score')(km, X_test, y_test)
    score2 = adjusted_rand_score(y_test, km.predict(X_test))
    assert_almost_equal(score1, score2)

@ignore_warnings
def test_raises_on_score_list():
    # Test that when a list of scores is returned, we raise proper errors.
    X, y = make_blobs(random_state=0)
    f1_scorer_no_average = make_scorer(f1_score, average=None)
    clf = DecisionTreeClassifier()
    assert_raises(ValueError, cross_val_score, clf, X, y,
                  scoring=f1_scorer_no_average)
    grid_search = GridSearchCV(clf, scoring=f1_scorer_no_average,
                               param_grid={'max_depth': [1, 2]})
    assert_raises(ValueError, grid_search.fit, X, y)

@ignore_warnings
def test_scorer_sample_weight():
    # Test that scorers support sample_weight or raise sensible errors

    # Unlike the metrics invariance test, in the scorer case it's harder
    # to ensure that, on the classifier output, weighted and unweighted
    # scores really should be unequal.
    X, y = make_classification(random_state=0)
    _, y_ml = make_multilabel_classification(n_samples=X.shape[0],
                                             random_state=0)
    split = train_test_split(X, y, y_ml, random_state=0)
    X_train, X_test, y_train, y_test, y_ml_train, y_ml_test = split

    sample_weight = np.ones_like(y_test)
    sample_weight[:10] = 0

    # get sensible estimators for each metric
    sensible_regr = DummyRegressor(strategy='median')
    sensible_regr.fit(X_train, y_train)
    sensible_clf = DecisionTreeClassifier(random_state=0)
    sensible_clf.fit(X_train, y_train)
    sensible_ml_clf = DecisionTreeClassifier(random_state=0)
    sensible_ml_clf.fit(X_train, y_ml_train)
    estimator = dict([(name, sensible_regr)
                      for name in REGRESSION_SCORERS] +
                     [(name, sensible_clf)
                      for name in CLF_SCORERS] +
                     [(name, sensible_ml_clf)
                      for name in MULTILABEL_ONLY_SCORERS])

    for name, scorer in SCORERS.items():
        if name in MULTILABEL_ONLY_SCORERS:
            target = y_ml_test
        else:
            target = y_test
        try:
            weighted = scorer(estimator[name], X_test, target,
                              sample_weight=sample_weight)
            ignored = scorer(estimator[name], X_test[10:], target[10:])
            unweighted = scorer(estimator[name], X_test, target)
            assert_not_equal(weighted, unweighted,
                             msg="scorer {0} behaves identically when "
                             "called with sample weights: {1} vs "
                             "{2}".format(name, weighted, unweighted))
            assert_almost_equal(weighted, ignored,
                                err_msg="scorer {0} behaves differently when "
                                "ignoring samples and setting sample_weight to"
                                " 0: {1} vs {2}".format(name, weighted,
                                                        ignored))

        except TypeError as e:
            assert_true("sample_weight" in str(e),
                        "scorer {0} raises unhelpful exception when called "
                        "with sample weights: {1}".format(name, str(e)))


"""
Reads in current year's Arctic sea ice extent from Sea Ice Index 3 (NSIDC)

Website   : ftp://sidads.colorado.edu/DATASETS/NOAA/G02135/north/daily/data/
Author    : Zachary M. Labe
Date      : 5 September 2016
"""

### Import modules
import numpy as np
import urllib.request
import urllib as UL
import datetime
import matplotlib.pyplot as plt

### Directory and time
directoryfigure = './Figures/'
now = datetime.datetime.now()
currentmn = str(now.month)
currentdy = str(now.day)
currentyr = str(now.year)
currenttime = currentmn + '_' + currentdy + '_' + currentyr
currentdoy = now.timetuple().tm_yday

### Load url
url = 'ftp://sidads.colorado.edu/DATASETS/NOAA/G02135/north/daily/data/' \
    'N_seaice_extent_daily_v3.0.csv'

### Read file
raw_data = UL.request.urlopen(url)
dataset = np.genfromtxt(raw_data, skip_header=2,delimiter=',',
                        usecols=[0,1,2,3,4])

print('\nCompleted: Read sea ice data!')

### Set missing data to nan
dataset[np.where(dataset==-9999)] = np.nan

### Variables
year = dataset[:,0]
month = dataset[:,1]
day = dataset[:,2]
ice = dataset[:,3]
missing = dataset[:,4]

### Call present year
yr2018 = np.where(year == 2018)[0]
ice18 = ice[yr2018]

### Ice Conversion
iceval = ice18 * 1e6

### Printing info
print('\n----- NSIDC Arctic Sea Ice -----')
print('Current Date =', now.strftime("%Y-%m-%d %H:%M"), '\n')

print('SIE Date    = %s/%s/%s' % (int(month[-1]),int(day[-1]),int(year[-1])))
print('Current SIE = %s km^2 \n' % (iceval[-1]))

print('1-day change SIE = %s km^2' % (iceval[-1]-iceval[-2]))
print('7-day change SIE = %s km^2 \n' % (iceval[-1]-iceval[-8]))

###########################################################################
###########################################################################
###########################################################################
### Reads in 1981-2010 means
### Load url
url2 = 'ftp://sidads.colorado.edu/DATASETS/NOAA/G02135/north/daily/data/' \
       'N_seaice_extent_climatology_1981-2010_v3.0.csv'

### Read file
raw_data2 = UL.request.urlopen(url2)
dataset2 = np.genfromtxt(raw_data2, skip_header=2,delimiter=',',
                        usecols=[0,1,2,3,4,5,6,7])

### Create variables
doy = dataset2[:,0]
meanice = dataset2[:,1] * 1e6
std = dataset2[:,2]

### Quartiles
quartile10 = dataset2[:,3]
quartile25 = dataset2[:,4]
quartile50 = dataset2[:,5]
quartile75 = dataset2[:,6]
quartile90 = dataset2[:,7]

### Anomalies
currentanom = iceval[-1]-meanice[currentdoy-2]

### Printing info
print('Current anomaly = %s km^2 \n' % currentanom)

### Selected other years for comparisons
yr2007 = np.where(year == 2007)[0]
yr2012 = np.where(year == 2012)[0]
yr2016 = np.where(year == 2016)[0]

sie7 = ice[yr2007]
sie12 = ice[yr2012]
sie16 = ice[yr2016]

###########################################################################
###########################################################################
###########################################################################
### Create plot
plt.rc('text',usetex=True)
plt.rc('font',**{'family':'sans-serif','sans-serif':['Avant Garde']})
plt.rc('savefig',facecolor='black')
plt.rc('axes',edgecolor='white')
plt.rc('xtick',color='white')
plt.rc('ytick',color='white')
plt.rc('axes',labelcolor='white')
plt.rc('axes',facecolor='black')

fig = plt.figure()
ax = plt.subplot(111)

xlabels = [r'Jan',r'Feb',r'Mar',r'Apr',r'May',r'Jun',r'Jul',
          r'Aug',r'Sep',r'Oct',r'Nov',r'Dec',r'Jan']
plt.xticks(np.arange(0,361,30.4),xlabels,rotation=0)
ylabels = map(str,np.arange(2,19,2))
plt.yticks(np.arange(2,19,2),ylabels)
plt.ylim([2,18])
plt.xlim([0,360])

strmonth = xlabels[int(currentmn)-1]
asof = strmonth + ' ' + currentdy + ', ' + currentyr

### Adjust axes in time series plots
def adjust_spines(ax, spines):
    for loc, spine in ax.spines.items():
        if loc in spines:
            spine.set_position(('outward', 5))
        else:
            spine.set_color('none')
    if 'left' in spines:
        ax.yaxis.set_ticks_position('left')
    else:
        ax.yaxis.set_ticks([])

    if 'bottom' in spines:
        ax.xaxis.set_ticks_position('bottom')
    else:
        ax.xaxis.set_ticks([])

ax.tick_params('both',length=5.5,width=2,which='major')
adjust_spines(ax, ['left','bottom'])
ax.spines['top'].set_color('none')
ax.spines['right'].set_color('none')
ax.spines['bottom'].set_linewidth(2)
ax.spines['left'].set_linewidth(2)

upper2std = (meanice/1e6)+(std*2)
lower2std = (meanice/1e6)-(std*2)

ax.grid(zorder=1,color='w',alpha=0.2)

plt.plot(ice18,linewidth=1.8,color='aqua',zorder=9,label=r'Current Year (2018)')

plt.plot(doy,upper2std,color='white',alpha=0.7,zorder=3,linewidth=0.1)
plt.plot(doy,lower2std,color='white',alpha=0.7,zorder=4,linewidth=0.1)
plt.plot(doy,quartile10,color='m',alpha=0.7,zorder=3,linewidth=0.4)
plt.plot(doy,quartile25,color='cornflowerblue',alpha=0.7,zorder=4,linewidth=0.4)
plt.plot(doy,quartile75,color='cornflowerblue',alpha=0.7,zorder=4,linewidth=0.4)
plt.plot(doy,quartile90,color='m',alpha=0.7,zorder=3,linewidth=0.4)

ax.fill_between(doy, lower2std, upper2std, facecolor='white', alpha=0.35,
                label=r'$\pm$2 standard deviations',zorder=2)
plt.plot(doy,quartile50,color='gold',alpha=1,zorder=3,linewidth=2,
         label=r'Median (1981-2010)')

ax.fill_between(doy, quartile90, quartile75, facecolor='m', alpha=0.55,
                label=r'10-90th percentiles',zorder=2)
ax.fill_between(doy, quartile10, quartile25, facecolor='m', alpha=0.55,
                zorder=2)
ax.fill_between(doy, quartile25, quartile50, facecolor='cornflowerblue', alpha=0.6,
                zorder=2)
ax.fill_between(doy, quartile50, quartile75, facecolor='cornflowerblue', alpha=0.6,
                label=r'25-75th percentiles',zorder=2)

plt.scatter(doy[currentdoy-3],ice[-1],s=10,color='aqua',zorder=9)

plt.ylabel(r'\textbf{Extent} [$\times$10$^{6}$ km$^2$]',fontsize=15,
           color='darkgrey')
le = plt.legend(shadow=False,fontsize=6,loc='upper left',
           bbox_to_anchor=(0.473, 1.011),fancybox=True,ncol=2)
for text in le.get_texts():
    text.set_color('w')
plt.title(r'\textbf{ARCTIC SEA ICE}',
                       fontsize=21,color='darkgrey')

plt.text(doy[currentdoy]-5,ice[-1]-1.35,r'\textbf{2018}',
         fontsize=13.5,rotation='horizontal',ha='left',color='aqua')
plt.text(0.5,3.1,r'\textbf{DATA:} National Snow \& Ice Data Center, Boulder CO',
         fontsize=5.5,rotation='horizontal',ha='left',color='darkgrey')
plt.text(0.5,2.6,r'\textbf{SOURCE:} ftp://sidads.colorado.edu/DATASETS/NOAA/G02135/',
         fontsize=5.5,rotation='horizontal',ha='left',color='darkgrey')
plt.text(0.5,2.1,r'\textbf{GRAPHIC:} Zachary Labe (@ZLabe)',
         fontsize=5.5,rotation='horizontal',ha='left',color='darkgrey')
fig.subplots_adjust(top=0.91)

### Save figure
plt.savefig(directoryfigure + 'nsidc_sie_quartiles_currentyear.png',dpi=300)


def input_jk_output_index(j,k,kmax):
    return j*(kmax) + k

def extrap1d(interpolator):

    xs = interpolator.x
    ys = interpolator.y

    def pointwise(x):
        if x < xs[0]:
            return ys[0]+(x-xs[0])*(ys[1]-ys[0])/(xs[1]-xs[0])
        elif x > xs[-1]:
            return ys[-1]+(x-xs[-1])*(ys[-1]-ys[-2])/(xs[-1]-xs[-2])
        else:
            return interpolator(x)

    def ufunclike(xs):
        from scipy import array
        return array(map(pointwise, array(xs)))

    return ufunclike

def solve_uref_both_bc(tstamp, zmum, FAWA_cos, ylat, ephalf2, Delta_PT,
                       zm_PT, Input_B0, Input_B1, use_real_Data=True,
                       plot_all_ref_quan=False):
    """
    Compute equivalent latitude and wave activity on a barotropic sphere.

    Parameters
    ----------
    tstamp : string
        Time stamp of the snapshot of the field.
    znum : ndarray
        Zonal mean wind.
    FAWA_cos : ndarray
        Zonal mean finite-amplitude wave activity.
    ylat : sequence or array_like
        1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat.
    ephalf2 : ndarray
        Epsilon in Nakamura and Solomon (2010).
    Delta_PT : ndarray
        \Delta \Theta in Nakamura and Solomon (2010); upper-boundary conditions.
    zm_PT : ndarray
        Zonal mean potential temperature.
    Input_B0 : sequence or array_like
        Zonal-mean surface wave activity for the lowest layer (k=0). Part of the lower-boundary condition.
    Input_B1 : sequence or array_like
        Zonal-mean surface wave activity for the second lowest layer (k=1). Part of the lower-boundary condition.
    use_real_Data : boolean
        Whether to use input data to compute the reference states. By detault True. If false, randomly generated arrays will be used.
    plot_all_ref_quan : boolean
        Whether to plot the solved reference states using matplotlib library. By default False. For debugging.

    Returns
    -------
    u_MassCorr_regular_noslip : ndarray
        2-d numpy array of mass correction \Delta u in NS10 with no-slip lower boundary conditions; dimension = (kmax,nlat).
    u_Ref_regular_noslip : ndarray
        2-d numpy array of zonal wind reference state u_ref in NS10 with no-slip lower boundary conditions; dimension = (kmax,nlat).
    T_MassCorr_regular_noslip : ndarray
        2-d numpy array of adjustment in reference temperature \Delta T in NS10 with no-slip lower boundary conditions; dimension = (kmax,nlat).
    T_Ref_regular_noslip : ndarray
        2-d numpy array of adjustment in reference temperature T_ref in NS10 with no-slip lower boundary conditions; dimension = (kmax,nlat).
    u_MassCorr_regular_adiab : ndarray
        2-d numpy array of mass correction \Delta u in NS10 with adiabatic lower boundary conditions; dimension = (kmax,nlat).
    u_Ref_regular_adiab : ndarray
        2-d numpy array of zonal wind reference state u_ref in NS10 with adiabatic lower boundary conditions; dimension = (kmax,nlat).
    T_MassCorr_regular_adiab : ndarray
        2-d numpy array of adjustment in reference temperature \Delta T in NS10 with adiabatic lower boundary conditions; dimension = (kmax,nlat).
    T_Ref_regular_adiab : ndarray
        2-d numpy array of adjustment in reference temperature T_ref in NS10 with adiabatic lower boundary conditions; dimension = (kmax,nlat).

    """

    # zm_PT = zonal mean potential temperature

    # Import necessary modules
    from math import pi, exp
    from scipy import interpolate
    from scipy.sparse import csc_matrix
    from scipy.sparse.linalg import spsolve
    from copy import copy
    import numpy as np
    import itertools
    if plot_all_ref_quan:
        import matplotlib.pyplot as plt

    # === Parameters (should be input externally. To be modified) ===
    dz = 1000.          # vertical z spacing (m)
    aa = 6378000.     # planetary radius
    r0 = 287.           # gas constant
    hh = 7000.          # scale height
    cp = 1004.          # specific heat
    rkappa = r0/cp
    om = 7.29e-5          # angular velocity of the earth

    # === These changes with input variables' dimensions ===
    nlat = FAWA_cos.shape[-1]
    jmax1 = nlat//4
    dm = 1./float(jmax1+1)  # gaussian latitude spacing
    gl = np.array([(j+1)*dm for j in range(jmax1)]) # This is sin / mu
    gl_2 = np.array([j*dm for j in range(jmax1+2)]) # This is sin / mu
    cosl = np.sqrt(1.-gl**2)
    #cosl_2 = np.sqrt(1.-gl_2**2)
    alat = np.arcsin(gl)*180./pi
    alat_2 = np.arcsin(gl_2)*180./pi
    dmdz = (dm/dz)

    # **** Get from input these parameters ****
    kmax = FAWA_cos.shape[0]
    #height = np.array([i for i in range(kmax)]) # in [km]

    # **** Initialize Coefficients ****
    c_a = np.zeros((jmax1, kmax))
    c_b = np.zeros((jmax1, kmax))
    c_c = np.zeros((jmax1, kmax))
    c_d = np.zeros((jmax1, kmax))
    c_e = np.zeros((jmax1, kmax))
    c_f = np.zeros((jmax1, kmax))

    # --- Initialize interpolated variables ---
    zmu1 = np.zeros((jmax1, kmax))
    cx1 = np.zeros((jmax1, kmax))
    cor1 = np.zeros((jmax1, kmax))
    ephalf = np.zeros((jmax1, kmax))
    Delta_PT1 = np.zeros((jmax1+2))
    zm_PT1 = np.zeros((jmax1, kmax))
    Input_B0_1 = np.zeros((jmax1+2))
    Input_B1_1 = np.zeros((jmax1+2))

    # --- Define Epsilon as a function of y and z ---

    # **** Interpolate to gaussian latitude ****
    if use_real_Data:
        # print 'use_real_Data'
        for vv1,vvm in zip([zmu1,cx1,zm_PT1] , [zmum,FAWA_cos,zm_PT]):
            f_toGaussian = interpolate.interp1d(ylat[:],vvm[:,:].T,axis=0, kind='linear')    #[jmax x kmax]
            vv1[:,:] = f_toGaussian(alat[:])
            #vv1[:,:] = vvm[:,:]
            #vv1[-1,:] = vvm[:,-1]

        # --- Interpolation of ephalf ---
        f_ep_toGaussian = interpolate.interp1d(ylat[:],ephalf2[:,:].T,axis=0, kind='linear')    #[jmax x kmax]
        ephalf[:,:] = f_ep_toGaussian(alat[:])

		# --- Interpolation of Delta_PT ---
        #f_DT_toGaussian = extrap1d( interpolate.interp1d(ylat[:],Delta_PT[:], kind='linear') )    # This is txt in Noboru's code
        f_DT_toGaussian = interpolate.interp1d(ylat[:],Delta_PT[:],
                                               kind='linear',fill_value='extrapolate')
        Delta_PT1[:] = f_DT_toGaussian(alat_2[:])

		# --- Interpolation of Input_B0_1 ---
        #f_B0_toGaussian = extrap1d( interpolate.interp1d(ylat[:],Input_B0[:], kind='linear') )    # This is txt in Noboru's code
        f_B0_toGaussian = interpolate.interp1d(ylat[:],Input_B0[:],
                                               kind='linear',fill_value='extrapolate')    # This is txt in Noboru's code
        Input_B0_1[:] = f_B0_toGaussian(alat_2[:])

		# --- Interpolation of Input_B1_1 ---
        # f_B1_toGaussian = extrap1d( interpolate.interp1d(ylat[:],Input_B1[:], kind='linear') )     # This is txt in Noboru's code
        f_B1_toGaussian = interpolate.interp1d(ylat[:],Input_B1[:],
                                               kind='linear',fill_value='extrapolate')    # This is txt in Noboru's code
        Input_B1_1[:] = f_B1_toGaussian(alat_2[:])

    else:
        # Use random matrix here just to test!
        zmu1 = np.random.rand(jmax1, kmax)+np.ones((jmax1, kmax))*1.e-8
        cx1 = np.random.rand(jmax1, kmax)+np.ones((jmax1, kmax))*1.e-8
        #cor1 = np.random.rand(jmax1, kmax)+np.ones((jmax1, kmax))*1.e-8

    # --- Added on Aug 1, 2016 ---
    cor1 = 2.*om*gl[:,np.newaxis] * np.ones((jmax1, kmax))
    #cor1[0] = cor1[1]*0.5

    # OLD: qxx0 = -cx1*cosl[:,np.newaxis]/cor1     #qxx0 = np.empty((jmax1, kmax))
    qxx0 = -cx1/cor1 # Input of LWA has cosine.
    c_f[0,:] = qxx0[1,:] - 2*qxx0[0,:]
    c_f[-1,:] = qxx0[-2,:] - 2*qxx0[-1,:]
    c_f[1:-1,:] = qxx0[:-2,:] + qxx0[2:,:] - 2*qxx0[1:-1,:]
    #c_f[:,0] = 0.0

    # --- Aug 9: Lower Adiabatic boundary conditions ---
    Input_dB0 = np.zeros((jmax1))
    Input_dB1 = np.zeros((jmax1))
    uz1 = np.zeros((jmax1))

    # prefac = - r0 * cosl[1:-1]**2 * dz / (cor1[1:-1,-2]**2 * aa**2 * hh * dm**2) * exp(-rkappa*(kmax-2.)/7.)

    # OLD: Input_dB0[:] = Input_B0_1[:-2]*cosl_2[:-2] + Input_B0_1[2:]*cosl_2[2:] - 2*Input_B0_1[1:-1]*cosl_2[1:-1]
    Input_dB0[:] = Input_B0_1[:-2] + Input_B0_1[2:] - 2*Input_B0_1[1:-1]

    # OLD: Input_dB1[:] = Input_B1_1[:-2]*cosl_2[:-2] + Input_B1_1[2:]*cosl_2[2:] - 2*Input_B1_1[1:-1]*cosl_2[1:-1]
    Input_dB1[:] = Input_B1_1[:-2] + Input_B1_1[2:] - 2*Input_B1_1[1:-1]

    # This is supposed to be correct but gave weird results.
    uz1[:] = - r0 * cosl[:]**2 * Input_dB1[:] * 2*dz / (cor1[:,1]**2 * aa**2 * hh * dm**2) * exp(-rkappa*(1.)/7.) \
    - r0 * cosl[:]**2 * Input_dB0[:] * 2*dz / (cor1[:,0]**2 * aa**2 * hh * dm**2) * exp(-rkappa*(0.)/7.)

    #  **** Upper Boundary Condition (Come back later) ****
    uz2 = np.zeros((jmax1))
    dDelta_PT1 = (Delta_PT1[2:]-Delta_PT1[:-2]) # Numerical trick: Replace uz2[1] with an extrapolated value

    # Original correct one:
    # uz2[1:-1] = - r0 * cosl[1:-1]**2 * exp(-rkappa*(kmax-2.)/7.) * dDelta_PT1 / (cor1[1:-1,-2]**2 * aa * hh * dmdz)
    uz2[:] = - r0 * cosl[:]**2 * exp(-rkappa*(kmax-2.)/7.) * dDelta_PT1 / (cor1[:,-2]**2 * aa * hh * dmdz)

    #  **** Initialize the coefficients a,b,c,d,e,f ****
    c_a[:,:] = 1.0
    c_b[:,:] = 1.0
    c_c[:,1:-1] = dmdz**2 *ephalf[:,1:-1]*exp(-dz/(2*hh))  # This one should be correct
    c_d[:,1:-1] = dmdz**2 *ephalf[:,0:-2]*exp(dz/(2*hh)) # Check convention of ephalf
    c_e[:,1:-1] = -(c_a[:,1:-1]+c_b[:,1:-1]+c_c[:,1:-1]+c_d[:,1:-1])

    b = np.zeros((jmax1*kmax))

    row_index=[]
    col_index=[]
    coeff = []

    jrange = range(jmax1)
    krange = range(1,kmax-1)
    for j, k in itertools.product(jrange, krange):
    # for j in range(jmax1):
    #     for k in range(1,kmax-1):
        ind = input_jk_output_index(j,k,kmax)
        b[ind] = c_f[j,k]
        if (j<jmax1-1):
            # A[ind,input_jk_output_index(j+1,k,kmax)] = c_a[j,k]
            row_index.append(ind)
            col_index.append(input_jk_output_index(j+1,k,kmax))
            coeff.append(c_a[j,k])
        if (j>0):
            # A[ind,input_jk_output_index(j-1,k,kmax)] = c_b[j,k]
            row_index.append(ind)
            col_index.append(input_jk_output_index(j-1,k,kmax))
            coeff.append(c_b[j,k])
        # A[ind,input_jk_output_index(j,k+1,kmax)] = c_c[j,k]
        row_index.append(ind)
        col_index.append(input_jk_output_index(j,k+1,kmax))
        coeff.append(c_c[j,k])
        # A[ind,input_jk_output_index(j,k-1,kmax)] = c_d[j,k]
        row_index.append(ind)
        col_index.append(input_jk_output_index(j,k-1,kmax))
        coeff.append(c_d[j,k])
        # A[ind,input_jk_output_index(j,k,kmax)] = c_e[j,k]
        row_index.append(ind)
        col_index.append(input_jk_output_index(j,k,kmax))
        coeff.append(c_e[j,k])

    # ==== Upper boundary condition - thermal wind ====
    # for j in range(1,jmax1-1):
    for j in range(jmax1):
        ind1 = input_jk_output_index(j,kmax-1,kmax)
        b[ind1] = uz2[j] #- r0 * cosl[j]**2  * exp(-rkappa*(kmax-2.)/7.) * (Delta_PT1[j+1]-Delta_PT1[j-1])/ (cor1[j,-2]**2 * aa * hh * dmdz)
        # A[ind1,ind1] = 1.0
        row_index.append(ind1)
        col_index.append(ind1)
        coeff.append(1.0)
        # A[ind1,input_jk_output_index(j,kmax-3,kmax)] = -1.0
        row_index.append(ind1)
        col_index.append(input_jk_output_index(j,kmax-3,kmax))
        coeff.append(-1.0)

    # Try sparse matrix
    # print 'try sparse matrix'
    # A = csc_matrix((coeff_noslip, (row_index, col_index)), shape=(jmax1*kmax,jmax1*kmax))
    # print 'shape of A=',A.shape
    # print 'Does it work?'
#
#     csc_matrix((data, (row_ind, col_ind)), [shape=(M, N)])
#     where data, row_ind and col_ind satisfy the relationship a[row_ind[k], col_ind[k]] = data[k].

        # A[ind1,input_jk_output_index(j,kmax-3,kmax)] = -1.0
        #uz2[1:-1] = - r0 * cosl[1:-1]**2 * exp(-rkappa*(kmax-2.)/7.) * (Delta_PT1[2:]-Delta_PT1[:-2]) / (cor1[1:-1,-2]**2 * aa * hh * dmdz)

    # === Make a copy to deal with adiabatic boundary condition ===
    # A: no-slip
    # A_adiab: adiabatic boundary conditions
    row_index_adiab = copy(row_index)
    col_index_adiab = copy(col_index)
    coeff_adiab = copy(coeff)
    b_adiab = np.copy(b)

    # print 'does it work till here?'

    # A_adiab = np.copy(A)

    # ==== Lower boundary condition - adiabatic (k=0) ====
    for j in range(jmax1):
        ind0 = input_jk_output_index(j,0,kmax)
        b_adiab[ind0] = uz1[j]
        # A_adiab[ind0,ind0] = -1.0 # k=0
        row_index_adiab.append(ind0)
        col_index_adiab.append(ind0)
        coeff_adiab.append(-1.0)
        # A_adiab[ind0,input_jk_output_index(j,2,kmax)] = 1.0 # k=2
        row_index_adiab.append(ind0)
        col_index_adiab.append(input_jk_output_index(j,2,kmax))
        coeff_adiab.append(1.0)

    A_adiab = csc_matrix((coeff_adiab, (row_index_adiab, col_index_adiab)), shape=(jmax1*kmax,jmax1*kmax))

    # ==== Lower boundary condition - no-slip (k=0) ====
    for j in range(jmax1):
        ind = input_jk_output_index(j,0,kmax)
        b[ind] = zmu1[j,0]*cosl[j]/cor1[j,0]
        # A[ind,ind] = 1.0
        row_index.append(ind)
        col_index.append(ind)
        coeff.append(1.0)

    A = csc_matrix((coeff, (row_index, col_index)), shape=(jmax1*kmax,jmax1*kmax))

    # print 'is it ok till here????'

    # === Solving the linear system ===
    u2_adiab =  spsolve(A_adiab, b_adiab)
    u2 = spsolve(A, b)

    # === Mapping back to 2D matrix ===
    u_adiab = np.zeros((jmax1+2,kmax))
    u = np.zeros((jmax1+2,kmax))
    for j in range(jmax1):
        for k in range(kmax):
            u_adiab[j+1,k] = u2_adiab[j*kmax + k]
            u[j+1,k] = u2[j*kmax + k]

    u_MassCorr_adiab = np.zeros_like(u_adiab)
    u_MassCorr_noslip = np.zeros_like(u)
    # u_MassCorr[1:-1,:] = u[1:-1,:] * cor1[1:-1,:] / cosl[1:-1,np.newaxis]
    u_MassCorr_adiab[1:-1,:] = u_adiab[1:-1,:] * cor1 / cosl[:,np.newaxis]
    u_MassCorr_noslip[1:-1,:] = u[1:-1,:] * cor1 / cosl[:,np.newaxis]

    # --- Initialize T_MassCorr to be output ---
    u_Ref_regular_adiab = np.zeros_like(zmum)
    u_Ref_regular_noslip = np.zeros_like(zmum)
    u_MassCorr_regular_adiab = np.zeros_like(zmum)
    u_MassCorr_regular_noslip = np.zeros_like(zmum)
    T_Ref_regular_adiab = np.zeros_like(zmum)
    T_Ref_regular_noslip = np.zeros_like(zmum)
    T_MassCorr_regular_adiab = np.zeros_like(zmum)
    T_MassCorr_regular_noslip = np.zeros_like(zmum)

    for u_MassCorr,u_MassCorr_regular,u_Ref_regular,T_MassCorr_regular,T_Ref_regular,BCstring in \
    zip([u_MassCorr_adiab,u_MassCorr_noslip],\
        [u_MassCorr_regular_adiab,u_MassCorr_regular_noslip],\
        [u_Ref_regular_adiab,u_Ref_regular_noslip],\
        [T_MassCorr_regular_adiab,T_MassCorr_regular_noslip],\
        [T_Ref_regular_adiab,T_Ref_regular_noslip],\
        ['Adiabatic','Noslip']):

        # ---- Back out temperature correction here -----
        T_MassCorr = np.zeros_like(u_MassCorr)
        for k in range(1,kmax-2):
            for j in range(2,jmax1,2): # This is temperature not potential temperature!!! Need to check.
                # print 'alat['+str(j)+']=',alat[j]
                # T_MassCorr[j,k] = T_MassCorr[j-2,k] - (2.*om*gl[j])*aa*hh*dmdz / (r0 * cosl[j]) * (u_MassCorr[j,k+1]-u_MassCorr[j,k-1])
                T_MassCorr[j,k] = T_MassCorr[j-2,k] - (2.*om*gl[j-1])*aa*hh*dmdz / (r0 * cosl[j-1]) * (u_MassCorr[j-1,k+1]-u_MassCorr[j-1,k-1])
            # ---- First do interpolation (gl is regular grid) ----
            # f_Todd = interpolate.interp1d(gl[:-1:2],T_MassCorr[1:-1:2,k])    #[jmax x kmax]
            #f_Todd = interpolate.interp1d(gl_2[::2],T_MassCorr[::2,k])    #[jmax x kmax]
            #f_Todd_ex = extrap1d(f_Todd)

            f_Todd = interpolate.interp1d(gl_2[::2],T_MassCorr[::2,k],
                                          kind='linear',fill_value='extrapolate')
            T_MassCorr[:,k] = f_Todd(gl_2[:])

            # T_MassCorr[:,k] = f_Todd_ex(gl_2[:]) # Get all the points interpolated

            # ---- Then do domain average ----
            T_MC_mean = np.mean(T_MassCorr[:,k])
            T_MassCorr[:,k] -= T_MC_mean

        # --- First, interpolate MassCorr back to regular grid first ---
        f_u_MassCorr = interpolate.interp1d(alat_2,u_MassCorr,axis=0, kind='linear')    #[jmax x kmax]
        u_MassCorr_regular[:,-nlat//2:] = f_u_MassCorr(ylat[-nlat//2:]).T
        f_T_MassCorr = interpolate.interp1d(alat_2,T_MassCorr,axis=0, kind='linear')    #[jmax x kmax]
        T_MassCorr_regular[:,-nlat//2:] = f_T_MassCorr(ylat[-nlat//2:]).T

        u_Ref = zmum[:,-nlat//2:] - u_MassCorr_regular[:,-nlat//2:]
        T_ref = zm_PT[:,-nlat//2:] * np.exp(-np.arange(kmax)/7. * rkappa)[:,np.newaxis] - T_MassCorr_regular[:,-nlat//2:]

        u_Ref_regular[:,-nlat//2:] = u_Ref
        T_Ref_regular[:,-nlat//2:] = T_ref
#
           #plot_all_ref_quan = False
        if plot_all_ref_quan:
            # --- height coordinate ---
            height = np.array([i for i in range(kmax)]) # in [km]
            # --- Colorbar scale ---
            contour_int = np.arange(-120,145,5)
            dT_contour_int = np.arange(-120,81,5)
            T_contour_int = np.arange(160,321,5)
            # --- Start plotting figure ---
            fig = plt.subplots(figsize=(12,12))
            plt.subplot(221)
            plt.contourf(ylat[-nlat//2:],height[:-2],u_MassCorr_regular[:-2,-nlat//2:],contour_int)
            plt.colorbar()
            c1=plt.contour(ylat[-nlat//2:],height[:-2],u_MassCorr_regular[:-2,-nlat//2:],contour_int[::2],colors='k')
            plt.clabel(c1,c1.levels,inline=True, fmt='%d', fontsize=10)
            plt.title('$\Delta$ u '+tstamp)
            plt.ylabel('height (km)')
            plt.subplot(222)
            plt.contourf(ylat[-nlat//2:],height[:-2],u_Ref[:-2,:],contour_int)
            plt.colorbar()
            c2=plt.contour(ylat[-nlat//2:],height[:-2],u_Ref[:-2,:],contour_int[::2],colors='k')
            plt.clabel(c2,c2.levels,inline=True, fmt='%d', fontsize=10)
            plt.title('$u_{REF}$ ('+BCstring+' BC)')
            plt.subplot(223)
            plt.contourf(ylat[-nlat//2:],height[:-2],T_MassCorr_regular[:-2,-nlat//2:],dT_contour_int)
            plt.colorbar()
            c3=plt.contour(ylat[-nlat//2:],height[:-2],T_MassCorr_regular[:-2,-nlat//2:],dT_contour_int,colors='k')
            plt.clabel(c3,c3.levels,inline=True, fmt='%d', fontsize=10)
            plt.title('$\Delta$ T')
            plt.ylabel('height (km)')
            plt.subplot(224)
            plt.contourf(ylat[-nlat//2:],height[:-2],T_ref[:-2,:],T_contour_int)
            plt.colorbar()
            c4=plt.contour(ylat[-nlat//2:],height[:-2],T_ref[:-2,:],T_contour_int[::2],colors='k')
            plt.clabel(c4,c4.levels,inline=True, fmt='%d', fontsize=10)
            plt.title('$T_{REF}$')
            plt.ylabel('height (km)')
            plt.tight_layout()
            plt.show()
            #plt.savefig('/home/csyhuang/Dropbox/Research-code/Sep12_test3_'+BCstring+'_'+tstamp+'.png')
            plt.close()

    # This is for only outputing Delta_u and Uref for no-slip and adiabatic boundary conditions.
    return u_MassCorr_regular_noslip,u_Ref_regular_noslip,T_MassCorr_regular_noslip,T_Ref_regular_noslip, u_MassCorr_regular_adiab,u_Ref_regular_adiab,T_MassCorr_regular_adiab,T_Ref_regular_adiab

# --- As a test whether the function Solve_Uref is working ---
if __name__ == "__main__":

    import matplotlib.pyplot as plt
    import numpy as np

    nlat = 121
    kmax = 49
    jmax1 = nlat

    # The codes below is just for testing purpose
    tstamp = 'random'
    ylat = np.linspace(-90,90,121,endpoint=True)
    t1 = np.random.rand(nlat,kmax)+np.ones((nlat,kmax))*0.001
    t2 = np.random.rand(nlat,kmax)+np.ones((nlat,kmax))*0.001
    t3 = np.random.rand(nlat,kmax)+np.ones((nlat,kmax))*0.001
    Delta_PT = np.random.rand(nlat)+np.ones((nlat))*0.001
    zm_PT = np.random.rand(nlat,kmax)+np.ones((nlat,kmax))*0.001
    Input_B0 = np.random.rand(nlat)+np.ones((nlat))*0.001
    Input_B1 = np.random.rand(nlat)+np.ones((nlat))*0.001
    eh = np.random.rand(jmax1, kmax)+np.ones((jmax1, kmax))*0.001
    Delta_PT = np.sort(np.random.rand(jmax1))

    xxx = solve_uref_both_bc(tstamp,t1,t2,ylat,t3,Delta_PT,zm_PT,Input_B0,Input_B1,use_real_Data=True)
    print(xxx)


import pandas as pd

from systems.accounts.pandl_calculators.pandl_calculation import pandlCalculation, apply_weighting

curve_types = ['gross', 'net', 'costs']
GROSS_CURVE = 'gross'
NET_CURVE = 'net'
COSTS_CURVE = 'costs'

class pandlCalculationWithGenericCosts(pandlCalculation):

    def weight(self, weight: pd.Series):

        weighted_capital = apply_weighting(weight, self.capital)
        weighted_positions = apply_weighting(weight, self.positions)

        return pandlCalculationWithGenericCosts(self.price,
                 positions = weighted_positions,
                 fx = self.fx,
                 capital = weighted_capital,
                 value_per_point = self.value_per_point,
                roundpositions = self.roundpositions,
                delayfill = self.delayfill)

    def as_pd_series(self, percent = False, curve_type=NET_CURVE):
        if curve_type==NET_CURVE:
            if percent:
                return  self.net_percentage_pandl()
            else:
                return self.net_pandl_in_base_currency()

        elif curve_type==GROSS_CURVE:
            if percent:
                return self.percentage_pandl()
            else:
                return self.pandl_in_base_currency()
        elif curve_type==COSTS_CURVE:
            if percent:
                return self.costs_percentage_pandl()
            else:
                return self.costs_pandl_in_base_currency()

        else:
            raise Exception("Curve type %s not recognised! Must be one of %s" % (curve_type, curve_types))

    def net_percentage_pandl(self) -> pd.Series:
        gross = self.percentage_pandl()
        costs = self.costs_percentage_pandl()
        net = _add_gross_and_costs(gross, costs)

        return net

    def net_pandl_in_base_currency(self) -> pd.Series:
        gross = self.pandl_in_base_currency()
        costs = self.costs_pandl_in_base_currency()
        net = _add_gross_and_costs(gross, costs)

        return net

    def net_pandl_in_instrument_currency(self) -> pd.Series:
        gross = self.pandl_in_instrument_currency()
        costs = self.costs_pandl_in_instrument_currency()
        net = _add_gross_and_costs(gross, costs)

        return net

    def net_pandl_in_points(self) -> pd.Series:
        gross = self.pandl_in_points()
        costs = self.costs_pandl_in_points()
        net = _add_gross_and_costs(gross, costs)

        return net

    def costs_percentage_pandl(self) -> pd.Series:
        costs_in_base = self.costs_pandl_in_base_currency()
        costs = self._percentage_pandl_given_pandl(costs_in_base)

        return costs

    def costs_pandl_in_base_currency(self) -> pd.Series:
        costs_in_instr_ccy = self.costs_pandl_in_instrument_currency()
        costs_in_base = self._base_pandl_given_currency_pandl(costs_in_instr_ccy)

        return costs_in_base

    def costs_pandl_in_instrument_currency(self) -> pd.Series:
        costs_in_points = self.costs_pandl_in_points()
        costs_in_instr_ccy = self._pandl_in_instrument_ccy_given_points_pandl(costs_in_points)

        return costs_in_instr_ccy

    def costs_pandl_in_points(self) -> pd.Series:
        raise NotImplementedError

def _add_gross_and_costs(gross: pd.Series,
                        costs: pd.Series):
    cumsum_costs = costs.cumsum()
    cumsum_costs_aligned = cumsum_costs.reindex(gross.index, method="ffill")
    costs_aligned = cumsum_costs_aligned.diff()

    net = gross + costs_aligned

    return net


# -*- coding: utf-8 -*-

# TAXIPP -- A French microsimulation model
# By: IPP <taxipp@ipp.eu>
#
# Copyright (C) 2012, 2013, 2014, 2015 IPP
# https://github.com/taxipp
#
# This file is part of TAXIPP.
#
# TAXIPP is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# TAXIPP is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import logging
import os
import pandas
import pkg_resources
from ipp_macro_series_parser.config import Config

config_parser = Config(
    config_files_directory = os.path.join(pkg_resources.get_distribution('ipp-macro-series-parser').location)
    )
xls_directory = os.path.join(config_parser.get('data', 'demographie_directory'), 'xls')

log = logging.getLogger(__name__)

def create_demographie_data_frame():
    data_frame = pandas.DataFrame()
    for year in range(1999, 2015 + 1):
        file_path = os.path.join(xls_directory, u'pyramide-des-ages-{}.xls'.format(year))
        skiprows = 5 - (year == 1999)
        parse_cols = "A:E"
        slice_start = 0
        slice_end = 101
        sheetname = 'France'

        if year <= 2010:
            sheetnames = ['France', u'France métropolitaine']
        elif year == 2011:
            sheetnames = ['{} France'.format(year), u"{} métropole".format(year)]
        else:
            sheetnames = ['Pyramide {} France'.format(year), u'Pyramide {} métropole'.format(year)]

        for sheetname in sheetnames:
            try:
                df = pandas.read_excel(
                    file_path,
                    # na_values = '-',
                    sheetname = sheetname,
                    skiprows = skiprows,
                    parse_cols = parse_cols).iloc[slice_start:slice_end]
                df['year'] = year
                if sheetname in ['France', u'France métropolitaine']:
                    df['champ'] = sheetname
                else:
                    df['champ'] = u'France métropolitaine' if u'métropole' in sheetname else 'France'
                # All column name on one line
                remove_cr = dict(
                    (column, column.replace(u"\n", " ").replace("  ", " ")) for column in df.columns)
                df.rename(columns = remove_cr, inplace = True)
                # Femmes _> Nombre de femmes etc
                df.rename(columns = dict(
                    Femmes = "Nombre de femmes",
                    Hommes = "Nombre d'hommes"), inplace = True)
                data_frame = pandas.concat((data_frame, df))
                del df
            except Exception, e:
                print year
                print sheetname
                raise(e)
    return pandas.melt(data_frame, id_vars = ['year', 'champ', u'Âge révolu', u'Année de naissance'])


"""
.. _tut-epochs-dataframe:

Exporting Epochs to Pandas DataFrames
=====================================

This tutorial shows how to export the data in :class:`~mne.Epochs` objects to a
:class:`Pandas DataFrame <pandas.DataFrame>`, and applies a typical Pandas
:doc:`split-apply-combine <pandas:user_guide/groupby>` workflow to examine the
latencies of the response maxima across epochs and conditions.

We'll use the :ref:`sample-dataset` dataset, but load a version of the raw file
that has already been filtered and downsampled, and has an average reference
applied to its EEG channels. As usual we'll start by importing the modules we
need and loading the data:
"""
import os
import seaborn as sns
import mne

sample_data_folder = mne.datasets.sample.data_path()
sample_data_raw_file = os.path.join(sample_data_folder, 'MEG', 'sample',
                                    'sample_audvis_filt-0-40_raw.fif')
raw = mne.io.read_raw_fif(sample_data_raw_file, verbose=False)

###############################################################################
# Next we'll load a list of events from file, map them to condition names with
# an event dictionary, set some signal rejection thresholds (cf.
# :ref:`tut-reject-epochs-section`), and segment the continuous data into
# epochs:

sample_data_events_file = os.path.join(sample_data_folder, 'MEG', 'sample',
                                       'sample_audvis_filt-0-40_raw-eve.fif')
events = mne.read_events(sample_data_events_file)

event_dict = {'auditory/left': 1, 'auditory/right': 2, 'visual/left': 3,
              'visual/right': 4}

reject_criteria = dict(mag=3000e-15,     # 3000 fT
                       grad=3000e-13,    # 3000 fT/cm
                       eeg=100e-6,       # 100 µV
                       eog=200e-6)       # 200 µV

tmin, tmax = (-0.2, 0.5)  # epoch from 200 ms before event to 500 ms after it
baseline = (None, 0)      # baseline period from start of epoch to time=0

epochs = mne.Epochs(raw, events, event_dict, tmin, tmax, proj=True,
                    baseline=baseline, reject=reject_criteria, preload=True)
del raw

###############################################################################
# Converting an ``Epochs`` object to a ``DataFrame``
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# Once we have our :class:`~mne.Epochs` object, converting it to a
# :class:`~pandas.DataFrame` is simple: just call :meth:`epochs.to_data_frame()
# <mne.Epochs.to_data_frame>`. Each channel's data will be a column of the new
# :class:`~pandas.DataFrame`, alongside three additional columns of event name,
# epoch number, and sample time. Here we'll just show the first few rows and
# columns:

df = epochs.to_data_frame()
df.iloc[:5, :10]

###############################################################################
# Scaling time and channel values
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# By default, time values are converted from seconds to milliseconds and
# then rounded to the nearest integer; if you don't want this, you can pass
# ``time_format=None`` to keep time as a :class:`float` value in seconds, or
# convert it to a :class:`~pandas.Timedelta` value via
# ``time_format='timedelta'``.
#
# Note also that, by default, channel measurement values are scaled so that EEG
# data are converted to µV, magnetometer data are converted to fT, and
# gradiometer data are converted to fT/cm. These scalings can be customized
# through the ``scalings`` parameter, or suppressed by passing
# ``scalings=dict(eeg=1, mag=1, grad=1)``.

df = epochs.to_data_frame(time_format=None,
                          scalings=dict(eeg=1, mag=1, grad=1))
df.iloc[:5, :10]

###############################################################################
# Notice that the time values are no longer integers, and the channel values
# have changed by several orders of magnitude compared to the earlier
# DataFrame.
#
#
# Setting the ``index``
# ~~~~~~~~~~~~~~~~~~~~~
#
# It is also possible to move one or more of the indicator columns (event name,
# epoch number, and sample time) into the :ref:`index <pandas:indexing>`, by
# passing a string or list of strings as the ``index`` parameter. We'll also
# demonstrate here the effect of ``time_format='timedelta'``, yielding
# :class:`~pandas.Timedelta` values in the "time" column.

df = epochs.to_data_frame(index=['condition', 'epoch'],
                          time_format='timedelta')
df.iloc[:5, :10]

###############################################################################
# Wide- versus long-format DataFrames
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Another parameter, ``long_format``, determines whether each channel's data is
# in a separate column of the :class:`~pandas.DataFrame`
# (``long_format=False``), or whether the measured values are pivoted into a
# single ``'value'`` column with an extra indicator column for the channel name
# (``long_format=True``). Passing ``long_format=True`` will also create an
# extra column ``ch_type`` indicating the channel type.

long_df = epochs.to_data_frame(time_format=None, index='condition',
                               long_format=True)
long_df.head()

###############################################################################
# Generating the :class:`~pandas.DataFrame` in long format can be helpful when
# using other Python modules for subsequent analysis or plotting. For example,
# here we'll take data from the "auditory/left" condition, pick a couple MEG
# channels, and use :func:`seaborn.lineplot` to automatically plot the mean and
# confidence band for each channel, with confidence computed across the epochs
# in the chosen condition:

channels = ['MEG 1332', 'MEG 1342']
data = long_df.loc['auditory/left'].query('channel in @channels')
# convert channel column (CategoryDtype → string; for a nicer-looking legend)
data['channel'] = data['channel'].astype(str)
sns.lineplot(x='time', y='value', hue='channel', data=data)

###############################################################################
# We can also now use all the power of Pandas for grouping and transforming our
# data. Here, we find the latency of peak activation of 2 gradiometers (one
# near auditory cortex and one near visual cortex), and plot the distribution
# of the timing of the peak in each channel as a :func:`~seaborn.violinplot`:

# sphinx_gallery_thumbnail_number = 2
df = epochs.to_data_frame(time_format=None)
peak_latency = (df.filter(regex=r'condition|epoch|MEG 1332|MEG 2123')
                .groupby(['condition', 'epoch'])
                .aggregate(lambda x: df['time'].iloc[x.idxmax()])
                .reset_index()
                .melt(id_vars=['condition', 'epoch'],
                      var_name='channel',
                      value_name='latency of peak')
                )

ax = sns.violinplot(x='channel', y='latency of peak', hue='condition',
                    data=peak_latency, palette='deep', saturation=1)