"""Restricted Boltzmann Machine
"""

# Authors: Yann N. Dauphin <dauphiya@iro.umontreal.ca>
#          Vlad Niculae
#          Gabriel Synnaeve
#          Lars Buitinck
# License: BSD 3 clause

import time

import numpy as np
import scipy.sparse as sp

from ..base import BaseEstimator
from ..base import TransformerMixin
from ..externals.six.moves import xrange
from ..utils import check_array
from ..utils import check_random_state
from ..utils import gen_even_slices
from ..utils import issparse
from ..utils.extmath import safe_sparse_dot
from ..utils.extmath import log_logistic
from ..utils.fixes import expit             # logistic function
from ..utils.validation import check_is_fitted

class BernoulliRBM(BaseEstimator, TransformerMixin):
    """Bernoulli Restricted Boltzmann Machine (RBM).

    A Restricted Boltzmann Machine with binary visible units and
    binary hidden units. Parameters are estimated using Stochastic Maximum
    Likelihood (SML), also known as Persistent Contrastive Divergence (PCD)
    [2].

    The time complexity of this implementation is ``O(d ** 2)`` assuming
    d ~ n_features ~ n_components.

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

    Parameters
    ----------
    n_components : int, optional
        Number of binary hidden units.

    learning_rate : float, optional
        The learning rate for weight updates. It is *highly* recommended
        to tune this hyper-parameter. Reasonable values are in the
        10**[0., -3.] range.

    batch_size : int, optional
        Number of examples per minibatch.

    n_iter : int, optional
        Number of iterations/sweeps over the training dataset to perform
        during training.

    verbose : int, optional
        The verbosity level. The default, zero, means silent mode.

    random_state : integer or numpy.RandomState, optional
        A random number generator instance to define the state of the
        random permutations generator. If an integer is given, it fixes the
        seed. Defaults to the global numpy random number generator.

    Attributes
    ----------
    intercept_hidden_ : array-like, shape (n_components,)
        Biases of the hidden units.

    intercept_visible_ : array-like, shape (n_features,)
        Biases of the visible units.

    components_ : array-like, shape (n_components, n_features)
        Weight matrix, where n_features in the number of
        visible units and n_components is the number of hidden units.

    Examples
    --------

    >>> import numpy as np
    >>> from sklearn.neural_network import BernoulliRBM
    >>> X = np.array([[0, 0, 0], [0, 1, 1], [1, 0, 1], [1, 1, 1]])
    >>> model = BernoulliRBM(n_components=2)
    >>> model.fit(X)
    BernoulliRBM(batch_size=10, learning_rate=0.1, n_components=2, n_iter=10,
           random_state=None, verbose=0)

    References
    ----------

    [1] Hinton, G. E., Osindero, S. and Teh, Y. A fast learning algorithm for
        deep belief nets. Neural Computation 18, pp 1527-1554.
        http://www.cs.toronto.edu/~hinton/absps/fastnc.pdf

    [2] Tieleman, T. Training Restricted Boltzmann Machines using
        Approximations to the Likelihood Gradient. International Conference
        on Machine Learning (ICML) 2008
    """
    def __init__(self, n_components=256, learning_rate=0.1, batch_size=10,
                 n_iter=10, verbose=0, random_state=None):
        self.n_components = n_components
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.n_iter = n_iter
        self.verbose = verbose
        self.random_state = random_state

    def transform(self, X):
        """Compute the hidden layer activation probabilities, P(h=1|v=X).

        Parameters
        ----------
        X : {array-like, sparse matrix} shape (n_samples, n_features)
            The data to be transformed.

        Returns
        -------
        h : array, shape (n_samples, n_components)
            Latent representations of the data.
        """
        check_is_fitted(self, "components_")

        X = check_array(X, accept_sparse='csr', dtype=np.float64)
        return self._mean_hiddens(X)

    def _mean_hiddens(self, v):
        """Computes the probabilities P(h=1|v).

        Parameters
        ----------
        v : array-like, shape (n_samples, n_features)
            Values of the visible layer.

        Returns
        -------
        h : array-like, shape (n_samples, n_components)
            Corresponding mean field values for the hidden layer.
        """
        p = safe_sparse_dot(v, self.components_.T)
        p += self.intercept_hidden_
        return expit(p, out=p)

    def _sample_hiddens(self, v, rng):
        """Sample from the distribution P(h|v).

        Parameters
        ----------
        v : array-like, shape (n_samples, n_features)
            Values of the visible layer to sample from.

        rng : RandomState
            Random number generator to use.

        Returns
        -------
        h : array-like, shape (n_samples, n_components)
            Values of the hidden layer.
        """
        p = self._mean_hiddens(v)
        return (rng.random_sample(size=p.shape) < p)

    def _sample_visibles(self, h, rng):
        """Sample from the distribution P(v|h).

        Parameters
        ----------
        h : array-like, shape (n_samples, n_components)
            Values of the hidden layer to sample from.

        rng : RandomState
            Random number generator to use.

        Returns
        -------
        v : array-like, shape (n_samples, n_features)
            Values of the visible layer.
        """
        p = np.dot(h, self.components_)
        p += self.intercept_visible_
        expit(p, out=p)
        return (rng.random_sample(size=p.shape) < p)

    def _free_energy(self, v):
        """Computes the free energy F(v) = - log sum_h exp(-E(v,h)).

        Parameters
        ----------
        v : array-like, shape (n_samples, n_features)
            Values of the visible layer.

        Returns
        -------
        free_energy : array-like, shape (n_samples,)
            The value of the free energy.
        """
        return (- safe_sparse_dot(v, self.intercept_visible_)
                - np.logaddexp(0, safe_sparse_dot(v, self.components_.T)
                               + self.intercept_hidden_).sum(axis=1))

    def gibbs(self, v):
        """Perform one Gibbs sampling step.

        Parameters
        ----------
        v : array-like, shape (n_samples, n_features)
            Values of the visible layer to start from.

        Returns
        -------
        v_new : array-like, shape (n_samples, n_features)
            Values of the visible layer after one Gibbs step.
        """
        check_is_fitted(self, "components_")
        if not hasattr(self, "random_state_"):
            self.random_state_ = check_random_state(self.random_state)
        h_ = self._sample_hiddens(v, self.random_state_)
        v_ = self._sample_visibles(h_, self.random_state_)

        return v_

    def partial_fit(self, X, y=None):
        """Fit the model to the data X which should contain a partial
        segment of the data.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training data.

        Returns
        -------
        self : BernoulliRBM
            The fitted model.
        """
        X = check_array(X, accept_sparse='csr', dtype=np.float64)
        if not hasattr(self, 'random_state_'):
            self.random_state_ = check_random_state(self.random_state)
        if not hasattr(self, 'components_'):
            self.components_ = np.asarray(
                self.random_state_.normal(
                    0,
                    0.01,
                    (self.n_components, X.shape[1])
                ),
                order='F')
        if not hasattr(self, 'intercept_hidden_'):
            self.intercept_hidden_ = np.zeros(self.n_components, )
        if not hasattr(self, 'intercept_visible_'):
            self.intercept_visible_ = np.zeros(X.shape[1], )
        if not hasattr(self, 'h_samples_'):
            self.h_samples_ = np.zeros((self.batch_size, self.n_components))

        self._fit(X, self.random_state_)

    def _fit(self, v_pos, rng):
        """Inner fit for one mini-batch.

        Adjust the parameters to maximize the likelihood of v using
        Stochastic Maximum Likelihood (SML).

        Parameters
        ----------
        v_pos : array-like, shape (n_samples, n_features)
            The data to use for training.

        rng : RandomState
            Random number generator to use for sampling.
        """
        h_pos = self._mean_hiddens(v_pos)
        v_neg = self._sample_visibles(self.h_samples_, rng)
        h_neg = self._mean_hiddens(v_neg)

        lr = float(self.learning_rate) / v_pos.shape[0]
        update = safe_sparse_dot(v_pos.T, h_pos, dense_output=True).T
        update -= np.dot(h_neg.T, v_neg)
        self.components_ += lr * update
        self.intercept_hidden_ += lr * (h_pos.sum(axis=0) - h_neg.sum(axis=0))
        self.intercept_visible_ += lr * (np.asarray(
                                         v_pos.sum(axis=0)).squeeze() -
                                         v_neg.sum(axis=0))

        h_neg[rng.uniform(size=h_neg.shape) < h_neg] = 1.0  # sample binomial
        self.h_samples_ = np.floor(h_neg, h_neg)

    def score_samples(self, X):
        """Compute the pseudo-likelihood of X.

        Parameters
        ----------
        X : {array-like, sparse matrix} shape (n_samples, n_features)
            Values of the visible layer. Must be all-boolean (not checked).

        Returns
        -------
        pseudo_likelihood : array-like, shape (n_samples,)
            Value of the pseudo-likelihood (proxy for likelihood).

        Notes
        -----
        This method is not deterministic: it computes a quantity called the
        free energy on X, then on a randomly corrupted version of X, and
        returns the log of the logistic function of the difference.
        """
        check_is_fitted(self, "components_")

        v = check_array(X, accept_sparse='csr')
        rng = check_random_state(self.random_state)

        # Randomly corrupt one feature in each sample in v.
        ind = (np.arange(v.shape[0]),
               rng.randint(0, v.shape[1], v.shape[0]))
        if issparse(v):
            data = -2 * v[ind] + 1
            v_ = v + sp.csr_matrix((data.A.ravel(), ind), shape=v.shape)
        else:
            v_ = v.copy()
            v_[ind] = 1 - v_[ind]

        fe = self._free_energy(v)
        fe_ = self._free_energy(v_)
        return v.shape[1] * log_logistic(fe_ - fe)

    def fit(self, X, y=None):
        """Fit the model to the data X.

        Parameters
        ----------
        X : {array-like, sparse matrix} shape (n_samples, n_features)
            Training data.

        Returns
        -------
        self : BernoulliRBM
            The fitted model.
        """
        X = check_array(X, accept_sparse='csr', dtype=np.float64)
        n_samples = X.shape[0]
        rng = check_random_state(self.random_state)

        self.components_ = np.asarray(
            rng.normal(0, 0.01, (self.n_components, X.shape[1])),
            order='F')
        self.intercept_hidden_ = np.zeros(self.n_components, )
        self.intercept_visible_ = np.zeros(X.shape[1], )
        self.h_samples_ = np.zeros((self.batch_size, self.n_components))

        n_batches = int(np.ceil(float(n_samples) / self.batch_size))
        batch_slices = list(gen_even_slices(n_batches * self.batch_size,
                                            n_batches, n_samples))
        verbose = self.verbose
        begin = time.time()
        for iteration in xrange(1, self.n_iter + 1):
            for batch_slice in batch_slices:
                self._fit(X[batch_slice], rng)

            if verbose:
                end = time.time()
                print("[%s] Iteration %d, pseudo-likelihood = %.2f,"
                      " time = %.2fs"
                      % (type(self).__name__, iteration,
                         self.score_samples(X).mean(), end - begin))
                begin = end

        return self


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button, RadioButtons

fig, ax = plt.subplots()
plt.subplots_adjust(left=0.25, bottom=0.25)
t = np.arange(0.0, 1.0, 0.001)
a0 = 5
f0 = 3
s = a0*np.sin(2*np.pi*f0*t)
l, = plt.plot(t,s, lw=2, color='red')
plt.axis([0, 1, -10, 10])

axcolor = 'lightgoldenrodyellow'
axfreq = plt.axes([0.25, 0.1, 0.65, 0.03], axisbg=axcolor)
axamp  = plt.axes([0.25, 0.15, 0.65, 0.03], axisbg=axcolor)

sfreq = Slider(axfreq, 'Freq', 0.1, 30.0, valinit=f0)
samp = Slider(axamp, 'Amp', 0.1, 10.0, valinit=a0)

def update(val):
    amp = samp.val
    freq = sfreq.val
    l.set_ydata(amp*np.sin(2*np.pi*freq*t))
    fig.canvas.draw_idle()
sfreq.on_changed(update)
samp.on_changed(update)

resetax = plt.axes([0.8, 0.025, 0.1, 0.04])
button = Button(resetax, 'Reset', color=axcolor, hovercolor='0.975')
def reset(event):
    sfreq.reset()
    samp.reset()
button.on_clicked(reset)

rax = plt.axes([0.025, 0.5, 0.15, 0.15], axisbg=axcolor)
radio = RadioButtons(rax, ('red', 'blue', 'green'), active=0)
def colorfunc(label):
    l.set_color(label)
    fig.canvas.draw_idle()
radio.on_clicked(colorfunc)

plt.show()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
# heliopy documentation build configuration file, created by
# sphinx-quickstart
#
# This file is execfile()d with the current directory set to its
# containing dir.
#
# Note that not all possible configuration values are present in this
# autogenerated file.
#
# All configuration values have a default; values that are commented out
# serve to show the default.

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
"""
import os
import sys
import unittest.mock as mock

import matplotlib

import heliopy

matplotlib.use('agg')
sys.path.insert(0, os.path.abspath('../../'))

html_favicon = '../../artwork/favicon.ico'
html_sidebars = {'**': ['docsidebar.html']}

# Pretend these modules exits so readthedocs builds
MOCK_MODULES = []
for mod_name in MOCK_MODULES:
    sys.modules[mod_name] = mock.Mock()

# -- General configuration ------------------------------------------------

# If your documentation needs a minimal Sphinx version, state it here.
#
# needs_sphinx = '1.0'

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    'sphinx.ext.autodoc',
    'sphinx.ext.doctest',
    'sphinx.ext.coverage',
    'sphinx.ext.mathjax',
    'sphinx.ext.napoleon',
    'sphinx.ext.intersphinx',
    'sphinx_gallery.gen_gallery',
    'sphinx_automodapi.automodapi',
    'sphinx_issues'
]

intersphinx_mapping = {
    'matplotlib': ('https://matplotlib.org', None),
    'python': ('https://docs.python.org/3', None),
    'numpy': ('https://numpy.org/doc/stable', None),
    'scipy': ('https://docs.scipy.org/doc/scipy/reference', None),
    'pandas': ('https://pandas.pydata.org/pandas-docs/stable', None),
    'astropy': ('https://docs.astropy.org/en/stable', None),
    'sunpy': ('https://docs.sunpy.org/en/stable', None)}

sphinx_gallery_conf = {
    'default_thumb_file': os.path.abspath(os.path.join('..', '..', 'artwork', 'logo_circle.png')),
    'examples_dirs': '../../examples',
    'gallery_dirs': 'auto_examples',
    'backreferences_dir': 'gen_modules/backreferences',
    'doc_module': ('sphinx_gallery', 'heliopy'),
    'min_reported_time': 0,
    'abort_on_example_error': False,
}

issues_github_path = 'heliopython/heliopy'

# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']

# The suffix(es) of source filenames.
# You can specify multiple suffix as a list of string:
#
# source_suffix = ['.rst', '.md']
source_suffix = '.rst'

# The encoding of source files.
#
# source_encoding = 'utf-8-sig'

# The master toctree document.
master_doc = 'index'

# General information about the project.
project = 'HelioPy'
author = 'David Stansby'

# The version info for the project you're documenting, acts as replacement for
# |version| and |release|, also used in various other places throughout the
# built documents.
#
# The short X.Y version.
version = heliopy.__version__
# The full version, including alpha/beta/rc tags.
release = version

# The language for content autogenerated by Sphinx. Refer to documentation
# for a list of supported languages.
#
# This is also used if you do content translation via gettext catalogs.
# Usually you set "language" from the command line for these cases.
language = None

# There are two options for replacing |today|: either, you set today to some
# non-false value, then it is used:
#
# today = ''
#
# Else, today_fmt is used as the format for a strftime call.
#
# today_fmt = '%B %d, %Y'

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This patterns also effect to html_static_path and html_extra_path
exclude_patterns = []

# The reST default role (used for this markup: `text`) to use for all
# documents.
#
default_role = 'py:obj'

# If true, '()' will be appended to :func: etc. cross-reference text.
#
# add_function_parentheses = True

# If true, the current module name will be prepended to all description
# unit titles (such as .. function::).
#
# add_module_names = True

# If true, sectionauthor and moduleauthor directives will be shown in the
# output. They are ignored by default.
#
# show_authors = False

# The name of the Pygments (syntax highlighting) style to use.
pygments_style = 'sphinx'

# A list of ignored prefixes for module index sorting.
# modindex_common_prefix = []

# If true, keep warnings as "system message" paragraphs in the built documents.
# keep_warnings = False

# If true, `todo` and `todoList` produce output, else they produce nothing.
todo_include_todos = False

# -- Options for HTML output ----------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#
html_theme = 'heliopy_theme'
html_theme_path = ['../']

# Theme options are theme-specific and customize the look and feel of a theme
# further.  For a list of options available for each theme, see the
# documentation.
#
# html_theme_options = {'canonical_url': 'http://docs.heliopy.org/en/stable/',
#                      'analytics_id': 'UA-112461508-1',
#                      'prev_next_buttons_location': 'None'}

# Add any paths that contain custom themes here, relative to this directory.
# html_theme_path = []

# The name for this set of Sphinx documents.
# "<project> v<release> documentation" by default.
#
# html_title = 'heliopy v0.1'

# A shorter title for the navigation bar.  Default is the same as html_title.
#
# html_short_title = None

# The name of an image file (relative to this directory) to place at the top
# of the sidebar.
#
# html_logo = None

# The name of an image file (relative to this directory) to use as a favicon of
# the docs.  This file should be a Windows icon file (.ico) being 16x16 or
# 32x32 pixels large.
#
# html_favicon = None

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
# html_static_path = ['_static']

# Add any extra paths that contain custom files (such as robots.txt or
# .htaccess) here, relative to this directory. These files are copied
# directly to the root of the documentation.
#
# html_extra_path = []

# If not None, a 'Last updated on:' timestamp is inserted at every page
# bottom, using the given strftime format.
# The empty string is equivalent to '%b %d, %Y'.
#
# html_last_updated_fmt = None

# If true, SmartyPants will be used to convert quotes and dashes to
# typographically correct entities.
#
# html_use_smartypants = True

# Custom sidebar templates, maps document names to template names.
#
# html_sidebars = {}

# Additional templates that should be rendered to pages, maps page names to
# template names.
#
# html_additional_pages = {}

# If false, no module index is generated.
#
# html_domain_indices = True

# If false, no index is generated.
#
# html_use_index = True

# If true, the index is split into individual pages for each letter.
#
# html_split_index = False

# If true, links to the reST sources are added to the pages.
#
# html_show_sourcelink = True

# If true, "Created using Sphinx" is shown in the HTML footer. Default is True.
#
# html_show_sphinx = True

# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True.
#
html_show_copyright = False

# If true, an OpenSearch description file will be output, and all pages will
# contain a <link> tag referring to it.  The value of this option must be the
# base URL from which the finished HTML is served.
#
# html_use_opensearch = ''

# This is the file name suffix for HTML files (e.g. ".xhtml").
# html_file_suffix = None

# Language to be used for generating the HTML full-text search index.
# Sphinx supports the following languages:
#   'da', 'de', 'en', 'es', 'fi', 'fr', 'h', 'it', 'ja'
#   'nl', 'no', 'pt', 'ro', 'r', 'sv', 'tr', 'zh'
#
# html_search_language = 'en'

# A dictionary with options for the search language support, empty by default.
# 'ja' uses this config value.
# 'zh' user can custom change `jieba` dictionary path.
#
# html_search_options = {'type': 'default'}

# The name of a javascript file (relative to the configuration directory) that
# implements a search results scorer. If empty, the default will be used.
#
# html_search_scorer = 'scorer.js'

# Output file base name for HTML help builder.
htmlhelp_basename = 'heliopydoc'

# -- Options for LaTeX output ---------------------------------------------

latex_elements = {
    # The paper size ('letterpaper' or 'a4paper').
    #
    # 'papersize': 'letterpaper',

    # The font size ('10pt', '11pt' or '12pt').
    #
    # 'pointsize': '10pt',

    # Additional stuff for the LaTeX preamble.
    #
    # 'preamble': '',

    # Latex figure (float) alignment
    #
    # 'figure_align': 'htbp',
}

# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title,
#  author, documentclass [howto, manual, or own class]).
latex_documents = [
    (master_doc, 'heliopy.tex', 'HelioPy Documentation',
     'David Stansby', 'manual'),
]

# The name of an image file (relative to this directory) to place at the top of
# the title page.
#
# latex_logo = None

# For "manual" documents, if this is true, then toplevel headings are parts,
# not chapters.
#
# latex_use_parts = False

# If true, show page references after internal links.
#
# latex_show_pagerefs = False

# If true, show URL addresses after external links.
#
# latex_show_urls = False

# Documents to append as an appendix to all manuals.
#
# latex_appendices = []

# It false, will not define \strong, \code, 	itleref, \crossref ... but only
# \sphinxstrong, ..., \sphinxtitleref, ... To help avoid clash with user added
# packages.
#
# latex_keep_old_macro_names = True

# If false, no module index is generated.
#
# latex_domain_indices = True

# -- Options for manual page output ---------------------------------------

# One entry per manual page. List of tuples
# (source start file, name, description, authors, manual section).
man_pages = [
    (master_doc, 'HelioPy', 'HelioPy Documentation',
     [author], 1)
]

# If true, show URL addresses after external links.
#
# man_show_urls = False

# -- Options for Texinfo output -------------------------------------------

# Grouping the document tree into Texinfo files. List of tuples
# (source start file, target name, title, author,
#  dir menu entry, description, category)
'''texinfo_documents = [
    (master_doc, 'HelioPy', 'HelioPy Documentation',
     author, 'HelioPy team', 'Python for space physics.',
     'Miscellaneous'),
]

html_theme_options = {
    "about_links": [
        ("About", "http://docs.heliopy.org/en/stable/guide/about.html", 1),
        (
            "Acknowledge HelioPy",
            "http://docs.heliopy.org/en/stable/guide/citing.html",
            1,
        ),
        ("Code of Conduct", "http://docs.heliopy.org/en/stable/guide/code-of-conduct.html", 1),
    ],
    "navbar_links": [
        ("Documentation", "http://docs.heliopy.org/en/stable/index.html", 1),
        ("Get Help", "http://docs.heliopy.org/en/stable/index.html", 1),
    ],
}'''


import matplotlib as mpl
mpl.use('Agg')

import collections
import sys
import numpy
import matplotlib.pyplot as plt
import vcf
from matplotlib.backends.backend_pdf import PdfPages
from fractions import Fraction

class TRIO:
    def __init__(self, child, mother, father):
        self.child  = child
        self.mother = mother
        self.father = father
    def __str__(self):
        return "%s\t%s\t%s"%(self.child, self.mother, self.father)

class FATHER_SON_PAIR:
    def __init__(self, son, father):
        self.son    = son
        self.father = father
    def __str__(self):
        return "%s\t%s"%(self.son, self.father)

def read_1kg_pedigree_file(input_file, header=True):
    data = open(input_file, "r")
    if header:
        data.readline()

    trios, father_son_pairs = [], []
    for line in data:
        tokens = line.strip().split()
        if tokens[2] != "0" and tokens[3] != "0":
            child, dad, mom = tokens[1:4]
            trios.append(TRIO(child,  dad, mom))
        if tokens[2] != "0" and tokens[4] == "1":
            father_son_pairs.append(FATHER_SON_PAIR(tokens[1], tokens[2]))
    data.close()
    print("There are %d trios and %d father-son-pairs in the pedigree file"%(len(trios), len(father_son_pairs)))
    return trios, father_son_pairs

# Find the index for the highest bin which is less than
# or equal to the provided value
def find_index(bins, value):
    low  = 0
    high = len(bins)-1
    while high > low + 1:
        midval = bins[(low+high)/2]
        if value > midval:
            low  = (low+high)/2
        elif value < midval:
            high = (low+high)/2 - 1
        else:
            return (low+high)/2
    if value < bins[low]:
        exit("Unable to find index. Exiting...")
    if value >= bins[high]:
        return high
    else:
        return low

def is_discordant(a11, a12, a21, a22):
    if (a11 == a21 and a12 == a22) or (a11 == a22 and a12 == a21):
        return False
    else:
        return True

def is_mendelian(a11, a12, a21, a22, a31, a32):
    if (a31 == a11 or a31 == a12) and (a32 == a21 or a32 == a22):
        return True
    elif (a31 == a21 or a31 == a22) and (a32 == a11 or a32 == a12):
        return True
    else:
        return False

def draw_bp_histogram(discordant_counts, pdfpage):
    # Create histogram of father-son differences
    bp_diff_counts     = [collections.defaultdict(int) for _ in xrange(6)]
    repeat_diff_counts = [collections.defaultdict(int) for _ in xrange(6)]
    out_frame_count    = 0
    in_frame_count     = 0
    for key,val in discordant_counts.items():
        bp_diff_counts[key[2]-1][key[1]-key[0]] += val
        repeat_diff_counts[key[2]-1][Fraction(key[1]-key[0], key[2])] += val
    for xlabel,diff_counts,in_frame in zip(["bps", "repeats"],
                                           [bp_diff_counts, repeat_diff_counts],
                                           [lambda bp,period: bp%period == 0, lambda rep,period: int(rep)==float(rep) ]):
        fig = plt.figure()
        ax  = fig.add_subplot(111)
        diffs = sorted(list(set(reduce(lambda x,y:x+y, map(lambda z: z.keys(), diff_counts)))))
        colors = ['c', 'r', 'g', 'y', 'b', 'm']
        heights = numpy.zeros(len(diffs))
        for i in xrange(6):
            vals = [diff_counts[i][x] for x in diffs]
            if sum(vals) == 0:
                continue

            in_frame_trips  = filter(lambda x: in_frame(x[0], i+1), zip(diffs, vals, heights))
            out_frame_trips = filter(lambda x: not in_frame(x[0], i+1), zip(diffs, vals, heights))
            if len(in_frame_trips) != 0:
                x,y,h = zip(*in_frame_trips)
                in_frame_count += sum(y)
                ax.bar(x, y, bottom=h, align='center', color=colors[i], width=0.25, label=str(i+1))
            if len(out_frame_trips) != 0:
                x,y,h = zip(*out_frame_trips)
                out_frame_count += sum(y)
                ax.bar(x, y, bottom=h, align='center', color=colors[i], width=0.25, label=str(i+1), hatch='//')
            heights += vals
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')
        ax.set_xlabel(r"$father-son ("+xlabel+")$")
        ax.set_ylabel(r"$n_{calls}$")
        ax.legend()
        pdfpage.savefig(fig)
    print("IN FRAME=%d, OUT FRAME=%d"%(in_frame_count/2, out_frame_count/2))

class CHRY_STATS:
    def __init__(self, father_son_pairs, call_output):
        self.pairs        = father_son_pairs
        self.output_calls = open(call_output, "w")

    def initialize(self, vcf_reader):
        sample_indices    = dict(zip(vcf_reader.samples, range(len(vcf_reader.samples))))
        self.pair_indices = []
        for i in xrange(len(self.pairs)):
            if self.pairs[i].son not in sample_indices:
                exit("Unable to assess chrY inheritance because no data was found for " + self.pairs[i].son)
            if self.pairs[i].father not in sample_indices:
                exit("Unable to assess chrY inheritance because no data was found for " + self.pairs[i].father)
            self.pair_indices.append([sample_indices[self.pairs[i].father], sample_indices[self.pairs[i].son]])
        self.missing_data_skip_counts = numpy.zeros(len(self.pair_indices))
        self.het_gt_skip_counts       = numpy.zeros(len(self.pair_indices))

        self.num_concordant    = 0
        self.num_discordant    = 0
        self.pair_info         = {}
        self.discordant_counts = collections.defaultdict(int)
        self.call_count        = 0

    def process_record(self, record):
        motif_len = len(record.INFO['MOTIF'])
        for i in xrange(len(self.pair_indices)):
            if any(map(lambda x: record.samples[x]['GT'] is None, self.pair_indices[i])):
                self.missing_data_skip_counts[i] += 1
                continue

            self.call_count += 1
            father = record.samples[self.pair_indices[i][0]]
            son    = record.samples[self.pair_indices[i][1]]
            gb_1a, gb_1b = map(int, father['GB'].split("/"))
            gb_2a, gb_2b = map(int, son['GB'].split("/"))
            self.output_calls.write("%d\t%s\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%s\t%s\n"%(self.call_count, record.CHROM, record.POS, record.INFO['END'],
                                                                                gb_1a + gb_1b, gb_2a + gb_2b,
                                                                                gb_1a,  gb_1b, gb_2a,  gb_2b, father.sample, son.sample))

            if gb_1a != gb_1b or gb_2a != gb_2b:
                self.het_gt_skip_counts[i] += 1
                if gb_1a != gb_1b:
                    print("chrY\t%d\t%d\t%s\t%s\t%s"%(record.POS, record.INFO["END"], father.sample, str(gb_1a) + "|" + str(gb_1b), "HET"))
                if gb_2a != gb_2b:
                    print("chrY\t%d\t%d\t%s\t%s\t%s"%(record.POS, record.INFO["END"], father.sample, str(gb_2a) + "|" + str(gb_2b), "HET"))
                continue

            if gb_1a != gb_2a:
                self.num_discordant += 1
                self.discordant_counts[(gb_1a, gb_2a, motif_len)] +=1
                print("chrY\t%d\t%d\t%s\t%s\t%s"%(record.POS, record.INFO["END"],
                                                  father.sample + "," + son.sample,
                                                  str(gb_1a) + "," + str(gb_2b), "DISCORDANT"))
            else:
                self.num_concordant += 1
            if (gb_1a, gb_2a) not in self.pair_info:
                self.pair_info[(gb_1a, gb_2a)] = []
            self.pair_info[(gb_1a, gb_2a)].append((record.CHROM, record.POS, record.INFO['END'], father.sample+"-"+son.sample))

    def finish(self, pdfpage, output_prefix):
        print("WARNING: Skipped " + str(self.missing_data_skip_counts) + " comparisons due to missing data for one or more individuals")
        print("WARNING: Skipped " + str(self.het_gt_skip_counts)       + " comparisons due to heterozygous genotypes for one or more individuals")

        if self.num_discordant + self.num_concordant != 0:
            print("%d vs. %d = %f Percent"%(self.num_discordant, self.num_concordant, 100.0*self.num_discordant/(self.num_discordant+self.num_concordant)))
        else:
            print("WARNING: No chrY calls were applicable for comparison")

        # Create bubble plot using all data
        fig  = plt.figure()
        ax   = fig.add_subplot(111)
        x, y = zip(*self.pair_info.keys())
        s    = numpy.array(map(len, self.pair_info.values()))*10
        ax.scatter(x, y, s=s, alpha=0.7)
        ax.set_xlabel("Father's genotype (bp)")
        ax.set_ylabel("Son's genotype (bp)")
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')
        ax.plot(numpy.arange(min(x)-5, max(x)+5, 1.0), numpy.arange(min(y)-5, max(y)+5, 1.0), linestyle='--', color='k')
        pdfpage.savefig(fig)

        # Create histogram of father-son differences
        draw_bp_histogram(self.discordant_counts, pdfpage)

        viz_output = open(output_prefix+"_chrY.csv", "w")
        viz_output.write(",".join(["X","Y", "CHROMS", "STARTS", "STOPS", "SAMPLES"]) + "\n")
        for key,val in self.pair_info.items():
            chroms, positions, ends, samples = map(list, zip(*val))
            viz_output.write(",".join([str(key[0]), str(key[1]), "_".join(chroms), "_".join(map(str, positions)), "_".join(map(str, ends)), "_".join(map(str, samples))]) + "\n")
        viz_output.close()

        self.output_calls.close()

class MENDELIAN_STATS:
    def __init__(self, trios, coverage_bins, quality_bins, max_coverage, quality_thresholds):
        self.trios         = trios
        self.coverage_bins = coverage_bins
        self.quality_bins  = quality_bins
        self.max_coverage  = max_coverage
        self.qual_thresh   = quality_thresholds

    def initialize(self, vcf_reader):
        sample_indices    = dict(zip(vcf_reader.samples, range(len(vcf_reader.samples))))
        self.trio_indices = []
        for i in xrange(len(self.trios)):
            if self.trios[i].child not in sample_indices:
                exit("Unable to calculate Mendelian inheritance because no data was found for " + self.trios[i].child)
            if self.trios[i].father not in sample_indices:
                exit("Unable to calculate Mendelian inheritance because no data was found for " + self.trios[i].father)
            if self.trios[i].mother not in sample_indices:
                exit("Unable to calculate Mendelian inheritance because no data was found for " + self.trios[i].mother)

            # Father, Mother, Child
            self.trio_indices.append(map(lambda x: sample_indices[x], [self.trios[i].father, self.trios[i].mother, self.trios[i].child]))

        self.coverage_bins = numpy.concatenate(([-100000], self.coverage_bins))
        self.quality_bins  = numpy.concatenate(([-100000], self.quality_bins))

        # Quality/Coverage x Trios x Period x Thresholds
        self.all_loci_nstrs  = [numpy.zeros((len(self.trios), 5, len(self.coverage_bins))), numpy.zeros((len(self.trios), 5, len(self.quality_bins)))]
        self.all_loci_nmend  = [numpy.zeros((len(self.trios), 5, len(self.coverage_bins))), numpy.zeros((len(self.trios), 5, len(self.quality_bins)))]
        self.disc_loci_nstrs = [numpy.zeros((len(self.trios), 5, len(self.coverage_bins))), numpy.zeros((len(self.trios), 5, len(self.quality_bins)))]
        self.disc_loci_nmend = [numpy.zeros((len(self.trios), 5, len(self.coverage_bins))), numpy.zeros((len(self.trios), 5, len(self.quality_bins)))]

        self.missing_data_skip_counts = numpy.zeros(len(self.trios))
        self.coverage_skip_counts     = numpy.zeros(len(self.trios))

        # Trios x Period x Thresholds
        self.all_loci_nstrs_min_q  = numpy.zeros((len(self.trios), 5, len(self.coverage_bins)))
        self.all_loci_nmend_min_q  = numpy.zeros((len(self.trios), 5, len(self.coverage_bins)))
        self.disc_loci_nstrs_min_q = numpy.zeros((len(self.trios), 5, len(self.coverage_bins)))
        self.disc_loci_nmend_min_q = numpy.zeros((len(self.trios), 5, len(self.coverage_bins)))

    def process_record(self, record):
        for i in xrange(len(self.trios)):
            if any(map(lambda x: record.samples[x]['GT'] is None, self.trio_indices[i])):
                self.missing_data_skip_counts[i] += 1
                continue
            if 'X' in record.CHROM or 'x' in record.CHROM or 'Y' in record.CHROM or 'y' in record.CHROM:
                continue

            q1, q2, q3 = map(lambda x: record.samples[x]["Q"],  self.trio_indices[i])
            c1, c2, c3 = map(lambda x: record.samples[x]["DP"], self.trio_indices[i])
            a11, a21   = record.samples[self.trio_indices[i][0]]["GT"].split("/")
            a21, a22   = record.samples[self.trio_indices[i][1]]["GT"].split("/")
            a31, a32   = record.samples[self.trio_indices[i][2]]["GT"].split("/")
            discordant = is_discordant(a11, a12, a21, a22)
            mendelian  = is_mendelian(a11, a12, a21, a22, a31, a32)

            # Filter out loci with too high of coverage
            if max(c1, c2, c3) > self.max_coverage:
                self.coverage_skip_counts[i] += 1
                continue

            coverage  = min(c1, c2, c3)
            bin_idx   = find_index(self.coverage_bins, coverage)
            motif_len = len(record.INFO["MOTIF"])-2
            self.all_loci_nstrs [0][i][motif_len][bin_idx] += 1
            self.all_loci_nmend [0][i][motif_len][bin_idx] += mendelian*1
            self.disc_loci_nstrs[0][i][motif_len][bin_idx] += discordant*1
            self.disc_loci_nmend[0][i][motif_len][bin_idx] += discordant*mendelian*1

            quality   = min(q1, q2, q3)
            bin_idx   = find_index(self.quality_bins, quality)
            self.all_loci_nstrs [1][i][motif_len][bin_idx] += 1
            self.all_loci_nmend [1][i][motif_len][bin_idx] += mendelian*1
            self.disc_loci_nstrs[1][i][motif_len][bin_idx] += discordant*1
            self.disc_loci_nmend[1][i][motif_len][bin_idx] += discordant*mendelian*1

            coverage  = min(c1, c2, c3)
            bin_idx   = find_index(self.coverage_bins, coverage)
            if quality > self.qual_thresh[motif_len]:
                self.all_loci_nstrs_min_q  [i][motif_len][bin_idx] += 1
                self.all_loci_nmend_min_q  [i][motif_len][bin_idx] += mendelian*1
                self.disc_loci_nstrs_min_q [i][motif_len][bin_idx] += discordant*1
                self.disc_loci_nmend_min_q [i][motif_len][bin_idx] += discordant*mendelian*1

    def finish(self, pdfpage):
        print("WARNING: Skipped " + str(self.missing_data_skip_counts) + " loci due to missing data for one or more individual")
        print("WARNING: Skipped " + str(self.coverage_skip_counts)     + " loci due to too high coverage")

        # Iterate over coverage and quality stats
        types = ['Coverage', 'Quality', 'Coverage']
        bins  = [self.coverage_bins, self.quality_bins, self.coverage_bins]
        for n in xrange(3):
            # Sum across all trios
            if n == 0 or n == 1:
                all_loci_nstrs  = numpy.sum(self.all_loci_nstrs [n],  axis=0)
                all_loci_nmend  = numpy.sum(self.all_loci_nmend [n],  axis=0)
                disc_loci_nstrs = numpy.sum(self.disc_loci_nstrs[n],  axis=0)
                disc_loci_nmend = numpy.sum(self.disc_loci_nmend[n],  axis=0)
            else:
                all_loci_nstrs  = numpy.sum(self.all_loci_nstrs_min_q,  axis=0)
                all_loci_nmend  = numpy.sum(self.all_loci_nmend_min_q,  axis=0)
                disc_loci_nstrs = numpy.sum(self.disc_loci_nstrs_min_q, axis=0)
                disc_loci_nmend = numpy.sum(self.disc_loci_nmend_min_q, axis=0)

            # Create plots for individual periods
            fig  = plt.figure()
            ax1 = fig.add_subplot(221)
            ax1.set_ylabel("Fraction Mendelian")
            ax1.set_title("All sites")

            ax2 = fig.add_subplot(222, sharey=ax1)
            ax2.set_title("Discordant parental sites")

            ax3 = fig.add_subplot(223, sharex=ax1)
            ax3.set_xlabel(types[n] + " threshold")
            ax3.set_ylabel("# genotypes")
            ax3.set_yscale('log')

            ax4 = fig.add_subplot(224, sharex=ax2, sharey=ax3)
            ax4.set_xlabel(types[n] + " threshold")
            ax4.set_yscale('log')

            box1 = ax1.get_position()
            ax1.set_position([box1.x0, box1.y0, box1.width*0.9, box1.height])
            ax2.set_position([box1.x0 + box1.width*1.15, box1.y0, box1.width*0.9, box1.height])
            box3 = ax3.get_position()
            ax3.set_position([box3.x0, box3.y0, box3.width*0.9, box3.height])
            ax4.set_position([box3.x0 + box3.width*1.15, box3.y0, box3.width*0.9, box3.height])

            font_size = 9
            for i in xrange(5):
                nstrs_all  = numpy.cumsum(all_loci_nstrs [i][::-1])[::-1]
                nmend_all  = numpy.cumsum(all_loci_nmend [i][::-1])[::-1]
                nstrs_disc = numpy.cumsum(disc_loci_nstrs[i][::-1])[::-1]
                nmend_disc = numpy.cumsum(disc_loci_nmend[i][::-1])[::-1]
                all_fracs  = (1.0*nmend_all/nstrs_all)[1:]
                disc_fracs = (1.0*nmend_disc/nstrs_disc)[1:]
                ax1.plot(bins[n][1:], all_fracs,      '-o', label=str(i+1))
                ax2.plot(bins[n][1:], disc_fracs,     '-o', label=str(i+1))
                ax3.plot(bins[n][1:], nstrs_all[1:],  '-o', label=str(i+1))
                ax4.plot(bins[n][1:], nstrs_disc[1:], '-o', label=str(i+1))
            ax4.legend(bbox_to_anchor=(1.05, 0.9, 0.25, 0.2), loc='center left')
            for ax in [ax1, ax2, ax3, ax4]:
                for tick in ax.xaxis.get_major_ticks():
                    tick.label.set_fontsize(font_size)
                for tick in ax.yaxis.get_major_ticks():
                    tick.label.set_fontsize(font_size)
            pdfpage.savefig(fig)

            # Create plots using all periods

            # Sum across all periods
            all_loci_nstrs  = numpy.sum(all_loci_nstrs,  axis=0)
            all_loci_nmend  = numpy.sum(all_loci_nmend,  axis=0)
            disc_loci_nstrs = numpy.sum(disc_loci_nstrs, axis=0)
            disc_loci_nmend = numpy.sum(disc_loci_nmend, axis=0)

            # Transform into running sums
            all_loci_nstrs  = numpy.cumsum(all_loci_nstrs[::-1])[::-1]
            all_loci_nmend  = numpy.cumsum(all_loci_nmend[::-1])[::-1]
            disc_loci_nstrs = numpy.cumsum(disc_loci_nstrs[::-1])[::-1]
            disc_loci_nmend = numpy.cumsum(disc_loci_nmend[::-1])[::-1]

            # Calculate the fraction of Mendelian inheritance for all loci and discordant loci
            all_loci_fracs  = (1.0*all_loci_nmend/all_loci_nstrs)[1:]
            disc_loci_fracs = (1.0*disc_loci_nmend/disc_loci_nstrs)[1:]

            fig = plt.figure()
            ax1  = fig.add_subplot(221)
            ax1.set_ylabel("Fraction Mendelian")
            ax1.set_title("All sites")
            ax1.plot(bins[n][1:], all_loci_fracs, '-o')

            ax2  = fig.add_subplot(222, sharey=ax1)
            ax2.plot(bins[n][1:], disc_loci_fracs, '-o')
            ax2.set_title("Discordant parental sites")

            ax3  = fig.add_subplot(223, sharex=ax1)
            ax3.set_xlabel(types[n] + " threshold")
            ax3.set_ylabel("# genotypes")
            ax3.set_yscale('log')
            ax3.plot(bins[n][1:], all_loci_nstrs[1:], '-o')

            ax4  = fig.add_subplot(224, sharex=ax2, sharey=ax3)
            ax4.set_xlabel(types[n] + " threshold")
            ax4.set_yscale('log')
            ax4.plot(bins[n][1:], disc_loci_nstrs[1:], '-o')

            for ax in [ax1, ax2, ax3, ax4]:
                for tick in ax.xaxis.get_major_ticks():
                    tick.label.set_fontsize(font_size)
                for tick in ax.yaxis.get_major_ticks():
                    tick.label.set_fontsize(font_size)
            pdfpage.savefig(fig)

            mpl.rcParams['xtick.labelsize'] = 10
            mpl.rcParams['ytick.labelsize'] = 10
            fig = plt.figure()
            ax1 = fig.add_subplot(111)
            ax1.plot(bins[n][1:], all_loci_fracs, '-o', color='b')
            ax1.set_ylabel("Fraction Mendelian")
            ax1.set_xlabel(types[n] + " threshold")
            ax2 = ax1.twinx()
            ax2.set_yscale('log')
            ax2.plot(bins[n][1:], all_loci_nstrs[1:], '-o', color='g')
            pdfpage.savefig(fig)
            ax1.axis('equal')
            pdfpage.savefig(fig)

            fig = plt.figure()
            ax1 = fig.add_subplot(211)
            ax2 = fig.add_subplot(212)
            ax1.plot(bins[n][1:], all_loci_fracs, '-o', color='b')
            ax1.set_ylabel("Fraction Mendelian")
            ax1.xaxis.set_ticks_position('bottom')
            ax1.yaxis.set_ticks_position('left')
            ax2.set_xlabel(types[n] + " threshold")
            ax2.plot(bins[n][1:], all_loci_nstrs[1:], '-o', color='g')
            ax2.set_yscale('log')
            ax2.set_ylabel("# Called loci across trios")
            ax2.xaxis.set_ticks_position('bottom')
            ax2.yaxis.set_ticks_position('left')
            pdfpage.savefig(fig)

def main():
    print("Invocation syntax: python pedigree_analysis.py 1kg_pedigree_file.txt vcf_file.vcf output_file.pdf")
    trios, father_son_pairs = read_1kg_pedigree_file(sys.argv[1], header=True)
    vcf_reader = vcf.Reader(filename=sys.argv[2])
    call_stats = sys.argv[3]
    samples    = vcf_reader.samples
    trios_with_data = []
    pairs_with_data = []

    for trio in trios:
        if trio.child in samples and trio.mother in samples and trio.father in samples:
            trios_with_data.append(trio)
    print("There are %d trios with data"%len(trios_with_data))

    for pair in father_son_pairs:
        if pair.father in samples and pair.son in samples:
            pairs_with_data.append(pair)
    print("There are %d father-son pairs with data"%(len(pairs_with_data)))

    coverage_bins  = numpy.append(numpy.arange(1.001, 5.0011, 1.0), numpy.arange(6.001, 18.0011, 2.0))
    quality_bins   = numpy.arange(0.0, 1.0, 0.1)
    quality_thresh = [0.9, 0.5, 0.5, 0.5, 0.5, 0.5]
    max_coverage   = 100
    processors     = [CHRY_STATS(pairs_with_data, call_stats)]
    #mend_stats     = MENDELIAN_STATS(trios_with_data, coverage_bins, quality_bins, max_coverage, quality_thresh)
    for proc in processors:
        proc.initialize(vcf_reader)
    for record in vcf_reader:
        for proc in processors:
            proc.process_record(record)
    pp = PdfPages(sys.argv[3]+".pdf")
    for proc in processors:
        proc.finish(pp, sys.argv[3])
    pp.close()
    return 0

if __name__ == "__main__":
    main()


import pytest

import numpy as np
import matplotlib.pyplot as plt

from aneris.control.factory import InterfaceFactory
from dtocean_core.core import (AutoFileInput,
                               AutoFileOutput,
                               AutoPlot,
                               Core)
from dtocean_core.data import CoreMetaData
from dtocean_core.data.definitions import XGrid2D

def test_XGrid2D_available():

    new_core = Core()
    all_objs = new_core.control._store._structures

    assert "XGrid2D" in all_objs.keys()

def test_XGrid2D():

    raw = {"values": np.random.randn(2, 3),
           "coords": [['a', 'b'], [-2, 0, 2]]}

    meta = CoreMetaData({"identifier": "test",
                         "structure": "test",
                         "title": "test",
                         "labels": ['x', 'y'],
                         "units": [None, 'm', 'POWER!']})

    test = XGrid2D()
    a = test.get_data(raw, meta)
    b = test.get_value(a)

    assert b.values.shape == (2,3)
    assert b.units == 'POWER!'
    assert b.y.units == 'm'

def test_get_None():

    test = XGrid2D()
    result = test.get_value(None)

    assert result is None

@pytest.mark.parametrize("fext", [".nc"])
def test_XGrid2D_auto_file(tmpdir, fext):

    test_path = tmpdir.mkdir("sub").join("test{}".format(fext))
    test_path_str = str(test_path)

    raw = {"values": np.random.randn(2, 3),
           "coords": [['a', 'b'], [-2, 0, 2]]}

    meta = CoreMetaData({"identifier": "test",
                         "structure": "test",
                         "title": "test",
                         "labels": ['x', 'y'],
                         "units": [None, 'm', 'POWER!']})

    test = XGrid2D()

    fout_factory = InterfaceFactory(AutoFileOutput)
    FOutCls = fout_factory(meta, test)

    fout = FOutCls()
    fout._path = test_path_str
    fout.data.result = test.get_data(raw, meta)

    fout.connect()

    assert len(tmpdir.listdir()) == 1

    fin_factory = InterfaceFactory(AutoFileInput)
    FInCls = fin_factory(meta, test)

    fin = FInCls()
    fin.meta.result = meta
    fin._path = test_path_str

    fin.connect()
    result = test.get_data(fin.data.result, meta)

    assert result.values.shape == (2,3)
    assert result.units == 'POWER!'
    assert result.y.units == 'm'

def test_XGrid2D_auto_plot(tmpdir):

    raw = {"values": np.random.randn(2, 3),
           "coords": [['a', 'b'], [-2, 0, 2]]}

    meta = CoreMetaData({"identifier": "test",
                         "structure": "test",
                         "title": "test",
                         "labels": ['x', 'y'],
                         "units": ['\sum_{n=1}^{\infty} 2^{-n} = 1',
                                   'm',
                                   'POWER!']})

    test = XGrid2D()

    fout_factory = InterfaceFactory(AutoPlot)
    PlotCls = fout_factory(meta, test)

    plot = PlotCls()
    plot.data.result = test.get_data(raw, meta)
    plot.meta.result = meta

    plot.connect()

    assert len(plt.get_fignums()) == 1
    plt.close("all")

def test_XGrid2D_auto_plot_reverse(tmpdir):

    raw = {"values": np.random.randn(3, 2),
           "coords": [[-2, 0, 2], ['a', 'b']]}

    meta = CoreMetaData({"identifier": "test",
                         "structure": "test",
                         "title": "test",
                         "labels": ['x', 'y'],
                         "units": ['\sum_{n=1}^{\infty} 2^{-n} = 1',
                                   'm',
                                   'POWER!']})

    test = XGrid2D()

    fout_factory = InterfaceFactory(AutoPlot)
    PlotCls = fout_factory(meta, test)

    plot = PlotCls()
    plot.data.result = test.get_data(raw, meta)
    plot.meta.result = meta

    plot.connect()

    assert len(plt.get_fignums()) == 1
    plt.close("all")


#
# Copyright 2014 Quantopian, Inc.
#
# 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.
import datetime
import hashlib
import os

import numpy as np
import pandas as pd
import pytz
import xlrd
import requests

from six.moves import map

def col_letter_to_index(col_letter):
    # Only supports single letter,
    # but answer key doesn't need multi-letter, yet.
    index = 0
    for i, char in enumerate(reversed(col_letter)):
        index += ((ord(char) - 65) + 1) * pow(26, i)
    return index

DIR = os.path.dirname(os.path.realpath(__file__))

ANSWER_KEY_CHECKSUMS_PATH = os.path.join(DIR, 'risk-answer-key-checksums')
ANSWER_KEY_CHECKSUMS = open(ANSWER_KEY_CHECKSUMS_PATH, 'r').read().splitlines()

ANSWER_KEY_FILENAME = 'risk-answer-key.xlsx'

ANSWER_KEY_PATH = os.path.join(DIR, ANSWER_KEY_FILENAME)

ANSWER_KEY_BUCKET_NAME = 'zipline-test_data'

ANSWER_KEY_DL_TEMPLATE = """
https://s3.amazonaws.com/zipline-test-data/risk/{md5}/risk-answer-key.xlsx
""".strip()

LATEST_ANSWER_KEY_URL = ANSWER_KEY_DL_TEMPLATE.format(
    md5=ANSWER_KEY_CHECKSUMS[-1])

def answer_key_signature():
    with open(ANSWER_KEY_PATH, 'rb') as f:
        md5 = hashlib.md5()
        buf = f.read(1024)
        md5.update(buf)
        while buf != b"":
            buf = f.read(1024)
            md5.update(buf)
    return md5.hexdigest()

def ensure_latest_answer_key():
    """
    Get the latest answer key from a publically available location.

    Logic for determining what and when to download is as such:

    - If there is no local spreadsheet file, then get the lastest answer key,
    as defined by the last row in the checksum file.
    - If there is a local spreadsheet file:
    -- If the spreadsheet's checksum is in the checksum file:
    --- If the spreadsheet's checksum does not match the latest, then grab the
    the latest checksum and replace the local checksum file.
    --- If the spreadsheet's checksum matches the latest, then skip download,
    and use the local spreadsheet as a cached copy.
    -- If the spreadsheet's checksum is not in the checksum file, then leave
    the local file alone, assuming that the local xls's md5 is not in the list
    due to local modifications during development.

    It is possible that md5's could collide, if that is ever case, we should
    then find an alternative naming scheme.

    The spreadsheet answer sheet is not kept in SCM, as every edit would
    increase the repo size by the file size, since it is treated as a binary.
    """

    answer_key_dl_checksum = None

    local_answer_key_exists = os.path.exists(ANSWER_KEY_PATH)
    if local_answer_key_exists:
        local_hash = answer_key_signature()

        if local_hash in ANSWER_KEY_CHECKSUMS:
            # Assume previously downloaded version.
            # Check for latest.
            if local_hash != ANSWER_KEY_CHECKSUMS[-1]:
                # More recent checksum, download
                answer_key_dl_checksum = ANSWER_KEY_CHECKSUMS[-1]
            else:
                # Assume local copy that is being developed on
                answer_key_dl_checksum = None
    else:
        answer_key_dl_checksum = ANSWER_KEY_CHECKSUMS[-1]

    if answer_key_dl_checksum:
        res = requests.get(
            ANSWER_KEY_DL_TEMPLATE.format(md5=answer_key_dl_checksum))
        with open(ANSWER_KEY_PATH, 'wb') as f:
            f.write(res.content)

# Get latest answer key on load.
ensure_latest_answer_key()

class DataIndex(object):
    """
    Coordinates for the spreadsheet, using the values as seen in the notebook.
    The python-excel libraries use 0 index, while the spreadsheet in a GUI
    uses a 1 index.
    """
    def __init__(self, sheet_name, col, row_start, row_end,
                 value_type='float'):
        self.sheet_name = sheet_name
        self.col = col
        self.row_start = row_start
        self.row_end = row_end
        self.value_type = value_type

    @property
    def col_index(self):
        return col_letter_to_index(self.col) - 1

    @property
    def row_start_index(self):
        return self.row_start - 1

    @property
    def row_end_index(self):
        return self.row_end - 1

    def __str__(self):
        return "'{sheet_name}'!{col}{row_start}:{col}{row_end}".format(
            sheet_name=self.sheet_name,
            col=self.col,
            row_start=self.row_start,
            row_end=self.row_end
        )

class AnswerKey(object):

    INDEXES = {
        'RETURNS': DataIndex('Sim Period', 'D', 4, 255),

        'BENCHMARK': {
            'Dates': DataIndex('s_p', 'A', 4, 254, value_type='date'),
            'Returns': DataIndex('s_p', 'H', 4, 254)
        },

        # Below matches the inconsistent capitalization in spreadsheet
        'BENCHMARK_PERIOD_RETURNS': {
            'Monthly': DataIndex('s_p', 'R', 8, 19),
            '3-Month': DataIndex('s_p', 'S', 10, 19),
            '6-month': DataIndex('s_p', 'T', 13, 19),
            'year': DataIndex('s_p', 'U', 19, 19),
        },

        'BENCHMARK_PERIOD_VOLATILITY': {
            'Monthly': DataIndex('s_p', 'V', 8, 19),
            '3-Month': DataIndex('s_p', 'W', 10, 19),
            '6-month': DataIndex('s_p', 'X', 13, 19),
            'year': DataIndex('s_p', 'Y', 19, 19),
        },

        'ALGORITHM_PERIOD_RETURNS': {
            'Monthly': DataIndex('Sim Period', 'Z', 23, 34),
            '3-Month': DataIndex('Sim Period', 'AA', 25, 34),
            '6-month': DataIndex('Sim Period', 'AB', 28, 34),
            'year': DataIndex('Sim Period', 'AC', 34, 34),
        },

        'ALGORITHM_PERIOD_VOLATILITY': {
            'Monthly': DataIndex('Sim Period', 'AH', 23, 34),
            '3-Month': DataIndex('Sim Period', 'AI', 25, 34),
            '6-month': DataIndex('Sim Period', 'AJ', 28, 34),
            'year': DataIndex('Sim Period', 'AK', 34, 34),
        },

        'ALGORITHM_PERIOD_SHARPE': {
            'Monthly': DataIndex('Sim Period', 'AL', 23, 34),
            '3-Month': DataIndex('Sim Period', 'AM', 25, 34),
            '6-month': DataIndex('Sim Period', 'AN', 28, 34),
            'year': DataIndex('Sim Period', 'AO', 34, 34),
        },

        'ALGORITHM_PERIOD_BETA': {
            'Monthly': DataIndex('Sim Period', 'AP', 23, 34),
            '3-Month': DataIndex('Sim Period', 'AQ', 25, 34),
            '6-month': DataIndex('Sim Period', 'AR', 28, 34),
            'year': DataIndex('Sim Period', 'AS', 34, 34),
        },

        'ALGORITHM_PERIOD_ALPHA': {
            'Monthly': DataIndex('Sim Period', 'AT', 23, 34),
            '3-Month': DataIndex('Sim Period', 'AU', 25, 34),
            '6-month': DataIndex('Sim Period', 'AV', 28, 34),
            'year': DataIndex('Sim Period', 'AW', 34, 34),
        },

        'ALGORITHM_PERIOD_BENCHMARK_VARIANCE': {
            'Monthly': DataIndex('Sim Period', 'BJ', 23, 34),
            '3-Month': DataIndex('Sim Period', 'BK', 25, 34),
            '6-month': DataIndex('Sim Period', 'BL', 28, 34),
            'year': DataIndex('Sim Period', 'BM', 34, 34),
        },

        'ALGORITHM_PERIOD_COVARIANCE': {
            'Monthly': DataIndex('Sim Period', 'BF', 23, 34),
            '3-Month': DataIndex('Sim Period', 'BG', 25, 34),
            '6-month': DataIndex('Sim Period', 'BH', 28, 34),
            'year': DataIndex('Sim Period', 'BI', 34, 34),
        },

        'ALGORITHM_PERIOD_DOWNSIDE_RISK': {
            'Monthly': DataIndex('Sim Period', 'BN', 23, 34),
            '3-Month': DataIndex('Sim Period', 'BO', 25, 34),
            '6-month': DataIndex('Sim Period', 'BP', 28, 34),
            'year': DataIndex('Sim Period', 'BQ', 34, 34),
        },

        'ALGORITHM_PERIOD_SORTINO': {
            'Monthly': DataIndex('Sim Period', 'BR', 23, 34),
            '3-Month': DataIndex('Sim Period', 'BS', 25, 34),
            '6-month': DataIndex('Sim Period', 'BT', 28, 34),
            'year': DataIndex('Sim Period', 'BU', 34, 34),
        },

        'ALGORITHM_RETURN_VALUES': DataIndex(
            'Sim Cumulative', 'D', 4, 254),

        'ALGORITHM_CUMULATIVE_VOLATILITY': DataIndex(
            'Sim Cumulative', 'P', 4, 254),

        'ALGORITHM_CUMULATIVE_SHARPE': DataIndex(
            'Sim Cumulative', 'R', 4, 254),

        'CUMULATIVE_DOWNSIDE_RISK': DataIndex(
            'Sim Cumulative', 'U', 4, 254),

        'CUMULATIVE_SORTINO': DataIndex(
            'Sim Cumulative', 'V', 4, 254),

        'CUMULATIVE_INFORMATION': DataIndex(
            'Sim Cumulative', 'AA', 4, 254),

        'CUMULATIVE_BETA': DataIndex(
            'Sim Cumulative', 'AD', 4, 254),

        'CUMULATIVE_ALPHA': DataIndex(
            'Sim Cumulative', 'AE', 4, 254),

        'CUMULATIVE_MAX_DRAWDOWN': DataIndex(
            'Sim Cumulative', 'AH', 4, 254),

    }

    def __init__(self):
        self.workbook = xlrd.open_workbook(ANSWER_KEY_PATH)

        self.sheets = {}
        self.sheets['Sim Period'] = self.workbook.sheet_by_name('Sim Period')
        self.sheets['Sim Cumulative'] = self.workbook.sheet_by_name(
            'Sim Cumulative')
        self.sheets['s_p'] = self.workbook.sheet_by_name('s_p')

        for name, index in self.INDEXES.items():
            if isinstance(index, dict):
                subvalues = {}
                for subkey, subindex in index.items():
                    subvalues[subkey] = self.get_values(subindex)
                setattr(self, name, subvalues)
            else:
                setattr(self, name, self.get_values(index))

    def parse_date_value(self, value):
        return xlrd.xldate_as_tuple(value, 0)

    def parse_float_value(self, value):
        return value if value != '' else np.nan

    def get_raw_values(self, data_index):
        return self.sheets[data_index.sheet_name].col_values(
            data_index.col_index,
            data_index.row_start_index,
            data_index.row_end_index + 1)

    @property
    def value_type_to_value_func(self):
        return {
            'float': self.parse_float_value,
            'date': self.parse_date_value,
        }

    def get_values(self, data_index):
        value_parser = self.value_type_to_value_func[data_index.value_type]
        return [value for value in
                map(value_parser, self.get_raw_values(data_index))]

ANSWER_KEY = AnswerKey()

BENCHMARK_DATES = ANSWER_KEY.BENCHMARK['Dates']
BENCHMARK_RETURNS = ANSWER_KEY.BENCHMARK['Returns']
DATES = [datetime.datetime(*x, tzinfo=pytz.UTC) for x in BENCHMARK_DATES]
BENCHMARK = pd.Series(dict(zip(DATES, BENCHMARK_RETURNS)))
ALGORITHM_RETURNS = pd.Series(
    dict(zip(DATES, ANSWER_KEY.ALGORITHM_RETURN_VALUES)))
RETURNS_DATA = pd.DataFrame({'Benchmark Returns': BENCHMARK,
                             'Algorithm Returns': ALGORITHM_RETURNS})
RISK_CUMULATIVE = pd.DataFrame({
    'volatility': pd.Series(dict(zip(
        DATES, ANSWER_KEY.ALGORITHM_CUMULATIVE_VOLATILITY))),
    'sharpe': pd.Series(dict(zip(
        DATES, ANSWER_KEY.ALGORITHM_CUMULATIVE_SHARPE))),
    'downside_risk': pd.Series(dict(zip(
        DATES, ANSWER_KEY.CUMULATIVE_DOWNSIDE_RISK))),
    'sortino': pd.Series(dict(zip(
        DATES, ANSWER_KEY.CUMULATIVE_SORTINO))),
    'information': pd.Series(dict(zip(
        DATES, ANSWER_KEY.CUMULATIVE_INFORMATION))),
    'alpha': pd.Series(dict(zip(
        DATES, ANSWER_KEY.CUMULATIVE_ALPHA))),
    'beta': pd.Series(dict(zip(
        DATES, ANSWER_KEY.CUMULATIVE_BETA))),
    'max_drawdown': pd.Series(dict(zip(
        DATES, ANSWER_KEY.CUMULATIVE_MAX_DRAWDOWN))),
})


__author__ = "Christian Kongsgaard"
__license__ = 'MIT'

# -------------------------------------------------------------------------------------------------------------------- #
# IMPORTS

# Modules
import matplotlib.pyplot as plt
import numpy as np
import os
import datetime
import matplotlib.dates as mdates
import pandas as pd

# RiBuild Modules
from delphin_6_automation.file_parsing import delphin_parser

# -------------------------------------------------------------------------------------------------------------------- #
# RIBuild

# Application
colors = {'top': '#FBBA00', 'mid': '#B81A5D', 'bottom': '#79C6C0', '1d_brick': '#000000', '1d_mortar': '#BDCCD4'}
project_dict = {'dresden_zp_high_ratio_uninsulated_4a':
                    {'map':
                         {'5ad9e0352e2cb22f2c4f15b4': 'brick_1d',
                          '5ad9e3bf2e2cb22f2c4f166b': 'mortar_1d',
                          '5adb0a102e2cb22f2c4f17e9': '2d'}
                     },
                'dresden_zd_high_ratio_uninsulated_4a':
                    {'map':
                         {'5ad9e0ba2e2cb22f2c4f15f1': 'brick_1d',
                          '5ad9e3bf2e2cb22f2c4f166b': 'mortar_1d',
                          '5adb2dc02e2cb22f2c4f1873': '2d'}
                     },
                'potsdam_high_ratio_uninsulated_4a':
                    {'map':
                         {'5ad9e3462e2cb22f2c4f162e': 'brick_1d',
                          '5ad9e3bf2e2cb22f2c4f166b': 'mortar_1d',
                          '5adcc9702e2cb22f2c4f18fd': '2d'}
                     },
                'dresden_zp_low_ratio_uninsulated_4a':
                    {'map':
                         {'5ad9e6192e2cb22f2c4f175f': 'brick_1d',
                          '5ad9e5812e2cb22f2c4f1722': 'mortar_1d',
                          '5adda7172e2cb20baca57c6e': '2d'}
                     },
                'dresden_zd_low_ratio_uninsulated_4a':
                    {'map':
                         {'5ad9e44f2e2cb22f2c4f16a8': 'brick_1d',
                          '5ad9e5812e2cb22f2c4f1722': 'mortar_1d',
                          '5adcd4402e2cb22f2c4f1987': '2d'}
                     },
                'potsdam_low_ratio_uninsulated_4a':
                    {'map': {'5ad9e4f22e2cb22f2c4f16e5': 'brick_1d',
                             '5ad9e5812e2cb22f2c4f1722': 'mortar_1d',
                             '5add9b902e2cb20baca57be4': '2d'}
                     },
                'dresden_zp_high_ratio_insulated_4a':
                    {'map': {'5ae824252e2cb22d48db5955': 'brick_1d',
                             '5ae82c222e2cb2156000902b': 'mortar_1d',
                             '5ae355cf2e2cb2201055c1a4': '2d'}
                     },
                'dresden_zd_high_ratio_insulated_4a':
                    {'map': {'5ae824d82e2cb22d48db5998': 'brick_1d',
                             '5ae82c222e2cb2156000902b': 'mortar_1d',
                             '5ae398f12e2cb2201055c263': '2d'}
                     },
                'potsdam_high_ratio_insulated_4a':
                    {'map':
                         {'5ae82bac2e2cb21560008fe8': 'brick_1d',
                          '5ae82c222e2cb2156000902b': 'mortar_1d',
                          '5ae6ca982e2cb2201055c322': '2d'}
                     },
                'dresden_zp_low_ratio_insulated_4a':
                    {'map':
                         {'5ae82e5d2e2cb21560009137': 'brick_1d',
                          '5ae82dc02e2cb215600090f4': 'mortar_1d',
                          '5ae6fdbf2e2cb20d5891272f': '2d'}
                     },
                'dresden_zd_low_ratio_insulated_4a':
                    {'map':
                         {'5ae82cb12e2cb2156000906e': 'brick_1d',
                          '5ae82dc02e2cb215600090f4': 'mortar_1d',
                          '5ae6d9bf2e2cb2201055c3e1': '2d'}
                     },
                'potsdam_low_ratio_insulated_4a':
                    {'map':
                         {'5ae82d3b2e2cb215600090b1': 'brick_1d',
                          '5ae82dc02e2cb215600090f4': 'mortar_1d',
                          '5ae6edaf2e2cb20d58912670': '2d'}
                     },
                }

result_folder = r'U:\RIBuild\2D_1D\Results'

files = ['temperature profile.d6o']

# Functions
def get_points(result: dict, geo: dict):

    points = []
    for index_ in result['indices']:
        x_ = geo['element_geometry'][index_][1]
        y_ = geo['element_geometry'][index_][2]
        points.append({'cell': index_, 'x': x_, 'y': y_})

    return points

def add_data_to_points(points: list, results: dict, result_name: str):
    for cell_ in results['result'].keys():
        cell_index = int(cell_.split('_')[1])

        for point in points:
            if point['cell'] == cell_index:
                point[result_name] = np.array(results['result'][cell_][8760:])
                break

def main(project_):
    projects = list(project_dict[project_]['map'].keys())
    parsed_dicts = {'brick_1d': {'temp': {}, 'geo': {}},
                    'mortar_1d': {'temp': {}, 'geo': {}},
                    '2d': {'temp': {}, 'geo': {}}, }

    for p_ in projects:
        for mp_key in project_dict[project_]['map'].keys():
            if p_ == mp_key:
                key = project_dict[project_]['map'][mp_key]

        folder = result_folder + f'/{p_}/results'
        geo_file = [file
                    for file in os.listdir(folder)
                    if file.endswith('.g6a')][0]

        parsed_dicts[key]['temp'], _ = delphin_parser.d6o_to_dict(folder, files[0])
        parsed_dicts[key]['geo'] = delphin_parser.g6a_to_dict(folder, geo_file)

    x_date = [datetime.datetime(2020, 1, 1) + datetime.timedelta(hours=i)
              for i in range(len(parsed_dicts['brick_1d']['temp']['result']['cell_0'][8760:]))]

    # Brick 1D
    brick_1d = get_points(parsed_dicts['brick_1d']['temp'], parsed_dicts['brick_1d']['geo'])
    brick_1d.sort(key=lambda point: point['x'])
    add_data_to_points(brick_1d, parsed_dicts['brick_1d']['temp'], 'temperature')

    # Mortar 1D
    mortar_1d = get_points(parsed_dicts['mortar_1d']['temp'], parsed_dicts['mortar_1d']['geo'])
    mortar_1d.sort(key=lambda point: point['x'])
    add_data_to_points(mortar_1d, parsed_dicts['mortar_1d']['temp'], 'temperature')

    # 2D
    sim_2d = get_points(parsed_dicts['2d']['temp'], parsed_dicts['2d']['geo'])
    sim_2d.sort(key=lambda point: (point['x'], point['y']))
    add_data_to_points(sim_2d, parsed_dicts['2d']['temp'], 'temperature')

    # Plots
    def plot_locations(quantity):
        # Axes 00
        plt.figure()
        plt.title(f"{quantity}\n1D-Location: {brick_1d[0]['x']:.4f} and 2D-Location: {sim_2d[0]['x']:.4f}")
        plt.plot(x_date, brick_1d[0][quantity], color=colors['1d_brick'], label=f"1D Brick")
        plt.plot(x_date, mortar_1d[0][quantity], color=colors['1d_mortar'], label=f"1D Mortar")
        plt.plot(x_date, sim_2d[0][quantity], color=colors['bottom'], label=f"2D Bottom")
        plt.plot(x_date, sim_2d[1][quantity], color=colors['mid'], label=f"2D Mid")
        plt.plot(x_date, sim_2d[2][quantity], color=colors['top'], label=f"2D Top")
        plt.legend()
        plt.gcf().autofmt_xdate()
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
        plt.ylabel(f'{quantity}')

        # Axes 01
        plt.figure()
        plt.title(f"{quantity}\n1D-Location: {brick_1d[1]['x']:.4f} and 2D-Location: {sim_2d[3]['x']:.4f}")
        plt.plot(x_date, brick_1d[1][quantity], color=colors['1d_brick'], label=f"1D Brick")
        plt.plot(x_date, mortar_1d[1][quantity], color=colors['1d_mortar'], label=f"1D Mortar")
        plt.plot(x_date, sim_2d[3][quantity], color=colors['bottom'], label=f"2D Bottom")
        plt.plot(x_date, sim_2d[4][quantity], color=colors['mid'], label=f"2D Mid")
        plt.plot(x_date, sim_2d[5][quantity], color=colors['top'], label=f"2D Top")
        plt.legend()
        plt.gcf().autofmt_xdate()
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
        plt.ylabel(f'{quantity}')

        # Axes 10
        plt.figure()
        plt.title(f"{quantity}\n1D-Location: {brick_1d[2]['x']:.4f} and 2D-Location: {sim_2d[6]['x']:.4f}")
        plt.plot(x_date, brick_1d[2][quantity], color=colors['1d_brick'], label=f"1D Brick")
        plt.plot(x_date, mortar_1d[2][quantity], color=colors['1d_mortar'], label=f"1D Mortar")
        plt.plot(x_date, sim_2d[6][quantity], color=colors['bottom'], label=f"2D Bottom")
        plt.plot(x_date, sim_2d[7][quantity], color=colors['mid'], label=f"2D Mid")
        plt.plot(x_date, sim_2d[8][quantity], color=colors['top'], label=f"2D Top")
        plt.legend()
        plt.gcf().autofmt_xdate()
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
        plt.ylabel(f'{quantity}')

        # Axes 11
        plt.figure()
        plt.title(f"{quantity}\n1D-Location: {brick_1d[3]['x']:.4f} and 2D-Location: {sim_2d[9]['x']:.4f}")
        plt.plot(x_date, brick_1d[3][quantity], color=colors['1d_brick'], label=f"1D Brick")
        plt.plot(x_date, mortar_1d[3][quantity], color=colors['1d_mortar'], label=f"1D Mortar")
        plt.plot(x_date, sim_2d[9][quantity], color=colors['bottom'], label=f"2D Bottom")
        plt.plot(x_date, sim_2d[10][quantity], color=colors['mid'], label=f"2D Mid")
        plt.plot(x_date, sim_2d[11][quantity], color=colors['top'], label=f"2D Top")
        plt.legend()
        plt.gcf().autofmt_xdate()
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
        plt.ylabel(f'{quantity}')

        # Axes 20
        plt.figure()
        plt.title(f"{quantity}\n1D-Location: {brick_1d[4]['x']:.4f} and 2D-Location: {sim_2d[12]['x']:.4f}")
        plt.plot(x_date, brick_1d[4][quantity], color=colors['1d_brick'], label=f"1D Brick")
        plt.plot(x_date, mortar_1d[4][quantity], color=colors['1d_mortar'], label=f"1D Mortar")
        plt.plot(x_date, sim_2d[12][quantity], color=colors['bottom'], label=f"2D Bottom")
        plt.plot(x_date, sim_2d[13][quantity], color=colors['mid'], label=f"2D Mid")
        plt.plot(x_date, sim_2d[14][quantity], color=colors['top'], label=f"2D Top")
        plt.legend()
        plt.gcf().autofmt_xdate()
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
        plt.ylabel(f'{quantity}')

        # Axes 21
        plt.figure()
        plt.title(f"{quantity}\n1D-Location: {brick_1d[5]['x']:.4f} and 2D-Location: {sim_2d[15]['x']:.4f}")
        plt.plot(x_date, brick_1d[5][quantity], color=colors['1d_brick'], label=f"1D Brick")
        plt.plot(x_date, mortar_1d[5][quantity], color=colors['1d_mortar'], label=f"1D Mortar")
        plt.plot(x_date, sim_2d[15][quantity], color=colors['bottom'], label=f"2D Bottom")
        plt.plot(x_date, sim_2d[16][quantity], color=colors['mid'], label=f"2D Mid")
        plt.plot(x_date, sim_2d[17][quantity], color=colors['top'], label=f"2D Top")
        plt.legend()
        plt.gcf().autofmt_xdate()
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
        plt.ylabel(f'{quantity}')

    #plot_locations(quantity='temperature')
    #plt.show()

    def abs_diff(x1, x2):
        return x2 - x1

    def rel_diff(x1, x2):
        return (abs(x2 - x1))/abs(x2) * 100

    def differences(i, plots=False):

        avg_2d = np.mean([sim_2d[i]['temperature'],  sim_2d[i+2]['temperature'],  sim_2d[i+2]['temperature']], axis=0)
        brick_abs = abs_diff(brick_1d[i]['temperature'], avg_2d)
        mortar_abs = abs_diff(mortar_1d[i]['temperature'], avg_2d)
        brick_rel = rel_diff(brick_1d[i]['temperature'], avg_2d)
        mortar_rel = rel_diff(mortar_1d[i]['temperature'], avg_2d)

        if plots:
            # Plot
            plt.figure()
            plt.title(f"Temperature - Absolute Difference\n"
                      f"1D-Location: {brick_1d[i]['x']:.4f} and 2D-Location: {sim_2d[i*3]['x']:.4f}")
            plt.plot(x_date, brick_abs, color=colors['1d_brick'], label=f"1D Brick")
            plt.plot(x_date, mortar_abs, color=colors['1d_mortar'], label=f"1D Mortar")
            plt.legend()
            plt.gcf().autofmt_xdate()
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
            plt.ylabel('C')

            plt.figure()
            plt.title(f"Temperature - Relative Difference\n"
                      f"1D-Location: {brick_1d[i]['x']:.4f} and 2D-Location: {sim_2d[i*3]['x']:.4f}")
            plt.plot(x_date, brick_rel, color=colors['1d_brick'], label=f"1D Brick")
            plt.plot(x_date, mortar_rel, color=colors['1d_mortar'], label=f"1D Mortar")
            plt.legend()
            plt.gcf().autofmt_xdate()
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
            plt.ylabel('%')

        local_df = pd.DataFrame(columns=[f"{brick_1d[i]['x']:.04f}", f"{brick_1d[i]['x']:.04f}",
                                         f"{brick_1d[i]['x']:.04f}", f"{brick_1d[i]['x']:.04f}"],
                                index=pd.DatetimeIndex(start=datetime.datetime(2020, 1, 1),
                                                       freq='h', periods=len(brick_rel)),
                                data=np.vstack([brick_rel, brick_abs, mortar_rel, mortar_abs]).T)

        local_df.columns = pd.MultiIndex.from_arrays([local_df.columns, ['brick', 'brick', 'mortar', 'mortar'],
                                                     ['relative', 'absolute', 'relative', 'absolute']],
                                                     names=['location', 'material', 'type'])

        return local_df

    def differences_weighted(i, plots=False):

        avg_2d = np.average(a=[sim_2d[i]['temperature'],
                               sim_2d[i+2]['temperature'],
                               sim_2d[i+2]['temperature']],
                            axis=0,
                            weights=[56, 24., 56])
        brick_abs = abs_diff(brick_1d[i]['temperature'], avg_2d)
        mortar_abs = abs_diff(mortar_1d[i]['temperature'], avg_2d)
        brick_rel = rel_diff(brick_1d[i]['temperature'], avg_2d)
        mortar_rel = rel_diff(mortar_1d[i]['temperature'], avg_2d)

        if plots:
            # Plot
            plt.figure()
            plt.title(f"Temperature - Weighted Absolute Difference\n"
                      f"1D-Location: {brick_1d[i]['x']:.4f} and 2D-Location: {sim_2d[i*3]['x']:.4f}")
            plt.plot(x_date, brick_abs, color=colors['1d_brick'], label=f"1D Brick")
            plt.plot(x_date, mortar_abs, color=colors['1d_mortar'], label=f"1D Mortar")
            plt.legend()
            plt.gcf().autofmt_xdate()
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
            plt.ylabel('%')

            plt.figure()
            plt.title(f"Temperature - Weighted Relative Difference\n"
                      f"1D-Location: {brick_1d[i]['x']:.4f} and 2D-Location: {sim_2d[i*3]['x']:.4f}")
            plt.plot(x_date, brick_rel, color=colors['1d_brick'], label=f"1D Brick")
            plt.plot(x_date, mortar_rel, color=colors['1d_mortar'], label=f"1D Mortar")
            plt.legend()
            plt.gcf().autofmt_xdate()
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%B'))
            plt.ylabel('%')

        local_df = pd.DataFrame(columns=[f"{brick_1d[i]['x']:.04f}", f"{brick_1d[i]['x']:.04f}",
                                         f"{brick_1d[i]['x']:.04f}", f"{brick_1d[i]['x']:.04f}"],
                                index=pd.DatetimeIndex(start=datetime.datetime(2020, 1, 1),
                                                       freq='h', periods=len(brick_rel)),
                                data=np.vstack([brick_rel, brick_abs, mortar_rel, mortar_abs]).T)

        local_df.columns = pd.MultiIndex.from_arrays([local_df.columns, ['brick', 'brick', 'mortar', 'mortar'],
                                                     ['relative', 'absolute', 'relative', 'absolute']],
                                                     names=['location', 'material', 'type'])

        return local_df

    dataframes = []
    weighted_dataframes = []
    for index in range(len(brick_1d)):
        dataframes.append(differences(index))
        weighted_dataframes.append(differences_weighted(index))
        #plt.show()

    result_dataframe = pd.concat(dataframes, axis=1)
    w_result_dataframe = pd.concat(weighted_dataframes, axis=1)

    absolute_df = result_dataframe.loc[:, pd.IndexSlice[:, :, 'absolute']]
    absolute_df.columns = absolute_df.columns.droplevel(level=2)
    relative_df = result_dataframe.loc[:, pd.IndexSlice[:, :, 'relative']]
    relative_df.columns = relative_df.columns.droplevel(level=2)

    w_absolute_df = w_result_dataframe.loc[:, pd.IndexSlice[:, :, 'absolute']]
    w_absolute_df.columns = w_absolute_df.columns.droplevel(level=2)
    w_relative_df = w_result_dataframe.loc[:, pd.IndexSlice[:, :, 'relative']]
    w_relative_df.columns = w_relative_df.columns.droplevel(level=2)

    plt.figure()
    ax = absolute_df.boxplot()
    ax.set_ylim(-20, 20)
    ax.set_ylabel('Temperature - C')
    ax.set_title('Absolute Differences')
    #plt.show()

    out_folder = r'C:\Users\ocni\PycharmProjects\delphin_6_automation\data_process\2d_1d\processed_data'

    def excel():
        writer = pd.ExcelWriter(out_folder + '/temperature.xlsx')
        relative_df.describe().to_excel(writer, 'relative')
        absolute_df.describe().to_excel(writer, 'absolute')
        writer.save()

    #excel()

    def save_relative():
        hdf_file = out_folder + '/relative_temperature.h5'
        w_relative_df.to_hdf(hdf_file, project_, append=True)

    save_relative()

for project_key in project_dict.keys():
    print(f'Processing {project_key}')
    main(project_key)


"""
========================================
Label Propagation digits active learning
========================================

Demonstrates an active learning technique to learn handwritten digits
using label propagation.

We start by training a label propagation model with only 10 labeled points,
then we select the top five most uncertain points to label. Next, we train
with 15 labeled points (original 10 + 5 new ones). We repeat this process
four times to have a model trained with 30 labeled examples. Note you can
increase this to label more than 30 by changing `max_iterations`. Labeling
more than 30 can be useful to get a sense for the speed of convergence of
this active learning technique.

A plot will appear showing the top 5 most uncertain digits for each iteration
of training. These may or may not contain mistakes, but we will train the next
model with their true labels.
"""
print(__doc__)

# Authors: Clay Woolam <clay@woolam.org>
# License: BSD

import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

from sklearn import datasets
from sklearn.semi_supervised import label_propagation
from sklearn.metrics import classification_report, confusion_matrix

digits = datasets.load_digits()
rng = np.random.RandomState(0)
indices = np.arange(len(digits.data))
rng.shuffle(indices)

X = digits.data[indices[:330]]
y = digits.target[indices[:330]]
images = digits.images[indices[:330]]

n_total_samples = len(y)
n_labeled_points = 10
max_iterations = 5

unlabeled_indices = np.arange(n_total_samples)[n_labeled_points:]
f = plt.figure()

for i in range(max_iterations):
    if len(unlabeled_indices) == 0:
        print("No unlabeled items left to label.")
        break
    y_train = np.copy(y)
    y_train[unlabeled_indices] = -1

    lp_model = label_propagation.LabelSpreading(gamma=0.25, max_iter=5)
    lp_model.fit(X, y_train)

    predicted_labels = lp_model.transduction_[unlabeled_indices]
    true_labels = y[unlabeled_indices]

    cm = confusion_matrix(true_labels, predicted_labels,
                          labels=lp_model.classes_)

    print("Iteration %i %s" % (i, 70 * "_"))
    print("Label Spreading model: %d labeled & %d unlabeled (%d total)"
          % (n_labeled_points, n_total_samples - n_labeled_points,
             n_total_samples))

    print(classification_report(true_labels, predicted_labels))

    print("Confusion matrix")
    print(cm)

    # compute the entropies of transduced label distributions
    pred_entropies = stats.distributions.entropy(
        lp_model.label_distributions_.T)

    # select up to 5 digit examples that the classifier is most uncertain about
    uncertainty_index = np.argsort(pred_entropies)[::-1]
    uncertainty_index = uncertainty_index[
        np.in1d(uncertainty_index, unlabeled_indices)][:5]

    # keep track of indices that we get labels for
    delete_indices = np.array([])

    # for more than 5 iterations, visualize the gain only on the first 5
    if i < 5:
        f.text(.05, (1 - (i + 1) * .183),
               "model %d\n\nfit with\n%d labels" %
               ((i + 1), i * 5 + 10), size=10)
    for index, image_index in enumerate(uncertainty_index):
        image = images[image_index]

        # for more than 5 iterations, visualize the gain only on the first 5
        if i < 5:
            sub = f.add_subplot(5, 5, index + 1 + (5 * i))
            sub.imshow(image, cmap=plt.cm.gray_r, interpolation='none')
            sub.set_title("predict: %i\ntrue: %i" % (
                lp_model.transduction_[image_index], y[image_index]), size=10)
            sub.axis('off')

        # labeling 5 points, remote from labeled set
        delete_index, = np.where(unlabeled_indices == image_index)
        delete_indices = np.concatenate((delete_indices, delete_index))

    unlabeled_indices = np.delete(unlabeled_indices, delete_indices)
    n_labeled_points += len(uncertainty_index)

f.suptitle("Active learning with Label Propagation.\nRows show 5 most "
           "uncertain labels to learn with the next model.", y=1.15)
plt.subplots_adjust(left=0.2, bottom=0.03, right=0.9, top=0.9, wspace=0.2,
                    hspace=0.85)
plt.show()