"""
==========================
Model Complexity Influence
==========================

Demonstrate how model complexity influences both prediction accuracy and
computational performance.

The dataset is the Boston Housing dataset (resp. 20 Newsgroups) for
regression (resp. classification).

For each class of models we make the model complexity vary through the choice
of relevant model parameters and measure the influence on both computational
performance (latency) and predictive power (MSE or Hamming Loss).
"""

print(__doc__)

# Author: Eustache Diemert <eustache@diemert.fr>
# License: BSD 3 clause

import time
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.parasite_axes import host_subplot
from mpl_toolkits.axisartist.axislines import Axes
from scipy.sparse.csr import csr_matrix

from sklearn import datasets
from sklearn.utils import shuffle
from sklearn.metrics import mean_squared_error
from sklearn.svm.classes import NuSVR
from sklearn.ensemble.gradient_boosting import GradientBoostingRegressor
from sklearn.linear_model.stochastic_gradient import SGDClassifier
from sklearn.metrics import hamming_loss

###############################################################################
# Routines

# initialize random generator
np.random.seed(0)

def generate_data(case, sparse=False):
    """Generate regression/classification data."""
    bunch = None
    if case == 'regression':
        bunch = datasets.load_boston()
    elif case == 'classification':
        bunch = datasets.fetch_20newsgroups_vectorized(subset='all')
    X, y = shuffle(bunch.data, bunch.target)
    offset = int(X.shape[0] * 0.8)
    X_train, y_train = X[:offset], y[:offset]
    X_test, y_test = X[offset:], y[offset:]
    if sparse:
        X_train = csr_matrix(X_train)
        X_test = csr_matrix(X_test)
    else:
        X_train = np.array(X_train)
        X_test = np.array(X_test)
    y_test = np.array(y_test)
    y_train = np.array(y_train)
    data = {'X_train': X_train, 'X_test': X_test, 'y_train': y_train,
            'y_test': y_test}
    return data

def benchmark_influence(conf):
    """
    Benchmark influence of :changing_param: on both MSE and latency.
    """
    prediction_times = []
    prediction_powers = []
    complexities = []
    for param_value in conf['changing_param_values']:
        conf['tuned_params'][conf['changing_param']] = param_value
        estimator = conf['estimator'](**conf['tuned_params'])
        print("Benchmarking %s" % estimator)
        estimator.fit(conf['data']['X_train'], conf['data']['y_train'])
        conf['postfit_hook'](estimator)
        complexity = conf['complexity_computer'](estimator)
        complexities.append(complexity)
        start_time = time.time()
        for _ in range(conf['n_samples']):
            y_pred = estimator.predict(conf['data']['X_test'])
        elapsed_time = (time.time() - start_time) / float(conf['n_samples'])
        prediction_times.append(elapsed_time)
        pred_score = conf['prediction_performance_computer'](
            conf['data']['y_test'], y_pred)
        prediction_powers.append(pred_score)
        print("Complexity: %d | %s: %.4f | Pred. Time: %fs\n" % (
            complexity, conf['prediction_performance_label'], pred_score,
            elapsed_time))
    return prediction_powers, prediction_times, complexities

def plot_influence(conf, mse_values, prediction_times, complexities):
    """
    Plot influence of model complexity on both accuracy and latency.
    """
    plt.figure(figsize=(12, 6))
    host = host_subplot(111, axes_class=Axes)
    plt.subplots_adjust(right=0.75)
    par1 = host.twinx()
    host.set_xlabel('Model Complexity (%s)' % conf['complexity_label'])
    y1_label = conf['prediction_performance_label']
    y2_label = "Time (s)"
    host.set_ylabel(y1_label)
    par1.set_ylabel(y2_label)
    p1, = host.plot(complexities, mse_values, 'b-', label="prediction error")
    p2, = par1.plot(complexities, prediction_times, 'r-',
                    label="latency")
    host.legend(loc='upper right')
    host.axis["left"].label.set_color(p1.get_color())
    par1.axis["right"].label.set_color(p2.get_color())
    plt.title('Influence of Model Complexity - %s' % conf['estimator'].__name__)
    plt.show()

def _count_nonzero_coefficients(estimator):
    a = estimator.coef_.toarray()
    return np.count_nonzero(a)

###############################################################################
# main code
regression_data = generate_data('regression')
classification_data = generate_data('classification', sparse=True)
configurations = [
    {'estimator': SGDClassifier,
     'tuned_params': {'penalty': 'elasticnet', 'alpha': 0.001, 'loss':
                      'modified_huber', 'fit_intercept': True},
     'changing_param': 'l1_ratio',
     'changing_param_values': [0.25, 0.5, 0.75, 0.9],
     'complexity_label': 'non_zero coefficients',
     'complexity_computer': _count_nonzero_coefficients,
     'prediction_performance_computer': hamming_loss,
     'prediction_performance_label': 'Hamming Loss (Misclassification Ratio)',
     'postfit_hook': lambda x: x.sparsify(),
     'data': classification_data,
     'n_samples': 30},
    {'estimator': NuSVR,
     'tuned_params': {'C': 1e3, 'gamma': 2 ** -15},
     'changing_param': 'nu',
     'changing_param_values': [0.1, 0.25, 0.5, 0.75, 0.9],
     'complexity_label': 'n_support_vectors',
     'complexity_computer': lambda x: len(x.support_vectors_),
     'data': regression_data,
     'postfit_hook': lambda x: x,
     'prediction_performance_computer': mean_squared_error,
     'prediction_performance_label': 'MSE',
     'n_samples': 30},
    {'estimator': GradientBoostingRegressor,
     'tuned_params': {'loss': 'ls'},
     'changing_param': 'n_estimators',
     'changing_param_values': [10, 50, 100, 200, 500],
     'complexity_label': 'n_trees',
     'complexity_computer': lambda x: x.n_estimators,
     'data': regression_data,
     'postfit_hook': lambda x: x,
     'prediction_performance_computer': mean_squared_error,
     'prediction_performance_label': 'MSE',
     'n_samples': 30},
]
for conf in configurations:
    prediction_performances, prediction_times, complexities = \
        benchmark_influence(conf)
    plot_influence(conf, prediction_performances, prediction_times,
                   complexities)


"""
===================================================================
Multi-output Decision Tree Regression
===================================================================

An example to illustrate multi-output regression with decision tree.

The :ref:`decision trees <tree>`
is used to predict simultaneously the noisy x and y observations of a circle
given a single underlying feature. As a result, it learns local linear
regressions approximating the circle.

We can see that if the maximum depth of the tree (controlled by the
`max_depth` parameter) is set too high, the decision trees learn too fine
details of the training data and learn from the noise, i.e. they overfit.
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeRegressor

# Create a random dataset
rng = np.random.RandomState(1)
X = np.sort(200 * rng.rand(100, 1) - 100, axis=0)
y = np.array([np.pi * np.sin(X).ravel(), np.pi * np.cos(X).ravel()]).T
y[::5, :] += (0.5 - rng.rand(20, 2))

# Fit regression model
clf_1 = DecisionTreeRegressor(max_depth=2)
clf_2 = DecisionTreeRegressor(max_depth=5)
clf_3 = DecisionTreeRegressor(max_depth=8)
clf_1.fit(X, y)
clf_2.fit(X, y)
clf_3.fit(X, y)

# Predict
X_test = np.arange(-100.0, 100.0, 0.01)[:, np.newaxis]
y_1 = clf_1.predict(X_test)
y_2 = clf_2.predict(X_test)
y_3 = clf_3.predict(X_test)

# Plot the results
plt.figure()
plt.scatter(y[:, 0], y[:, 1], c="k", label="data")
plt.scatter(y_1[:, 0], y_1[:, 1], c="g", label="max_depth=2")
plt.scatter(y_2[:, 0], y_2[:, 1], c="r", label="max_depth=5")
plt.scatter(y_3[:, 0], y_3[:, 1], c="b", label="max_depth=8")
plt.xlim([-6, 6])
plt.ylim([-6, 6])
plt.xlabel("data")
plt.ylabel("target")
plt.title("Multi-output Decision Tree Regression")
plt.legend()
plt.show()


#  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.
"""Example of using Estimator with multiple GPUs to distribute one model.

This example only runs if you have multiple GPUs to assign to.
"""

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

from sklearn import cross_validation
from sklearn import datasets
from sklearn import metrics
import tensorflow as tf

layers = tf.contrib.layers
learn = tf.contrib.learn

def my_model(features, target):
  """DNN with three hidden layers, and dropout of 0.1 probability.

  Note: If you want to run this example with multiple GPUs, Cuda Toolkit 7.0 and
  CUDNN 6.5 V2 from NVIDIA need to be installed beforehand.

  Args:
    features: `Tensor` of input features.
    target: `Tensor` of targets.

  Returns:
    Tuple of predictions, loss and training op.
  """
  # Convert the target to a one-hot tensor of shape (length of features, 3) and
  # with a on-value of 1 for each one-hot vector of length 3.
  target = tf.one_hot(target, 3, 1, 0)

  # Create three fully connected layers respectively of size 10, 20, and 10 with
  # each layer having a dropout probability of 0.1.
  normalizer_fn = layers.dropout
  normalizer_params = {'keep_prob': 0.5}
  with tf.device('/gpu:1'):
    features = layers.stack(
        features,
        layers.fully_connected, [10, 20, 10],
        normalizer_fn=normalizer_fn,
        normalizer_params=normalizer_params)

  with tf.device('/gpu:2'):
    # Compute logits (1 per class) and compute loss.
    logits = layers.fully_connected(features, 3, activation_fn=None)
    loss = tf.losses.softmax_cross_entropy(target, logits)

    # Create a tensor for training op.
    train_op = tf.contrib.layers.optimize_loss(
        loss,
        tf.contrib.framework.get_global_step(),
        optimizer='Adagrad',
        learning_rate=0.1)

  return ({
      'class': tf.argmax(logits, 1),
      'prob': tf.nn.softmax(logits)
  }, loss, train_op)

def main(unused_argv):
  iris = datasets.load_iris()
  x_train, x_test, y_train, y_test = cross_validation.train_test_split(
      iris.data, iris.target, test_size=0.2, random_state=42)

  classifier = learn.Estimator(model_fn=my_model)
  classifier.fit(x_train, y_train, steps=1000)

  y_predicted = [
      p['class'] for p in classifier.predict(
          x_test, as_iterable=True)
  ]
  score = metrics.accuracy_score(y_test, y_predicted)
  print('Accuracy: {0:f}'.format(score))

if __name__ == '__main__':
  tf.app.run()


'''Calculate bandstructure diagrams in jasp'''

from jasp import *
import os
import matplotlib.pyplot as plt
from ase.dft import DOS

def get_bandstructure(self,
                      kpts_path=None,
                      kpts_nintersections=10):
    """Calculate band structure along :param kpts_path:

    :param list kpts_path: list of tuples of (label, k-point) to
    calculate path on.

    :param int kpts_nintersections: is the number of points between
    points in band structures. More makes the bands smoother. See
    :func:`jasp_kpts.write_kpoints`.

    >>> from jasp import *
    >>> from jasp.jasp_bandstructure import *

    >>> with jasp('bulk/tio2/step3') as calc:
    ...     n, bands, p = calc.get_bandstructure(kpts_path=[('$\Gamma$',[0.0, 0.0, 0.0]),
                                                            ('X',[0.5, 0.5, 0.0]),
                                                            ('X',[0.5, 0.5, 0.0]),
                                                            ('M',[0.0, 0.5, 0.5]),
                                                            ('M',[0.0, 0.5, 0.5]),
                                                            ('$\Gamma$',[0.0, 0.0, 0.0])])

    >>> p.savefig('images/tio2-bandstructure-dos.png')

    returns (npoints, band_energies, fighandle)

    """

    kpts = [k[1] for k in kpts_path]
    labels = [k[0] for k in kpts_path]

    dos = DOS(self, width=0.2)
    d = dos.get_dos()
    e = dos.get_energies()

    ef = self.get_fermi_level()

    # run in non-selfconsistent directory
    cwd = os.getcwd()
    base, end = os.path.split(cwd)
    wd = cwd + '/bandstructure'
    self.clone(wd)

    with jasp(wd,
              kpts=kpts,
              kpts_nintersections=kpts_nintersections,
              reciprocal=True,
              nsw=0,  # no ionic updates required
              isif=None,
              ibrion=None,
              debug=logging.DEBUG,
              icharg=11) as calc:

        calc.calculate()

        fig = plt.figure()
        with open('EIGENVAL') as f:
            line1 = f.readline()
            line2 = f.readline()
            line3 = f.readline()
            line4 = f.readline()
            comment = f.readline()
            unknown, npoints, nbands = [int(x) for x in f.readline().split()]

            blankline = f.readline()

            band_energies = [[] for i in range(nbands)]

            for i in range(npoints):
                x, y, z, weight = [float(x) for x in f.readline().split()]

                for j in range(nbands):
                    fields = f.readline().split()
                    id, energy = int(fields[0]), float(fields[1])
                    band_energies[id-1].append(energy)
                blankline = f.readline()
            f.close()
            ax1 = plt.subplot(121)
            for i in range(nbands):
                plt.plot(range(npoints), np.array(band_energies[i]) - ef)

            ax = plt.gca()
            ax.set_xticks([])  # no tick marks
            plt.xlabel('k-vector')
            plt.ylabel('Energy (eV)')

            nticks = len(labels)/2 + 1
            ax.set_xticks(np.linspace(0, npoints, nticks))
            L = []
            L.append(labels[0])
            for i in range(2, len(labels)):
                if i % 2 == 0:
                    L.append(labels[i])
                else:
                    pass
            L.append(labels[-1])
            ax.set_xticklabels(L)
            plt.axhline(0, c='r')

    plt.subplot(122, sharey=ax1)
    plt.plot(d, e)
    plt.axhline(0, c='r')
    plt.ylabel('energy (eV)')
    plt.xlabel('DOS')

    plt.subplots_adjust(wspace=0.26)
    plt.show()
    return (npoints, band_energies, fig)

Vasp.get_bandstructure = get_bandstructure


# 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.
# ==============================================================================

"""TensorFlowDataFrame implements convenience functions using TensorFlow."""

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

import collections
import csv

import numpy as np

from tensorflow.contrib.learn.python.learn.dataframe import dataframe as df
from tensorflow.contrib.learn.python.learn.dataframe.transforms import batch
from tensorflow.contrib.learn.python.learn.dataframe.transforms import csv_parser
from tensorflow.contrib.learn.python.learn.dataframe.transforms import example_parser
from tensorflow.contrib.learn.python.learn.dataframe.transforms import in_memory_source
from tensorflow.contrib.learn.python.learn.dataframe.transforms import reader_source
from tensorflow.contrib.learn.python.learn.dataframe.transforms import sparsify
from tensorflow.contrib.learn.python.learn.dataframe.transforms import split_mask
from tensorflow.python.client import session as sess
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.ops import io_ops
from tensorflow.python.ops import parsing_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.training import coordinator
from tensorflow.python.training import queue_runner as qr

def _expand_file_names(filepatterns):
  """Takes a list of file patterns and returns a list of resolved file names."""
  if not isinstance(filepatterns, (list, tuple, set)):
    filepatterns = [filepatterns]
  filenames = set()
  for filepattern in filepatterns:
    names = set(gfile.Glob(filepattern))
    filenames |= names
  return list(filenames)

def _dtype_to_nan(dtype):
  if dtype is dtypes.string:
    return b""
  elif dtype.is_integer:
    return np.nan
  elif dtype.is_floating:
    return np.nan
  elif dtype is dtypes.bool:
    return np.nan
  else:
    raise ValueError("Can't parse type without NaN into sparse tensor: %s" %
                     dtype)

def _get_default_value(feature_spec):
  if isinstance(feature_spec, parsing_ops.FixedLenFeature):
    return feature_spec.default_value
  else:
    return _dtype_to_nan(feature_spec.dtype)

class TensorFlowDataFrame(df.DataFrame):
  """TensorFlowDataFrame implements convenience functions using TensorFlow."""

  def run(self,
          num_batches=None,
          graph=None,
          session=None,
          start_queues=True,
          initialize_variables=True,
          **kwargs):
    """Builds and runs the columns of the `DataFrame` and yields batches.

    This is a generator that yields a dictionary mapping column names to
    evaluated columns.

    Args:
      num_batches: the maximum number of batches to produce. If none specified,
        the returned value will iterate through infinite batches.
      graph: the `Graph` in which the `DataFrame` should be built.
      session: the `Session` in which to run the columns of the `DataFrame`.
      start_queues: if true, queues will be started before running and halted
        after producting `n` batches.
      initialize_variables: if true, variables will be initialized.
      **kwargs: Additional keyword arguments e.g. `num_epochs`.

    Yields:
      A dictionary, mapping column names to the values resulting from running
      each column for a single batch.
    """
    if graph is None:
      graph = ops.get_default_graph()
    with graph.as_default():
      if session is None:
        session = sess.Session()
      self_built = self.build(**kwargs)
      keys = list(self_built.keys())
      cols = list(self_built.values())
      if initialize_variables:
        if variables.local_variables():
          session.run(variables.local_variables_initializer())
        if variables.global_variables():
          session.run(variables.global_variables_initializer())
      if start_queues:
        coord = coordinator.Coordinator()
        threads = qr.start_queue_runners(sess=session, coord=coord)
      i = 0
      while num_batches is None or i < num_batches:
        i += 1
        try:
          values = session.run(cols)
          yield collections.OrderedDict(zip(keys, values))
        except errors.OutOfRangeError:
          break
      if start_queues:
        coord.request_stop()
        coord.join(threads)

  def select_rows(self, boolean_series):
    """Returns a `DataFrame` with only the rows indicated by `boolean_series`.

    Note that batches may no longer have consistent size after calling
    `select_rows`, so the new `DataFrame` may need to be rebatched.
    For example:
    '''
    filtered_df = df.select_rows(df["country"] == "jp").batch(64)
    '''

    Args:
      boolean_series: a `Series` that evaluates to a boolean `Tensor`.

    Returns:
      A new `DataFrame` with the same columns as `self`, but selecting only the
      rows where `boolean_series` evaluated to `True`.
    """
    result = type(self)()
    for key, col in self._columns.items():
      try:
        result[key] = col.select_rows(boolean_series)
      except AttributeError as e:
        raise NotImplementedError((
            "The select_rows method is not implemented for Series type {}. "
            "Original error: {}").format(type(col), e))
    return result

  def split(self, index_series, proportion, batch_size=None):
    """Deterministically split a `DataFrame` into two `DataFrame`s.

    Note this split is only as deterministic as the underlying hash function;
    see `tf.string_to_hash_bucket_fast`.  The hash function is deterministic
    for a given binary, but may change occasionally.  The only way to achieve
    an absolute guarantee that the split `DataFrame`s do not change across runs
    is to materialize them.

    Note too that the allocation of a row to one partition or the
    other is evaluated independently for each row, so the exact number of rows
    in each partition is binomially distributed.

    Args:
      index_series: a `Series` of unique strings, whose hash will determine the
        partitioning; or the name in this `DataFrame` of such a `Series`.
        (This `Series` must contain strings because TensorFlow provides hash
        ops only for strings, and there are no number-to-string converter ops.)
      proportion: The proportion of the rows to select for the 'left'
        partition; the remaining (1 - proportion) rows form the 'right'
        partition.
      batch_size: the batch size to use when rebatching the left and right
        `DataFrame`s.  If None (default), the `DataFrame`s are not rebatched;
        thus their batches will have variable sizes, according to which rows
        are selected from each batch of the original `DataFrame`.

    Returns:
      Two `DataFrame`s containing the partitioned rows.
    """
    if isinstance(index_series, str):
      index_series = self[index_series]
    left_mask, = split_mask.SplitMask(proportion)(index_series)
    right_mask = ~left_mask
    left_rows = self.select_rows(left_mask)
    right_rows = self.select_rows(right_mask)

    if batch_size:
      left_rows = left_rows.batch(batch_size=batch_size, shuffle=False)
      right_rows = right_rows.batch(batch_size=batch_size, shuffle=False)

    return left_rows, right_rows

  def split_fast(self, index_series, proportion, batch_size,
                 base_batch_size=1000):
    """Deterministically split a `DataFrame` into two `DataFrame`s.

    Note this split is only as deterministic as the underlying hash function;
    see `tf.string_to_hash_bucket_fast`.  The hash function is deterministic
    for a given binary, but may change occasionally.  The only way to achieve
    an absolute guarantee that the split `DataFrame`s do not change across runs
    is to materialize them.

    Note too that the allocation of a row to one partition or the
    other is evaluated independently for each row, so the exact number of rows
    in each partition is binomially distributed.

    Args:
      index_series: a `Series` of unique strings, whose hash will determine the
        partitioning; or the name in this `DataFrame` of such a `Series`.
        (This `Series` must contain strings because TensorFlow provides hash
        ops only for strings, and there are no number-to-string converter ops.)
      proportion: The proportion of the rows to select for the 'left'
        partition; the remaining (1 - proportion) rows form the 'right'
        partition.
      batch_size: the batch size to use when rebatching the left and right
        `DataFrame`s.  If None (default), the `DataFrame`s are not rebatched;
        thus their batches will have variable sizes, according to which rows
        are selected from each batch of the original `DataFrame`.
      base_batch_size: the batch size to use for materialized data, prior to the
        split.

    Returns:
      Two `DataFrame`s containing the partitioned rows.
    """
    if isinstance(index_series, str):
      index_series = self[index_series]
    left_mask, = split_mask.SplitMask(proportion)(index_series)
    right_mask = ~left_mask
    self["left_mask__"] = left_mask
    self["right_mask__"] = right_mask
    # TODO(soergel): instead of base_batch_size can we just do one big batch?
    # avoid computing the hashes twice
    m = self.materialize_to_memory(batch_size=base_batch_size)
    left_rows_df = m.select_rows(m["left_mask__"])
    right_rows_df = m.select_rows(m["right_mask__"])
    del left_rows_df[["left_mask__", "right_mask__"]]
    del right_rows_df[["left_mask__", "right_mask__"]]

    # avoid recomputing the split repeatedly
    left_rows_df = left_rows_df.materialize_to_memory(batch_size=batch_size)
    right_rows_df = right_rows_df.materialize_to_memory(batch_size=batch_size)
    return left_rows_df, right_rows_df

  def run_one_batch(self):
    """Creates a new 'Graph` and `Session` and runs a single batch.

    Returns:
      A dictionary mapping column names to numpy arrays that contain a single
      batch of the `DataFrame`.
    """
    return list(self.run(num_batches=1))[0]

  def run_one_epoch(self):
    """Creates a new 'Graph` and `Session` and runs a single epoch.

    Naturally this makes sense only for DataFrames that fit in memory.

    Returns:
      A dictionary mapping column names to numpy arrays that contain a single
      epoch of the `DataFrame`.
    """
    # batches is a list of dicts of numpy arrays
    batches = [b for b in self.run(num_epochs=1)]

    # first invert that to make a dict of lists of numpy arrays
    pivoted_batches = {}
    for k in batches[0].keys():
      pivoted_batches[k] = []
    for b in batches:
      for k, v in b.items():
        pivoted_batches[k].append(v)

    # then concat the arrays in each column
    result = {k: np.concatenate(column_batches)
              for k, column_batches in pivoted_batches.items()}
    return result

  def materialize_to_memory(self, batch_size):
    unordered_dict_of_arrays = self.run_one_epoch()

    # there may already be an 'index' column, in which case from_ordereddict)
    # below will complain because it wants to generate a new one.
    # for now, just remove it.
    # TODO(soergel): preserve index history, potentially many levels deep
    del unordered_dict_of_arrays["index"]

    # the order of the columns in this dict is arbitrary; we just need it to
    # remain consistent.
    ordered_dict_of_arrays = collections.OrderedDict(unordered_dict_of_arrays)
    return TensorFlowDataFrame.from_ordereddict(ordered_dict_of_arrays,
                                                batch_size=batch_size)

  def batch(self,
            batch_size,
            shuffle=False,
            num_threads=1,
            queue_capacity=None,
            min_after_dequeue=None,
            seed=None):
    """Resize the batches in the `DataFrame` to the given `batch_size`.

    Args:
      batch_size: desired batch size.
      shuffle: whether records should be shuffled. Defaults to true.
      num_threads: the number of enqueueing threads.
      queue_capacity: capacity of the queue that will hold new batches.
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` with `batch_size` rows.
    """
    column_names = list(self._columns.keys())
    if shuffle:
      batcher = batch.ShuffleBatch(batch_size,
                                   output_names=column_names,
                                   num_threads=num_threads,
                                   queue_capacity=queue_capacity,
                                   min_after_dequeue=min_after_dequeue,
                                   seed=seed)
    else:
      batcher = batch.Batch(batch_size,
                            output_names=column_names,
                            num_threads=num_threads,
                            queue_capacity=queue_capacity)

    batched_series = batcher(list(self._columns.values()))
    dataframe = type(self)()
    dataframe.assign(**(dict(zip(column_names, batched_series))))
    return dataframe

  @classmethod
  def _from_csv_base(cls, filepatterns, get_default_values, has_header,
                     column_names, num_threads, enqueue_size,
                     batch_size, queue_capacity, min_after_dequeue, shuffle,
                     seed):
    """Create a `DataFrame` from CSV files.

    If `has_header` is false, then `column_names` must be specified. If
    `has_header` is true and `column_names` are specified, then `column_names`
    overrides the names in the header.

    Args:
      filepatterns: a list of file patterns that resolve to CSV files.
      get_default_values: a function that produces a list of default values for
        each column, given the column names.
      has_header: whether or not the CSV files have headers.
      column_names: a list of names for the columns in the CSV files.
      num_threads: the number of readers that will work in parallel.
      enqueue_size: block size for each read operation.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed lines.
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` that has columns corresponding to `features` and is filled
      with examples from `filepatterns`.

    Raises:
      ValueError: no files match `filepatterns`.
      ValueError: `features` contains the reserved name 'index'.
    """
    filenames = _expand_file_names(filepatterns)
    if not filenames:
      raise ValueError("No matching file names.")

    if column_names is None:
      if not has_header:
        raise ValueError("If column_names is None, has_header must be true.")
      with gfile.GFile(filenames[0]) as f:
        column_names = csv.DictReader(f).fieldnames

    if "index" in column_names:
      raise ValueError(
          "'index' is reserved and can not be used for a column name.")

    default_values = get_default_values(column_names)

    reader_kwargs = {"skip_header_lines": (1 if has_header else 0)}
    index, value = reader_source.TextFileSource(
        filenames,
        reader_kwargs=reader_kwargs,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        num_threads=num_threads,
        seed=seed)()
    parser = csv_parser.CSVParser(column_names, default_values)
    parsed = parser(value)

    column_dict = parsed._asdict()
    column_dict["index"] = index

    dataframe = cls()
    dataframe.assign(**column_dict)
    return dataframe

  @classmethod
  def from_csv(cls,
               filepatterns,
               default_values,
               has_header=True,
               column_names=None,
               num_threads=1,
               enqueue_size=None,
               batch_size=32,
               queue_capacity=None,
               min_after_dequeue=None,
               shuffle=True,
               seed=None):
    """Create a `DataFrame` from CSV files.

    If `has_header` is false, then `column_names` must be specified. If
    `has_header` is true and `column_names` are specified, then `column_names`
    overrides the names in the header.

    Args:
      filepatterns: a list of file patterns that resolve to CSV files.
      default_values: a list of default values for each column.
      has_header: whether or not the CSV files have headers.
      column_names: a list of names for the columns in the CSV files.
      num_threads: the number of readers that will work in parallel.
      enqueue_size: block size for each read operation.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed lines.
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` that has columns corresponding to `features` and is filled
      with examples from `filepatterns`.

    Raises:
      ValueError: no files match `filepatterns`.
      ValueError: `features` contains the reserved name 'index'.
    """

    def get_default_values(column_names):
      # pylint: disable=unused-argument
      return default_values

    return cls._from_csv_base(filepatterns, get_default_values, has_header,
                              column_names, num_threads,
                              enqueue_size, batch_size, queue_capacity,
                              min_after_dequeue, shuffle, seed)

  @classmethod
  def from_csv_with_feature_spec(cls,
                                 filepatterns,
                                 feature_spec,
                                 has_header=True,
                                 column_names=None,
                                 num_threads=1,
                                 enqueue_size=None,
                                 batch_size=32,
                                 queue_capacity=None,
                                 min_after_dequeue=None,
                                 shuffle=True,
                                 seed=None):
    """Create a `DataFrame` from CSV files, given a feature_spec.

    If `has_header` is false, then `column_names` must be specified. If
    `has_header` is true and `column_names` are specified, then `column_names`
    overrides the names in the header.

    Args:
      filepatterns: a list of file patterns that resolve to CSV files.
      feature_spec: a dict mapping column names to `FixedLenFeature` or
          `VarLenFeature`.
      has_header: whether or not the CSV files have headers.
      column_names: a list of names for the columns in the CSV files.
      num_threads: the number of readers that will work in parallel.
      enqueue_size: block size for each read operation.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed lines.
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` that has columns corresponding to `features` and is filled
      with examples from `filepatterns`.

    Raises:
      ValueError: no files match `filepatterns`.
      ValueError: `features` contains the reserved name 'index'.
    """

    def get_default_values(column_names):
      return [_get_default_value(feature_spec[name]) for name in column_names]

    dataframe = cls._from_csv_base(filepatterns, get_default_values, has_header,
                                   column_names, num_threads,
                                   enqueue_size, batch_size, queue_capacity,
                                   min_after_dequeue, shuffle, seed)

    # replace the dense columns with sparse ones in place in the dataframe
    for name in dataframe.columns():
      if name != "index" and isinstance(feature_spec[name],
                                        parsing_ops.VarLenFeature):
        strip_value = _get_default_value(feature_spec[name])
        (dataframe[name],) = sparsify.Sparsify(strip_value)(dataframe[name])

    return dataframe

  @classmethod
  def from_examples(cls,
                    filepatterns,
                    features,
                    reader_cls=io_ops.TFRecordReader,
                    num_threads=1,
                    enqueue_size=None,
                    batch_size=32,
                    queue_capacity=None,
                    min_after_dequeue=None,
                    shuffle=True,
                    seed=None):
    """Create a `DataFrame` from `tensorflow.Example`s.

    Args:
      filepatterns: a list of file patterns containing `tensorflow.Example`s.
      features: a dict mapping feature names to `VarLenFeature` or
        `FixedLenFeature`.
      reader_cls: a subclass of `tensorflow.ReaderBase` that will be used to
        read the `Example`s.
      num_threads: the number of readers that will work in parallel.
      enqueue_size: block size for each read operation.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed `Example`s
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` that has columns corresponding to `features` and is filled
      with `Example`s from `filepatterns`.

    Raises:
      ValueError: no files match `filepatterns`.
      ValueError: `features` contains the reserved name 'index'.
    """
    filenames = _expand_file_names(filepatterns)
    if not filenames:
      raise ValueError("No matching file names.")

    if "index" in features:
      raise ValueError(
          "'index' is reserved and can not be used for a feature name.")

    index, record = reader_source.ReaderSource(
        reader_cls,
        filenames,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        num_threads=num_threads,
        seed=seed)()
    parser = example_parser.ExampleParser(features)
    parsed = parser(record)

    column_dict = parsed._asdict()
    column_dict["index"] = index

    dataframe = cls()
    dataframe.assign(**column_dict)
    return dataframe

  @classmethod
  def from_pandas(cls,
                  pandas_dataframe,
                  num_threads=None,
                  enqueue_size=None,
                  batch_size=None,
                  queue_capacity=None,
                  min_after_dequeue=None,
                  shuffle=True,
                  seed=None,
                  data_name="pandas_data"):
    """Create a `tf.learn.DataFrame` from a `pandas.DataFrame`.

    Args:
      pandas_dataframe: `pandas.DataFrame` that serves as a data source.
      num_threads: the number of threads to use for enqueueing.
      enqueue_size: the number of rows to enqueue per step.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed `Example`s
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.
      data_name: a scope name identifying the data.

    Returns:
      A `tf.learn.DataFrame` that contains batches drawn from the given
      `pandas_dataframe`.
    """
    pandas_source = in_memory_source.PandasSource(
        pandas_dataframe,
        num_threads=num_threads,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        seed=seed,
        data_name=data_name)
    dataframe = cls()
    dataframe.assign(**(pandas_source()._asdict()))
    return dataframe

  @classmethod
  def from_numpy(cls,
                 numpy_array,
                 num_threads=None,
                 enqueue_size=None,
                 batch_size=None,
                 queue_capacity=None,
                 min_after_dequeue=None,
                 shuffle=True,
                 seed=None,
                 data_name="numpy_data"):
    """Creates a `tf.learn.DataFrame` from a `numpy.ndarray`.

    The returned `DataFrame` contains two columns: 'index' and 'value'. The
    'value' column contains a row from the array. The 'index' column contains
    the corresponding row number.

    Args:
      numpy_array: `numpy.ndarray` that serves as a data source.
      num_threads: the number of threads to use for enqueueing.
      enqueue_size: the number of rows to enqueue per step.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed `Example`s
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.
      data_name: a scope name identifying the data.

    Returns:
      A `tf.learn.DataFrame` that contains batches drawn from the given
      array.
    """
    numpy_source = in_memory_source.NumpySource(
        numpy_array,
        num_threads=num_threads,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        seed=seed,
        data_name=data_name)
    dataframe = cls()
    dataframe.assign(**(numpy_source()._asdict()))
    return dataframe

  @classmethod
  def from_ordereddict(cls,
                       ordered_dict_of_arrays,
                       num_threads=None,
                       enqueue_size=None,
                       batch_size=None,
                       queue_capacity=None,
                       min_after_dequeue=None,
                       shuffle=True,
                       seed=None,
                       data_name="numpy_data"):
    """Creates a `tf.learn.DataFrame` from an `OrderedDict` of `numpy.ndarray`.

    The returned `DataFrame` contains a column for each key of the dict plus an
    extra 'index' column. The 'index' column contains the row number. Each of
    the other columns contains a row from the corresponding array.

    Args:
      ordered_dict_of_arrays: `OrderedDict` of `numpy.ndarray` that serves as a
          data source.
      num_threads: the number of threads to use for enqueueing.
      enqueue_size: the number of rows to enqueue per step.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed `Example`s
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.
      data_name: a scope name identifying the data.

    Returns:
      A `tf.learn.DataFrame` that contains batches drawn from the given arrays.

    Raises:
      ValueError: `ordered_dict_of_arrays` contains the reserved name 'index'.
    """
    numpy_source = in_memory_source.OrderedDictNumpySource(
        ordered_dict_of_arrays,
        num_threads=num_threads,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        seed=seed,
        data_name=data_name)
    dataframe = cls()
    dataframe.assign(**(numpy_source()._asdict()))
    return dataframe


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jun 14 05:54:11 2020

@author: mathieumoog
"""

import cpmd
import filexyz
import numpy as np
import matplotlib.pyplot as plt
# MSMbuilder ( lacks CK validation )
from msmbuilder.msm import MarkovStateModel
from msmbuilder.msm import BayesianMarkovStateModel
from msmbuilder.utils import dump
# PyEMMMA ( has CK validation )
import pyemma as pe
from pyemma.datasets import double_well_discrete

def getDistance1Dsq( position1, position2, length):
    dist = position1-position2
    half_length = length*0.5
    if dist > half_length :
        dist -= length
    elif dist < -half_length:
        dist += length
    return dist*dist
def getDistanceOrtho( positions, index1, index2, cell_lengths ):
    dist=0
    for i in range(3):
        dist += getDistance1Dsq( positions[index1,i], positions[index2,i], cell_lengths[i] )
    return np.sqrt(dist)
def computeContactMatrix( positions, cell_lengths, cut_off ):
    nb_atoms = len(positions[:,0])
    matrix = np.zeros(( nb_atoms, nb_atoms ))
    for atom in range(nb_atoms):
        for atom2 in range(atom+1,nb_atoms):
            if getDistanceOrtho( positions, atom, atom2, cell_lengths ) < cut_off :
                matrix[atom,atom2] = 1
                matrix[atom2,atom] = 1
    return matrix
def computeTransitionMatrix( states, nb_states, tau, step_max ):
    nb_step = len(states)
    matrix = np.zeros((nb_states,nb_states))
    for step in range( nb_step-step_max ):
        matrix[ states[step], states[step+tau] ] += 1
    return matrix
def computeChapmanKolmogorov( matrix, nb_states ):
    matrix_ck = np.zeros((nb_states,nb_states),dtype=float)
    for state_i in range( nb_states ):
        for state_j in range( nb_states ):
            for i in range(nb_states):
                matrix_ck[ state_i, state_j ] += matrix[state_i,i]*matrix[i,state_j]
    return matrix_ck

volume=8.82
temperature=3000
# run_nb=1

path_sim = str( "/Users/mathieumoog/Documents/CO2/" +
               str(volume) + "/" +
               str(temperature) + "K/"
               # + str(run_nb) + "-run/"
               )

cell_lengths = np.ones(3)*volume

traj_path = str( path_sim + "TRAJEC_fdb_wrapped.xyz" )
traj = filexyz.readAsArray( traj_path )

nbC=32
nbO=64
nb_atoms=nbC+nbO
max_neigh=5
nb_step=len(traj[:,0,0])
cut_off = 1.75
min_stat=1000

# Build States
coordC = np.zeros( (nb_step,nbC), dtype=int )
coordO = np.zeros( (nb_step,nbO), dtype=int )
for step in range(nb_step):
    matrix = computeContactMatrix( traj[step,:,:], cell_lengths, cut_off)
    for carbon in range(0,nbC):
        coordC[ step, carbon ] = int( sum(matrix[carbon,:]) )
    for oxygen in range(nbC,nb_atoms):
        coordO[ step, oxygen-nbC ] = int( sum(matrix[oxygen,:]) )

c_min = coordC.min()
o_min = coordO.min()
# Adapting the labels to make sure they are in the 0-nb_states range
coordC -= c_min
coordO -= c_min

msm = MarkovStateModel( lag_time=1, n_timescales=6)
msm.fit( coordC[:,0] )
msm.timescales_

# Computing nb of states (max)
nb_states_C = coordC.max()+1
nb_states_O = coordO.max()+1

# Computing Equilibrium States Probabilities
coordC_hist = np.zeros( nb_states_C )
ones_ = np.ones((nb_step,nbC), dtype=int )
for i in range( nb_states_C ):
    coordC_hist[i] = sum( ones_[ coordC == i ] )
# Clean marginal states
# for state in range( nb_states_C ):
#     if coordC_hist[state] < min_stat:
#         mask_to_clean = coordC[ :, : ]
coordC_hist /= sum(coordC_hist[:])

# Computing Equilibrium States Probabilities, cleaning marginals
ones_ = np.ones((nb_step,nbO), dtype=int )
coordO_hist = np.zeros( nb_states_O )
for i in range( nb_states_O ):
    coordO_hist[i] = sum( ones_[ coordO == i ] )
coordO_hist /= sum(coordO_hist[:])

# Plotting Oxygens
plt.figure()
plt.plot(coordC_hist,"b.-")
plt.plot(coordO_hist,"r.-")
plt.legend(["C states","O states"])
plt.show()

dt=5*0.001
frac = 0.75
max_step=int(nb_step*frac)
nb_tau_min=int(250)
nb_tau_max=int(2*nb_tau_min)

# Computing Transition Matrix for a given tau
matrix_tot=np.zeros((nb_states_C,nb_states_C,nb_tau_max), dtype=float )
matrix_tot_ck=np.zeros((nb_states_C,nb_states_C,nb_tau_min), dtype=float )
for tau in range(nb_tau_max):
    matrix = np.zeros((nb_states_C,nb_states_C),dtype=float)
    for carbon in range(nbC):
        matrix += computeTransitionMatrix( coordC[:,carbon], nb_states_C, tau+1, max_step )
    for state in range(nb_states_C):
        matrix[state,:] /= sum( matrix[state,:] )
    matrix_tot[:,:,tau] = matrix[:,:]
    if tau < nb_tau_min:
        matrix_tot_ck[:,:,tau] = computeChapmanKolmogorov( matrix_tot[:,:,tau], nb_states_C )

carbon_target=3
matrix_markov = np.zeros( (4,4,nb_tau_min), dtype=float )
matrix_markov_ck = np.zeros( (4,4,nb_tau_min), dtype=float )
for tau in range(1,nb_tau_min+1):
    msm_matrix = MarkovStateModel( lag_time=tau, reversible_type="mle" ,n_timescales=nb_states_C, ergodic_cutoff="on", sliding_window=True, verbose=True)
    msm_matrix.fit( coordC[:,carbon_target] )
    matrix_markov[:,:,tau-1] = msm_matrix.transmat_
    for state_i in range( len(matrix_markov) ):
        for state_j in range( len(matrix_markov) ):
            for i in range( len(matrix_markov) ):
                matrix_markov_ck[ state_i, state_j, tau-1 ] +=  matrix_markov[state_i,i,tau-1]*matrix_markov[i,state_j,tau-1]

# PyEMMA
lags = [1,5,10,15,20,50,100,200]
implied_timescales = pe.msm.its(dtrajs=coordC[:,carbon_target].tolist(),lags=lags)
pe.plots.plot_implied_timescales(implied_timescales,units='time-steps', ylog=False)

M = pe.msm.estimate_markov_model(dtrajs=coordC[:,carbon_target].tolist(), lag = 10 )
cktest = M.cktest(nsets=3)
cktplt = pe.plots.plot_cktest(cktest)

plt.figure()
plt.xlabel("Time lag (ps)")
plt.ylabel("P_ij, P_ij^CK")
# plt.plot( np.arange(0,dt*nb_tau_max,dt*1), matrix_tot[0,0,:], "k-" )
# plt.plot( np.arange(0,dt*nb_tau_max,dt*2), matrix_tot_ck[0,0,:], "k--" )
plt.plot( np.arange(0,dt*nb_tau_min,dt*1), matrix_markov[0,0,:], "k-" )
plt.plot( np.arange(0,2*dt*nb_tau_min,dt*2), matrix_markov_ck[0,0,:], "k--" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*1), matrix_tot[1,1,:], "r-" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), matrix_tot_ck[1,1,:], "r--" )
plt.plot( np.arange(0,dt*nb_tau_min,dt*1), matrix_markov[0,1,:], "k-" )
plt.plot( np.arange(0,2*dt*nb_tau_min,dt*2), matrix_markov_ck[0,1,:], "k--" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*1), matrix_tot[1,2,:], "b-" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), matrix_tot_ck[1,2,:], "b--" )
plt.plot( np.arange(0,dt*nb_tau_min,dt*1), matrix_markov[0,2,:], "k-" )
plt.plot( np.arange(0,2*dt*nb_tau_min,dt*2), matrix_markov_ck[0,2,:], "k--" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*1), matrix_tot[1,3,:], "g-" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), matrix_tot_ck[1,3,:], "g--" )
plt.plot( np.arange(0,dt*nb_tau_min,dt*1), matrix_markov[0,3,:], "k-" )
plt.plot( np.arange(0,2*dt*nb_tau_min,dt*2), matrix_markov_ck[0,3,:], "k--" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*1), matrix_tot[1,4,:], "m-" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), matrix_tot_ck[1,4,:], "m--" )
plt.show()

rmseC = np.zeros(nb_tau_min, dtype=float)
for tau in range(nb_tau_min):
    mat = matrix_tot[:,:,2*tau]-matrix_tot_ck[:,:,tau]
    rmseC[tau] = sum(sum( mat*mat ))/(nb_states_C*nb_states_C)

plt.figure()
plt.xlabel("Time lag (ps)")
plt.ylabel("RMSE C (%)")
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), rmseC*100 )
plt.show()

matrix_tot=np.zeros((nb_states_O,nb_states_O,nb_tau_max), dtype=float )
matrix_tot_ck=np.zeros((nb_states_O,nb_states_O,nb_tau_min), dtype=float )
for tau in range(nb_tau_max):
    matrix = np.zeros((nb_states_O,nb_states_O),dtype=float)
    for carbon in range(nbC):
        matrix += computeTransitionMatrix( coordO[:,carbon], nb_states_O, tau, max_step )
    for state in range(nb_states_O):
        matrix[state,:] /= sum( matrix[state,:] )
    matrix_tot[:,:,tau] = matrix[:,:]
    if tau < nb_tau_min:
        matrix_tot_ck[:,:,tau] = computeChapmanKolmogorov( matrix_tot[:,:,tau], nb_states_O )

plt.figure()
plt.xlabel("Time lag (ps)")
plt.ylabel("P_ij, P_ij^CK")
plt.plot( np.arange(0,dt*nb_tau_max,dt*1), matrix_tot[0,0,:], "k-" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), matrix_tot_ck[0,0,:], "k--" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*1), matrix_tot[1,1,:], "r-" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), matrix_tot_ck[1,1,:], "r--" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*1), matrix_tot[2,2,:], "b-" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), matrix_tot_ck[2,2,:], "b--" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*1), matrix_tot[3,3,:], "g-" )
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), matrix_tot_ck[3,3,:], "g--" )
plt.show()

rmseO = np.zeros(nb_tau_min, dtype=float)
for tau in range(nb_tau_min):
    mat = matrix_tot[:,:,2*tau]-matrix_tot_ck[:,:,tau]
    rmseO[tau] = sum(sum( mat*mat ))/(nb_states_O*nb_states_O)

plt.figure()
plt.xlabel("Time lag (ps)")
plt.ylabel("RMSE O (%)")
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), rmseO*100 )
plt.show()

plt.figure()
plt.xlabel("Time lag (ps)")
plt.ylabel("RMSE all (%)")
plt.plot( np.arange(0,dt*nb_tau_max,dt*2), (rmseO+rmseC)*100*0.5 )
plt.show()


# http://github.com/timestocome

# Attempt to predict gold prices and find outliers
# http://web.ipac.caltech.edu/staff/fmasci/home/astro_refs/TestForRandomness_RunsTest.pdf

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

######################################################################
# load data
########################################################################
# read in gold file
data = pd.read_csv('data/Gold_all.csv', parse_dates=True, index_col=0)
data = data[['Open']]

# convert to log values
#data['Open'] = np.log(data['Open'])

data['Open'] = pd.to_numeric(data['Open'], errors='coerce')
data['Volatility'] = data['Open'] - data['Open'].shift(1)
data = data.dropna()

gold_standard = data.loc[data.index < '01-01-1971']
gold = data.loc[data.index > '01-01-1971']

print(len(gold_standard), len(gold))
########################################################################

# try to fit linear regression model
from sklearn import linear_model

x1 = np.arange(1, len(gold)+ 1)
x2 = x1 **2
x3 = x1 **3
x4 = x1 **4     # best so far

x = [x1, x2, x3, x4]
x = np.reshape(x, (4, len(gold))).T
print(x.shape)

regression = linear_model.LinearRegression()
regression.fit(x, gold['Open'])
coeffs = regression.coef_
intercept = regression.intercept_

print(coeffs[0], coeffs[1])
gold['Regression'] = intercept + coeffs[0] * x1 + coeffs[1] * x2 + coeffs[2] * x3 + coeffs[3] * x4
gold['Residuals'] = gold['Open'] - gold['Regression']
std_regression = gold['Regression'].std()
std_open = gold['Open'].std()

##################################################################
# Run's Test, part 3 of paper

gold_mean = gold['Open'].mean()
runs = gold['Open'] > gold['Regression']

# convert runs data to number of runs
R = 0
r_prev = runs[0]
for r in runs:
    if r != r_prev: R += 1
    r_prev = r

T = len(runs)
Ta = runs.sum()
Tb = T - Ta

E = (T + 2 * Ta * Tb) / T           # expected runs
V = (2 * Ta * Tb * (2*Ta*Tb - T)) / (T **2 * (T - 1))  # variance of runs

Z1 = (R - E) / std_open
Z2 = (R -E) / std_regression
print("Run's Test Results")
print("R %lf, E %lf, V %lf" % (R, E, V))
print("Z (not random if Z > +/- 2.5)", Z1, Z2)

print("Regression:")
print("Start date", gold.ix[-1])
print("Start step", len(x))
print("intercept", intercept)
print("coeff", coeffs)

#######################################################################
# predict next 12 months ~253 trading days

dates = pd.bdate_range('1971-01-01', '2018-12-31')

x1 = np.arange(1, len(dates) + 1)
x2 = x1 **2
x3 = x1 **3
x4 = x1 **4

gold_futures = intercept + coeffs[0] * x1 + coeffs[1] * x2 + coeffs[2] * x3 + coeffs[3] * x4
std_regression = gold['Regression'].std()

predicted = pd.DataFrame(data=gold_futures, index=dates)
predicted.index.name = 'Date'
predicted.columns = ['Open']

actual = pd.read_csv('data/Gold_all.csv', parse_dates=True, index_col=0)
actual = actual.loc[actual.index > '01-01-1971']
actual = actual['Open']

plt.figure(figsize=(18, 16))
plt.plot(actual, label="Actual")
plt.plot(predicted, label="Predicted")
plt.plot(predicted - std_regression, label='Predicted - std')
plt.plot(predicted + std_regression, label='Predicted + std')
plt.legend(loc='best')
plt.title("Gold 1971 - predicted 2019")
plt.savefig("Gold_Predictions_2018.png")
plt.show()


# Author: Vlad Niculae
# Licence: BSD 3 clause

import numpy as np

from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_warns
from sklearn.utils.testing import ignore_warnings

from sklearn.linear_model import (orthogonal_mp, orthogonal_mp_gram,
                                  OrthogonalMatchingPursuit,
                                  OrthogonalMatchingPursuitCV,
                                  LinearRegression)
from sklearn.utils import check_random_state
from sklearn.datasets import make_sparse_coded_signal

n_samples, n_features, n_nonzero_coefs, n_targets = 20, 30, 5, 3
y, X, gamma = make_sparse_coded_signal(n_targets, n_features, n_samples,
                                       n_nonzero_coefs, random_state=0)
G, Xy = np.dot(X.T, X), np.dot(X.T, y)
# this makes X (n_samples, n_features)
# and y (n_samples, 3)

def test_correct_shapes():
    assert_equal(orthogonal_mp(X, y[:, 0], n_nonzero_coefs=5).shape,
                 (n_features,))
    assert_equal(orthogonal_mp(X, y, n_nonzero_coefs=5).shape,
                 (n_features, 3))

def test_correct_shapes_gram():
    assert_equal(orthogonal_mp_gram(G, Xy[:, 0], n_nonzero_coefs=5).shape,
                 (n_features,))
    assert_equal(orthogonal_mp_gram(G, Xy, n_nonzero_coefs=5).shape,
                 (n_features, 3))

def test_n_nonzero_coefs():
    assert_true(np.count_nonzero(orthogonal_mp(X, y[:, 0],
                                 n_nonzero_coefs=5)) <= 5)
    assert_true(np.count_nonzero(orthogonal_mp(X, y[:, 0], n_nonzero_coefs=5,
                                               precompute=True)) <= 5)

def test_tol():
    tol = 0.5
    gamma = orthogonal_mp(X, y[:, 0], tol=tol)
    gamma_gram = orthogonal_mp(X, y[:, 0], tol=tol, precompute=True)
    assert_true(np.sum((y[:, 0] - np.dot(X, gamma)) ** 2) <= tol)
    assert_true(np.sum((y[:, 0] - np.dot(X, gamma_gram)) ** 2) <= tol)

def test_with_without_gram():
    assert_array_almost_equal(
        orthogonal_mp(X, y, n_nonzero_coefs=5),
        orthogonal_mp(X, y, n_nonzero_coefs=5, precompute=True))

def test_with_without_gram_tol():
    assert_array_almost_equal(
        orthogonal_mp(X, y, tol=1.),
        orthogonal_mp(X, y, tol=1., precompute=True))

def test_unreachable_accuracy():
    assert_array_almost_equal(
        orthogonal_mp(X, y, tol=0),
        orthogonal_mp(X, y, n_nonzero_coefs=n_features))

    assert_array_almost_equal(
        assert_warns(RuntimeWarning, orthogonal_mp, X, y, tol=0,
                     precompute=True),
        orthogonal_mp(X, y, precompute=True,
                      n_nonzero_coefs=n_features))

def test_bad_input():
    assert_raises(ValueError, orthogonal_mp, X, y, tol=-1)
    assert_raises(ValueError, orthogonal_mp, X, y, n_nonzero_coefs=-1)
    assert_raises(ValueError, orthogonal_mp, X, y,
                  n_nonzero_coefs=n_features + 1)
    assert_raises(ValueError, orthogonal_mp_gram, G, Xy, tol=-1)
    assert_raises(ValueError, orthogonal_mp_gram, G, Xy, n_nonzero_coefs=-1)
    assert_raises(ValueError, orthogonal_mp_gram, G, Xy,
                  n_nonzero_coefs=n_features + 1)

def test_perfect_signal_recovery():
    idx, = gamma[:, 0].nonzero()
    gamma_rec = orthogonal_mp(X, y[:, 0], 5)
    gamma_gram = orthogonal_mp_gram(G, Xy[:, 0], 5)
    assert_array_equal(idx, np.flatnonzero(gamma_rec))
    assert_array_equal(idx, np.flatnonzero(gamma_gram))
    assert_array_almost_equal(gamma[:, 0], gamma_rec, decimal=2)
    assert_array_almost_equal(gamma[:, 0], gamma_gram, decimal=2)

def test_estimator():
    omp = OrthogonalMatchingPursuit(n_nonzero_coefs=n_nonzero_coefs)
    omp.fit(X, y[:, 0])
    assert_equal(omp.coef_.shape, (n_features,))
    assert_equal(omp.intercept_.shape, ())
    assert_true(np.count_nonzero(omp.coef_) <= n_nonzero_coefs)

    omp.fit(X, y)
    assert_equal(omp.coef_.shape, (n_targets, n_features))
    assert_equal(omp.intercept_.shape, (n_targets,))
    assert_true(np.count_nonzero(omp.coef_) <= n_targets * n_nonzero_coefs)

    omp.set_params(fit_intercept=False, normalize=False)

    omp.fit(X, y[:, 0])
    assert_equal(omp.coef_.shape, (n_features,))
    assert_equal(omp.intercept_, 0)
    assert_true(np.count_nonzero(omp.coef_) <= n_nonzero_coefs)

    omp.fit(X, y)
    assert_equal(omp.coef_.shape, (n_targets, n_features))
    assert_equal(omp.intercept_, 0)
    assert_true(np.count_nonzero(omp.coef_) <= n_targets * n_nonzero_coefs)

def test_identical_regressors():
    newX = X.copy()
    newX[:, 1] = newX[:, 0]
    gamma = np.zeros(n_features)
    gamma[0] = gamma[1] = 1.
    newy = np.dot(newX, gamma)
    assert_warns(RuntimeWarning, orthogonal_mp, newX, newy, 2)

def test_swapped_regressors():
    gamma = np.zeros(n_features)
    # X[:, 21] should be selected first, then X[:, 0] selected second,
    # which will take X[:, 21]'s place in case the algorithm does
    # column swapping for optimization (which is the case at the moment)
    gamma[21] = 1.0
    gamma[0] = 0.5
    new_y = np.dot(X, gamma)
    new_Xy = np.dot(X.T, new_y)
    gamma_hat = orthogonal_mp(X, new_y, 2)
    gamma_hat_gram = orthogonal_mp_gram(G, new_Xy, 2)
    assert_array_equal(np.flatnonzero(gamma_hat), [0, 21])
    assert_array_equal(np.flatnonzero(gamma_hat_gram), [0, 21])

def test_no_atoms():
    y_empty = np.zeros_like(y)
    Xy_empty = np.dot(X.T, y_empty)
    gamma_empty = ignore_warnings(orthogonal_mp)(X, y_empty, 1)
    gamma_empty_gram = ignore_warnings(orthogonal_mp)(G, Xy_empty, 1)
    assert_equal(np.all(gamma_empty == 0), True)
    assert_equal(np.all(gamma_empty_gram == 0), True)

def test_omp_path():
    path = orthogonal_mp(X, y, n_nonzero_coefs=5, return_path=True)
    last = orthogonal_mp(X, y, n_nonzero_coefs=5, return_path=False)
    assert_equal(path.shape, (n_features, n_targets, 5))
    assert_array_almost_equal(path[:, :, -1], last)
    path = orthogonal_mp_gram(G, Xy, n_nonzero_coefs=5, return_path=True)
    last = orthogonal_mp_gram(G, Xy, n_nonzero_coefs=5, return_path=False)
    assert_equal(path.shape, (n_features, n_targets, 5))
    assert_array_almost_equal(path[:, :, -1], last)

def test_omp_return_path_prop_with_gram():
    path = orthogonal_mp(X, y, n_nonzero_coefs=5, return_path=True,
                         precompute=True)
    last = orthogonal_mp(X, y, n_nonzero_coefs=5, return_path=False,
                         precompute=True)
    assert_equal(path.shape, (n_features, n_targets, 5))
    assert_array_almost_equal(path[:, :, -1], last)

def test_omp_cv():
    y_ = y[:, 0]
    gamma_ = gamma[:, 0]
    ompcv = OrthogonalMatchingPursuitCV(normalize=True, fit_intercept=False,
                                        max_iter=10, cv=5)
    ompcv.fit(X, y_)
    assert_equal(ompcv.n_nonzero_coefs_, n_nonzero_coefs)
    assert_array_almost_equal(ompcv.coef_, gamma_)
    omp = OrthogonalMatchingPursuit(normalize=True, fit_intercept=False,
                                    n_nonzero_coefs=ompcv.n_nonzero_coefs_)
    omp.fit(X, y_)
    assert_array_almost_equal(ompcv.coef_, omp.coef_)

def test_omp_reaches_least_squares():
    # Use small simple data; it's a sanity check but OMP can stop early
    rng = check_random_state(0)
    n_samples, n_features = (10, 8)
    n_targets = 3
    X = rng.randn(n_samples, n_features)
    Y = rng.randn(n_samples, n_targets)
    omp = OrthogonalMatchingPursuit(n_nonzero_coefs=n_features)
    lstsq = LinearRegression()
    omp.fit(X, Y)
    lstsq.fit(X, Y)
    assert_array_almost_equal(omp.coef_, lstsq.coef_)


from __future__ import absolute_import

from collections import Iterable, OrderedDict, Sequence
import difflib
import itertools
import re
import textwrap
import warnings

import numpy as np
import sys
from six import string_types, reraise

from ..models import (
    BoxSelectTool, BoxZoomTool, CategoricalAxis,
    TapTool, CrosshairTool, DataRange1d, DatetimeAxis,
    FactorRange, Grid, HelpTool, HoverTool, LassoSelectTool, Legend, LegendItem, LinearAxis,
    LogAxis, PanTool, ZoomInTool, ZoomOutTool, PolySelectTool, ContinuousTicker,
    SaveTool, Range, Range1d, UndoTool, RedoTool, ResetTool, ResizeTool, Tool,
    WheelPanTool, WheelZoomTool, ColumnarDataSource, ColumnDataSource, GlyphRenderer,
    LogScale, LinearScale, CategoricalScale)

from ..core.properties import ColorSpec, Datetime, value, field
from ..transform import stack
from ..util.dependencies import import_optional
from ..util.deprecation import deprecated
from ..util.string import nice_join

pd = import_optional('pandas')

DEFAULT_PALETTE = ["#f22c40", "#5ab738", "#407ee7", "#df5320", "#00ad9c", "#c33ff3"]

def _stack(stackers, spec0, spec1, **kw):
    for name in (spec0, spec1):
        if name in kw:
            raise ValueError("Stack property '%s' cannot appear in keyword args" % name)

    lengths = { len(x) for x in kw.values() if isinstance(x, (list, tuple)) }

    # lengths will be empty if there are no kwargs supplied at all
    if len(lengths) > 0:
        if len(lengths) != 1:
            raise ValueError("Keyword argument sequences for broadcasting must all be the same lengths. Got lengths: %r" % sorted(list(lengths)))
        if lengths.pop() != len(stackers):
            raise ValueError("Keyword argument sequences for broadcasting must be the same length as stackers")

    s0 = []
    s1 = []

    _kw = []

    for i, val in enumerate(stackers):
        d  = {}
        s0 = list(s1)
        s1.append(val)

        d[spec0] = stack(*s0)
        d[spec1] = stack(*s1)

        for k, v in kw.items():
            if isinstance(v, (list, tuple)):
                d[k] = v[i]
            else:
                d[k] = v

        _kw.append(d)

    return _kw

def get_default_color(plot=None):
    colors = [
        "#1f77b4",
        "#ff7f0e", "#ffbb78",
        "#2ca02c", "#98df8a",
        "#d62728", "#ff9896",
        "#9467bd", "#c5b0d5",
        "#8c564b", "#c49c94",
        "#e377c2", "#f7b6d2",
        "#7f7f7f",
        "#bcbd22", "#dbdb8d",
        "#17becf", "#9edae5"
    ]
    if plot:
        renderers = plot.renderers
        renderers = [x for x in renderers if x.__view_model__ == "GlyphRenderer"]
        num_renderers = len(renderers)
        return colors[num_renderers]
    else:
        return colors[0]

def get_default_alpha(plot=None):
    return 1.0

def _pop_renderer_args(kwargs):
    result = dict(data_source=kwargs.pop('source', ColumnDataSource()))
    for attr in ['name', 'x_range_name', 'y_range_name', 'level', 'view', 'visible', 'muted']:
        val = kwargs.pop(attr, None)
        if val:
            result[attr] = val
    return result

def _pop_colors_and_alpha(glyphclass, kwargs, prefix="", default_alpha=1.0):
    """
    Given a kwargs dict, a prefix, and a default value, looks for different
    color and alpha fields of the given prefix, and fills in the default value
    if it doesn't exist.
    """
    result = dict()

    # TODO: The need to do this and the complexity of managing this kind of
    # thing throughout the codebase really suggests that we need to have
    # a real stylesheet class, where defaults and Types can declaratively
    # substitute for this kind of imperative logic.
    color = kwargs.pop(prefix + "color", get_default_color())
    for argname in ("fill_color", "line_color"):
        if argname not in glyphclass.properties():
            continue
        result[argname] = kwargs.pop(prefix + argname, color)

    # NOTE: text fill color should really always default to black, hard coding
    # this here now until the stylesheet solution exists
    if "text_color" in glyphclass.properties():
        result["text_color"] = kwargs.pop(prefix + "text_color", "black")

    alpha = kwargs.pop(prefix + "alpha", default_alpha)
    for argname in ("fill_alpha", "line_alpha", "text_alpha"):
        if argname not in glyphclass.properties():
            continue
        result[argname] = kwargs.pop(prefix + argname, alpha)

    return result

def _get_legend_item_label(kwargs):
    legend = kwargs.pop('legend', None)
    source = kwargs.get('source')
    legend_item_label = None
    if legend:
        if isinstance(legend, string_types):
            # Do the simple thing first
            legend_item_label = value(legend)
            # But if there's a source - try and do something smart
            if source and hasattr(source, 'column_names'):
                if legend in source.column_names:
                    legend_item_label = field(legend)
        else:
            legend_item_label = legend
    return legend_item_label

_GLYPH_SOURCE_MSG = """
Supplying a user-defined data source AND iterable values to glyph methods is deprecated.

See https://github.com/bokeh/bokeh/issues/2056 for more information.
"""

def _process_sequence_literals(glyphclass, kwargs, source, is_user_source):
    dataspecs = glyphclass.dataspecs_with_props()
    for var, val in kwargs.items():

        # ignore things that are not iterable
        if not isinstance(val, Iterable):
            continue
        # pass dicts (i.e., values or fields) on as-is
        if isinstance(val, dict):
            continue
        # let any non-dataspecs do their own validation (e.g., line_dash properties)
        if var not in dataspecs:
            continue
        # strings sequences are handled by the dataspec as-is
        if isinstance(val, string_types):
            continue
        # similarly colorspecs handle color tuple sequences as-is
        if (isinstance(dataspecs[var].property, ColorSpec) and isinstance(val, tuple)):
            continue

        if isinstance(val, np.ndarray) and val.ndim != 1:
            raise RuntimeError("Columns need to be 1D (%s is not)" % var)

        if is_user_source:
            deprecated(_GLYPH_SOURCE_MSG)

        source.add(val, name=var)
        kwargs[var] = var

def _make_glyph(glyphclass, kws, extra):
    if extra is None:
        return None
    kws = kws.copy()
    kws.update(extra)
    return glyphclass(**kws)

def _update_legend(plot, legend_item_label, glyph_renderer):
    # Get the plot's legend
    legends = plot.select(type=Legend)
    if not legends:
        legend = Legend()
        plot.add_layout(legend)
    elif len(legends) == 1:
        legend = legends[0]
    else:
        raise RuntimeError("Plot %s configured with more than one legend renderer" % plot)

    # If there is an existing legend with a matching label, then put the
    # renderer on that (if the source matches). Otherwise add a new one.
    added = False
    for item in legend.items:
        if item.label == legend_item_label:
            if item.label.get('value'):
                item.renderers.append(glyph_renderer)
                added = True
                break
            if item.label.get('field') and \
                    glyph_renderer.data_source is item.renderers[0].data_source:
                item.renderers.append(glyph_renderer)
                added = True
                break
    if not added:
        new_item = LegendItem(label=legend_item_label, renderers=[glyph_renderer])
        legend.items.append(new_item)

def _get_range(range_input):
    if range_input is None:
        return DataRange1d()
    if pd and isinstance(range_input, pd.core.groupby.GroupBy):
        return FactorRange(factors=sorted(list(range_input.groups.keys())))
    if isinstance(range_input, Range):
        return range_input
    if isinstance(range_input, Sequence):
        if all(isinstance(x, string_types) for x in range_input):
            return FactorRange(factors=list(range_input))
        if len(range_input) == 2:
            try:
                return Range1d(start=range_input[0], end=range_input[1])
            except ValueError:  # @mattpap suggests ValidationError instead
                pass
    raise ValueError("Unrecognized range input: '%s'" % str(range_input))

def _get_scale(range_input, axis_type):
    if isinstance(range_input, (DataRange1d, Range1d)) and axis_type in ["linear", "datetime", "auto", None]:
        return LinearScale()
    elif isinstance(range_input, (DataRange1d, Range1d)) and axis_type == "log":
        return LogScale()
    elif isinstance(range_input, FactorRange):
        return CategoricalScale()
    else:
        raise ValueError("Unable to determine proper scale for: '%s'" % str(range_input))

def _get_axis_class(axis_type, range_input):
    if axis_type is None:
        return None
    elif axis_type == "linear":
        return LinearAxis
    elif axis_type == "log":
        return LogAxis
    elif axis_type == "datetime":
        return DatetimeAxis
    elif axis_type == "auto":
        if isinstance(range_input, FactorRange):
            return CategoricalAxis
        elif isinstance(range_input, Range1d):
            try:
                # Easier way to validate type of Range1d parameters
                Datetime.validate(Datetime(), range_input.start)
                return DatetimeAxis
            except ValueError:
                pass
        return LinearAxis
    else:
        raise ValueError("Unrecognized axis_type: '%r'" % axis_type)

def _get_num_minor_ticks(axis_class, num_minor_ticks):
    if isinstance(num_minor_ticks, int):
        if num_minor_ticks <= 1:
            raise ValueError("num_minor_ticks must be > 1")
        return num_minor_ticks
    if num_minor_ticks is None:
        return 0
    if num_minor_ticks == 'auto':
        if axis_class is LogAxis:
            return 10
        return 5

_known_tools = {
    "pan": lambda: PanTool(dimensions='both'),
    "xpan": lambda: PanTool(dimensions='width'),
    "ypan": lambda: PanTool(dimensions='height'),
    "wheel_zoom": lambda: WheelZoomTool(dimensions='both'),
    "xwheel_zoom": lambda: WheelZoomTool(dimensions='width'),
    "ywheel_zoom": lambda: WheelZoomTool(dimensions='height'),
    "zoom_in": lambda: ZoomInTool(dimensions='both'),
    "xzoom_in": lambda: ZoomInTool(dimensions='width'),
    "yzoom_in": lambda: ZoomInTool(dimensions='height'),
    "zoom_out": lambda: ZoomOutTool(dimensions='both'),
    "xzoom_out": lambda: ZoomOutTool(dimensions='width'),
    "yzoom_out": lambda: ZoomOutTool(dimensions='height'),
    "xwheel_pan": lambda: WheelPanTool(dimension="width"),
    "ywheel_pan": lambda: WheelPanTool(dimension="height"),
    "resize": lambda: ResizeTool(),
    "click": lambda: TapTool(behavior="inspect"),
    "tap": lambda: TapTool(),
    "crosshair": lambda: CrosshairTool(),
    "box_select": lambda: BoxSelectTool(),
    "xbox_select": lambda: BoxSelectTool(dimensions='width'),
    "ybox_select": lambda: BoxSelectTool(dimensions='height'),
    "poly_select": lambda: PolySelectTool(),
    "lasso_select": lambda: LassoSelectTool(),
    "box_zoom": lambda: BoxZoomTool(dimensions='both'),
    "xbox_zoom": lambda: BoxZoomTool(dimensions='width'),
    "ybox_zoom": lambda: BoxZoomTool(dimensions='height'),
    "hover": lambda: HoverTool(tooltips=[
        ("index", "$index"),
        ("data (x, y)", "($x, $y)"),
        ("canvas (x, y)", "($sx, $sy)"),
    ]),
    "save": lambda: SaveTool(),
    "previewsave": "save",
    "undo": lambda: UndoTool(),
    "redo": lambda: RedoTool(),
    "reset": lambda: ResetTool(),
    "help": lambda: HelpTool(),
}

def _tool_from_string(name):
    """ Takes a string and returns a corresponding `Tool` instance. """
    known_tools = sorted(_known_tools.keys())

    if name in known_tools:
        tool_fn = _known_tools[name]

        if isinstance(tool_fn, string_types):
            tool_fn = _known_tools[tool_fn]

        return tool_fn()
    else:
        matches, text = difflib.get_close_matches(name.lower(), known_tools), "similar"

        if not matches:
            matches, text = known_tools, "possible"

        raise ValueError("unexpected tool name '%s', %s tools are %s" % (name, text, nice_join(matches)))

def _process_axis_and_grid(plot, axis_type, axis_location, minor_ticks, axis_label, rng, dim):
    axiscls = _get_axis_class(axis_type, rng)
    if axiscls:

        if axiscls is LogAxis:
            if dim == 0:
                plot.x_scale = LogScale()
            elif dim == 1:
                plot.y_scale = LogScale()
            else:
                raise ValueError("received invalid dimension value: %r" % dim)

        # this is so we can get a ticker off the axis, even if we discard it
        axis = axiscls(plot=plot if axis_location else None)

        if isinstance(axis.ticker, ContinuousTicker):
            axis.ticker.num_minor_ticks = _get_num_minor_ticks(axiscls, minor_ticks)

        axis_label = axis_label
        if axis_label:
            axis.axis_label = axis_label

        grid = Grid(plot=plot, dimension=dim, ticker=axis.ticker); grid

        if axis_location is not None:
            getattr(plot, axis_location).append(axis)

def _process_tools_arg(plot, tools):
    """ Adds tools to the plot object

    Args:
        plot (Plot): instance of a plot object
        tools (seq[Tool or str]|str): list of tool types or string listing the
            tool names. Those are converted using the _tool_from_string
            function. I.e.: `wheel_zoom,box_zoom,reset`.

    Returns:
        list of Tools objects added to plot, map of supplied string names to tools
    """
    tool_objs = []
    tool_map = {}
    temp_tool_str = ""
    repeated_tools = []

    if isinstance(tools, (list, tuple)):
        for tool in tools:
            if isinstance(tool, Tool):
                tool_objs.append(tool)
            elif isinstance(tool, string_types):
                temp_tool_str += tool + ','
            else:
                raise ValueError("tool should be a string or an instance of Tool class")
        tools = temp_tool_str

    for tool in re.split(r"\s*,\s*", tools.strip()):
        # re.split will return empty strings; ignore them.
        if tool == "":
            continue

        tool_obj = _tool_from_string(tool)
        tool_objs.append(tool_obj)
        tool_map[tool] = tool_obj

    for typename, group in itertools.groupby(
            sorted([tool.__class__.__name__ for tool in tool_objs])):
        if len(list(group)) > 1:
            repeated_tools.append(typename)

    if repeated_tools:
        warnings.warn("%s are being repeated" % ",".join(repeated_tools))

    return tool_objs, tool_map

def _process_active_tools(toolbar, tool_map, active_drag, active_inspect, active_scroll, active_tap):
    """ Adds tools to the plot object

    Args:
        toolbar (Toolbar): instance of a Toolbar object
        tools_map (dict[str]|Tool): tool_map from _process_tools_arg
        active_drag (str or Tool): the tool to set active for drag
        active_inspect (str or Tool): the tool to set active for inspect
        active_scroll (str or Tool): the tool to set active for scroll
        active_tap (str or Tool): the tool to set active for tap

    Returns:
        None

    Note:
        This function sets properties on Toolbar
    """
    if active_drag in ['auto', None] or isinstance(active_drag, Tool):
        toolbar.active_drag = active_drag
    elif active_drag in tool_map:
        toolbar.active_drag = tool_map[active_drag]
    else:
        raise ValueError("Got unknown %r for 'active_drag', which was not a string supplied in 'tools' argument" % active_drag)

    if active_inspect in ['auto', None] or isinstance(active_inspect, Tool) or all([isinstance(t, Tool) for t in active_inspect]):
        toolbar.active_inspect = active_inspect
    elif active_inspect in tool_map:
        toolbar.active_inspect = tool_map[active_inspect]
    else:
        raise ValueError("Got unknown %r for 'active_inspect', which was not a string supplied in 'tools' argument" % active_scroll)

    if active_scroll in ['auto', None] or isinstance(active_scroll, Tool):
        toolbar.active_scroll = active_scroll
    elif active_scroll in tool_map:
        toolbar.active_scroll = tool_map[active_scroll]
    else:
        raise ValueError("Got unknown %r for 'active_scroll', which was not a string supplied in 'tools' argument" % active_scroll)

    if active_tap in ['auto', None] or isinstance(active_tap, Tool):
        toolbar.active_tap = active_tap
    elif active_tap in tool_map:
        toolbar.active_tap = tool_map[active_tap]
    else:
        raise ValueError("Got unknown %r for 'active_tap', which was not a string supplied in 'tools' argument" % active_tap)

def _get_argspecs(glyphclass):
    argspecs = OrderedDict()
    for arg in glyphclass._args:
        spec = {}
        descriptor = getattr(glyphclass, arg)

        # running python with -OO will discard docstrings -> __doc__ is None
        if descriptor.__doc__:
            spec['desc'] = "\n        ".join(textwrap.dedent(descriptor.__doc__).split("\n"))
        else:
            spec['desc'] = ""
        spec['default'] = descriptor.class_default(glyphclass)
        spec['type'] = descriptor.property._sphinx_type()
        argspecs[arg] = spec
    return argspecs

# This template generates the following:
#
# def foo(self, x, y=10, kwargs):
#     kwargs['x'] = x
#     kwargs['y'] = y
#     return func(self, **kwargs)
_sigfunc_template = """
def %s(self, %s, **kwargs):
%s
    return func(self, **kwargs)
"""

def _get_sigfunc(func_name, func, argspecs):
    # This code is to wrap the generic func(*args, **kw) glyph method so that
    # a much better signature is available to users. E.g., for ``square`` we have:
    #
    # Signature: p.square(x, y, size=4, angle=0.0, **kwargs)
    #
    # which provides descriptive names for positional args, as well as any defaults
    func_args_with_defaults = []
    for arg, spec in argspecs.items():
        if spec['default'] is None:
            func_args_with_defaults.append(arg)
        else:
            func_args_with_defaults.append("%s=%r" % (arg, spec['default']))
    args_text = ", ".join(func_args_with_defaults)
    kwargs_assign_text = "\n".join("    kwargs[%r] = %s" % (x, x) for x in argspecs)
    func_text = _sigfunc_template % (func_name, args_text, kwargs_assign_text)
    func_code = compile(func_text, "fakesource", "exec")
    func_globals = {}
    eval(func_code, {"func": func}, func_globals)
    return func_globals[func_name]

_arg_template = """    %s (%s) : %s
        (default: %r)
"""
_doc_template = """ Configure and add %s glyphs to this Figure.

Args:
%s

Keyword Args:
%s

Other Parameters:
    alpha (float) : an alias to set all alpha keyword args at once
    color (Color) : an alias to set all color keyword args at once
    source (ColumnDataSource) : a user supplied data source
    legend (str) : a legend tag for this glyph
    x_range_name (str) : name an extra range to use for mapping x-coordinates
    y_range_name (str) : name an extra range to use for mapping y-coordinates
    level (Enum) : control the render level order for this glyph

It is also possible to set the color and alpha parameters of a "nonselection"
glyph. To do so, prefix any visual parameter with ``'nonselection_'``.
For example, pass ``nonselection_alpha`` or ``nonselection_fill_alpha``.

Returns:
    GlyphRenderer
"""

def _add_sigfunc_info(func, argspecs, glyphclass, extra_docs):
    func.__name__ = glyphclass.__name__.lower()

    omissions = {'js_event_callbacks', 'js_property_callbacks', 'subscribed_events'}

    kwlines = []
    kws = glyphclass.properties() - set(argspecs)
    for kw in kws:
        # these are not really useful, and should also really be private, just skip them
        if kw in omissions: continue

        descriptor = getattr(glyphclass, kw)
        typ = descriptor.property._sphinx_type()
        if descriptor.__doc__:
            desc = "\n        ".join(textwrap.dedent(descriptor.__doc__).split("\n"))
        else:
            desc = ""
        kwlines.append(_arg_template % (kw, typ, desc, descriptor.class_default(glyphclass)))
    extra_kws = getattr(glyphclass, '_extra_kws', {})
    for kw, (typ, desc) in extra_kws.items():
        kwlines.append("    %s (%s) : %s" % (kw, typ, desc))
    kwlines.sort()

    arglines = []
    for arg, spec in argspecs.items():
        arglines.append(_arg_template % (arg, spec['type'], spec['desc'], spec['default']))

    func.__doc__ = _doc_template % (func.__name__, "\n".join(arglines), "\n".join(kwlines))
    if extra_docs:
        func.__doc__ += extra_docs

def _glyph_function(glyphclass, extra_docs=None):

    def func(self, **kwargs):

        # Process legend kwargs and remove legend before we get going
        legend_item_label = _get_legend_item_label(kwargs)

        # Need to check if user source is present before _pop_renderer_args
        is_user_source = kwargs.get('source', None) is not None
        renderer_kws = _pop_renderer_args(kwargs)
        source = renderer_kws['data_source']
        if not isinstance(source, ColumnarDataSource):
            try:
                # try converting the soruce to ColumnDataSource
                source = ColumnDataSource(source)
            except ValueError as err:
                msg = "Failed to auto-convert {curr_type} to ColumnDataSource.\n Original error: {err}".format(
                    curr_type=str(type(source)),
                    err=err.message
                )
                reraise(ValueError, ValueError(msg), sys.exc_info()[2])

            # update reddered_kws so that others can use the new source
            renderer_kws['data_source'] = source

        # handle the main glyph, need to process literals
        glyph_ca = _pop_colors_and_alpha(glyphclass, kwargs)
        _process_sequence_literals(glyphclass, kwargs, source, is_user_source)
        _process_sequence_literals(glyphclass, glyph_ca, source, is_user_source)

        # handle the nonselection glyph, we always set one
        nsglyph_ca = _pop_colors_and_alpha(glyphclass, kwargs, prefix='nonselection_', default_alpha=0.1)

        # handle the selection glyph, if any properties were given
        if any(x.startswith('selection_') for x in kwargs):
            sglyph_ca = _pop_colors_and_alpha(glyphclass, kwargs, prefix='selection_')
        else:
            sglyph_ca = None

        # handle the hover glyph, if any properties were given
        if any(x.startswith('hover_') for x in kwargs):
            hglyph_ca = _pop_colors_and_alpha(glyphclass, kwargs, prefix='hover_')
        else:
            hglyph_ca = None

        # handle the mute glyph, if any properties were given
        if any(x.startswith('muted_') for x in kwargs):
            mglyph_ca = _pop_colors_and_alpha(glyphclass, kwargs, prefix='muted_')
        else:
            mglyph_ca = None

        glyph = _make_glyph(glyphclass, kwargs, glyph_ca)
        nsglyph = _make_glyph(glyphclass, kwargs, nsglyph_ca)
        sglyph = _make_glyph(glyphclass, kwargs, sglyph_ca)
        hglyph = _make_glyph(glyphclass, kwargs, hglyph_ca)
        mglyph = _make_glyph(glyphclass, kwargs, mglyph_ca)

        glyph_renderer = GlyphRenderer(glyph=glyph,
                                       nonselection_glyph=nsglyph,
                                       selection_glyph=sglyph,
                                       hover_glyph=hglyph,
                                       muted_glyph=mglyph,
                                       **renderer_kws)

        if legend_item_label:
            _update_legend(self, legend_item_label, glyph_renderer)

        for tool in self.select(type=BoxSelectTool):
            tool.renderers.append(glyph_renderer)

        self.renderers.append(glyph_renderer)

        return glyph_renderer

    argspecs = _get_argspecs(glyphclass)

    sigfunc = _get_sigfunc(glyphclass.__name__.lower(), func, argspecs)

    sigfunc.glyph_method = True

    _add_sigfunc_info(sigfunc, argspecs, glyphclass, extra_docs)

    return sigfunc


import unittest
import numpy as np
import pandas as pd
from sklearn_helpers.preprocessing import \
    EnhancedLabelEncoder, MultiColumnLabelEncoder

class EnhancedLabelEncoderTest(unittest.TestCase):
    def test_accepts_only_1d(self):
        """It should only accept only a 1d array"""
        ehe = EnhancedLabelEncoder()
        train = np.array([
            [1,2],
            [2,1]
        ])

        self.assertRaises(ValueError, lambda: ehe.fit(train))

        # If it is flattened, it should not raise.
        train = train.flatten()
        ehe.fit(train)

    def test_handle_unknown_error(self):
        """If handle_unkown is 'error' it should throw on unseen labels"""
        ehe = EnhancedLabelEncoder(handle_unknown='error')
        train = np.array(['a', 'b', 'a'])
        test = np.array(['a','c'])

        ehe.fit(train)

        # Check that a ValueError is raised on transform
        self.assertRaises(ValueError, lambda: ehe.transform(test))

    def test_handle_unknown_ignore(self):
        """If handle_unknown is 'ignore' it should map unseen labels to a new value"""
        ehe = EnhancedLabelEncoder(handle_unknown='ignore')
        train = np.array(['a', 'b', 'a'])
        test = np.array(['a','c'])

        ehe.fit(train)

        # Check that the new label is mapped to the next value
        self.assertTrue(
            (np.array([0,2]) == ehe.transform(test)).all()
        )

class MultiColumnLabelEncoderTest(unittest.TestCase):
    def test_handle_ignore(self):
        """If handle_unknown is 'ignore' it should map unseen labels to a new value"""

        mce = MultiColumnLabelEncoder(handle_unknown='ignore')
        train = np.array([
            ['a', 'b'],
            ['c', 'a']
        ])
        test = np.array([
            ['a', 'd'],
            ['c', 'd']
        ])

        mce.fit(train)

        test_transformed = np.array([
            [0.,2.],
            [1.,2.]
        ])

        self.assertTrue(
            (mce.transform(test) == test_transformed).all()
        )

    def test_accepts_pandas(self):
        """It shouold accept a Pandas dataframe"""
        mce = MultiColumnLabelEncoder(handle_unknown='ignore')
        train = pd.DataFrame(
            np.array([
                ['a', 'b'],
                ['c', 'a']
            ]),
            columns=['col1', 'col2']
        )

        # This should not throw
        mce.fit_transform(train, np.array([1,2]))

    def test_classes(self):
        """It should return classes for each column"""

        def test_accepts_pandas(self):
            """It shouold accept a Pandas dataframe"""
        mce = MultiColumnLabelEncoder(
            handle_unknown='ignore'
        )
        train = pd.DataFrame(
            np.array([
                ['a', 'b'],
                ['c', 'a']
            ]),
            columns=['col1', 'col2']
        )

        mce.fit(train, np.array([1,2]))

        self.assertEqual(
            mce.classes_[0][0],
            'a'
        )

        self.assertEqual(
            mce.classes_[1][1],
            'b'
        )


from nose.tools import nottest, make_decorator
from functools import wraps

# tools without any external dependencies
from mtools.mlogfilter.mlogfilter import MLogFilterTool
from mtools.mlogvis.mlogvis import MLogVisTool
from mtools.mloginfo.mloginfo import MLogInfoTool

tools = [MLogFilterTool, MLogVisTool, MLogInfoTool]

# mlaunch depends on pymongo
try:
    from mtools.mlaunch.mlaunch import MLaunchTool
    tools.append(MLaunchTool)
except ImportError:
    pass

# mplotqueries depends on matplotlib
try:
    from mtools.mplotqueries.mplotqueries import MPlotQueriesTool
    tools.append(MPlotQueriesTool)
except ImportError:
    pass

def all_tools(fn):
    """ This is a decorator for test functions, that runs a loop over all command line tool
        classes imported above and passes each class to the test function.

        To use this decorator, the test function must accept a single parameter. Example:

        @all_tools
        def test_something(tool_cls):
            tool = tool_cls()
            # test tool here ...
    """
    @wraps(fn)     # copies __name__ of the original function, nose requires the name to start with "test_"
    def new_func():
        for tool in tools:
            fn(tool)
    return new_func

def test_import_all():
    """ Import all tools from mtools module.
        The tools that have external dependencies will only be imported if the dependencies are fulfilled.
        This test just passes by default because the imports are tested implicitly by loading this file.
    """
    pass


import itertools
import xml.etree.cElementTree as et

import networkx as nx
import pandas as pd
import numpy as np

def trackmate_peak_import(trackmate_xml_path, get_tracks=False):
    """Import detected peaks with TrackMate Fiji plugin.

    Parameters
    ----------
    trackmate_xml_path : str
        TrackMate XML file path.
    get_tracks : boolean
        Add tracks to label
    """

    root = et.fromstring(open(trackmate_xml_path).read())

    objects = []
    object_labels = {'FRAME': 't_stamp',
                     'POSITION_T': 't',
                     'POSITION_X': 'x',
                     'POSITION_Y': 'y',
                     'POSITION_Z': 'z',
                     'MEAN_INTENSITY': 'I',
                     'ESTIMATED_DIAMETER': 'w',
                     'QUALITY': 'q',
                     'ID': 'spot_id',
                     'MEAN_INTENSITY': 'mean_intensity',
                     'MEDIAN_INTENSITY': 'median_intensity',
                     'MIN_INTENSITY': 'min_intensity',
                     'MAX_INTENSITY': 'max_intensity',
                     'TOTAL_INTENSITY': 'total_intensity',
                     'STANDARD_DEVIATION': 'std_intensity',
                     'CONTRAST': 'contrast',
                     'SNR': 'snr'}

    features = root.find('Model').find('FeatureDeclarations').find('SpotFeatures')
    features = [c.get('feature') for c in features.getchildren()] + ['ID']

    spots = root.find('Model').find('AllSpots')
    trajs = pd.DataFrame([])
    objects = []
    for frame in spots.findall('SpotsInFrame'):
        for spot in frame.findall('Spot'):

            single_object = []
            for label in features:
                single_object.append(spot.get(label))

            objects.append(single_object)
    trajs = pd.DataFrame(objects, columns=features)
    trajs = trajs.astype(np.float)

    # Apply initial filtering
    initial_filter = root.find("Settings").find("InitialSpotFilter")

    trajs = filter_spots(trajs,
                         name=initial_filter.get('feature'),
                         value=float(initial_filter.get('value')),
                         isabove=True if initial_filter.get('isabove') == 'true' else False)

    # Apply filters
    spot_filters = root.find("Settings").find("SpotFilterCollection")

    for spot_filter in spot_filters.findall('Filter'):

        trajs = filter_spots(trajs,
                             name=spot_filter.get('feature'),
                             value=float(spot_filter.get('value')),
                             isabove=True if spot_filter.get('isabove') == 'true' else False)

    trajs = trajs.loc[:, object_labels.keys()]
    trajs.columns = [object_labels[k] for k in object_labels.keys()]
    trajs['label'] = np.arange(trajs.shape[0])

    # Get tracks
    if get_tracks:
        filtered_track_ids = [int(track.get('TRACK_ID')) for track in root.find('Model').find('FilteredTracks').findall('TrackID')]

        new_trajs = pd.DataFrame()
        label_id = 0
        trajs = trajs.set_index('spot_id')

        tracks = root.find('Model').find('AllTracks')
        for track in tracks.findall('Track'):

            track_id = int(track.get("TRACK_ID"))
            if track_id in filtered_track_ids:

                spot_ids = [(edge.get('SPOT_SOURCE_ID'), edge.get('SPOT_TARGET_ID'), edge.get('EDGE_TIME')) for edge in track.findall('Edge')]
                spot_ids = np.array(spot_ids).astype('float')
                spot_ids = pd.DataFrame(spot_ids, columns=['source', 'target', 'time'])
                spot_ids = spot_ids.sort_values(by='time')
                spot_ids = spot_ids.set_index('time')

                # Build graph
                graph = nx.Graph()
                for t, spot in spot_ids.iterrows():
                    graph.add_edge(int(spot['source']), int(spot['target']), attr_dict=dict(t=t))

                # Find graph extremities by checking if number of neighbors is equal to 1
                tracks_extremities = [node for node in graph.nodes() if len(graph.neighbors(node)) == 1]

                paths = []
                # Find all possible paths between extremities
                for source, target in itertools.combinations(tracks_extremities, 2):

                    # Find all path between two nodes
                    for path in nx.all_simple_paths(graph, source=source, target=target):

                        # Now we need to check wether this path respect the time logic contraint
                        # edges can only go in one direction of the time

                        # Build times vector according to path
                        t = []
                        for i, node_srce in enumerate(path[:-1]):
                            node_trgt = path[i+1]
                            t.append(graph.edge[node_srce][node_trgt]['t'])

                        # Will be equal to 1 if going to one time direction
                        if len(np.unique(np.sign(np.diff(t)))) == 1:
                            paths.append(path)

                # Add each individual trajectory to a new DataFrame called new_trajs
                for path in paths:
                    traj = trajs.loc[path].copy()
                    traj['label'] = label_id
                    label_id += 1

                    new_trajs = new_trajs.append(traj)

        trajs = new_trajs

    trajs.set_index(['t_stamp', 'label'], inplace=True)
    trajs = trajs.sort_index()

    return trajs

def filter_spots(spots, name, value, isabove):
    if isabove:
        spots = spots[spots[name] > value]
    else:
        spots = spots[spots[name] < value]

    return spots