#!/usr/bin/env python

import os
import click
from matplotlib import use
use('Agg')  # noqa
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from skbio import read, DistanceMatrix
from skbio.stats import isubsample
from skbio.stats.ordination import OrdinationResults
from collections import defaultdict
from collections import OrderedDict

ALPHA = 1.0
LINE_WIDTH = 0.3
LINE_WIDTH_WHITE = 2.0
LINE_WIDTH_BLACK = 1.0

@click.group()
def mod2_pcoa():
    pass

@mod2_pcoa.command()
@click.option('--coords', required=True, type=click.Path(
              resolve_path=True, readable=True, exists=True),
              help='Coordinates file')
@click.option('--mapping_file', required=True, type=click.Path(
              resolve_path=True, readable=True, exists=True),
              help='Mapping file')
@click.option('--output', required=True, type=click.Path(exists=True,
              writable=True, resolve_path=True), help='Output directory')
@click.option('--prefix', required=True, type=str, help='Output file prefix')
@click.option('--samples', required=False, type=str,
              help='Comma separated list of samples to print')
def body_site(coords, mapping_file, output, prefix, samples):
    """Generates as many figures as samples in the coordinates file"""
    o = read(coords, into=OrdinationResults)

    # coordinates
    c_df = pd.DataFrame(o.site, o.site_ids)

    # mapping file
    mf = pd.read_csv(mapping_file, '\t', converters=defaultdict(str),
                     index_col='#SampleID')
    mf = mf.loc[o.site_ids]

    if samples is None:
        samples = mf.index
    else:
        samples = set(samples.split(',')).intersection(set(o.site_ids))
        samples = mf.loc[samples].index

    color_hmp_fecal = sns.color_palette('Paired', 12)[10]  # light brown
    color_agp_fecal = sns.color_palette('Paired', 12)[11]  # dark brown
    color_hmp_oral = sns.color_palette('Paired', 12)[0]    # light blue
    color_agp_oral = sns.color_palette('Paired', 12)[1]    # dark blue
    color_hmp_skin = sns.color_palette('Paired', 12)[2]    # light green
    color_agp_skin = sns.color_palette('Paired', 12)[3]    # dark green

    grp_colors = {'AGP-FECAL': color_agp_fecal,
                  'AGP-ORAL':  color_agp_oral,
                  'AGP-SKIN':  color_agp_skin,
                  'HMP-FECAL': color_hmp_fecal,
                  'GG-FECAL':  color_hmp_fecal,
                  'PGP-FECAL': color_hmp_fecal,
                  'HMP-ORAL':  color_hmp_oral,
                  'PGP-ORAL':  color_hmp_oral,
                  'HMP-SKIN':  color_hmp_skin,
                  'PGP-SKIN':  color_hmp_skin}

    for sample in samples:

        # plot categories as 50 slices with random zorder
        for grp, color in grp_colors.iteritems():
            sub_coords = c_df[mf.TITLE_BODY_SITE == grp].values
            for i in np.array_split(sub_coords, 50):
                plt.scatter(i[:, 0], i[:, 1], color=color,
                            edgecolor=np.asarray(color)*0.6, lw=LINE_WIDTH,
                            alpha=ALPHA, zorder=np.random.rand())

        # plot participant's dot
        plt.scatter(c_df.loc[sample][0], c_df.loc[sample][1],
                    color=grp_colors[mf.loc[sample]['TITLE_BODY_SITE']],
                    s=270, edgecolor='w', zorder=1, lw=LINE_WIDTH_WHITE)
        plt.scatter(c_df.loc[sample][0], c_df.loc[sample][1],
                    color=grp_colors[mf.loc[sample]['TITLE_BODY_SITE']],
                    s=250, edgecolor=np.asarray(
                    grp_colors[mf.loc[sample]['TITLE_BODY_SITE']])*0.6,
                    zorder=2, lw=LINE_WIDTH_BLACK)

        plt.axis('off')
        my_dpi = 72
        figsize = (1000 / my_dpi, 1000 / my_dpi)
        out_file = os.path.join(output, '.'.join([prefix, sample, 'pdf']))
        plt.savefig(out_file, figsize=figsize, dpi=my_dpi)
        plt.close()

@mod2_pcoa.command()
@click.option('--distmat', required=True, type=click.Path(resolve_path=True,
                                                          readable=True,
                                                          exists=True),
              help='Input distance matrix to subsample nearest sample')
@click.option('--mapping_file', required=True, type=click.Path(
              resolve_path=True, readable=True, exists=True),
              help='Mapping file')
@click.option('--max', required=True, type=int,
              help='Max number of samples per category value')
@click.option('--category', required=True, type=str,
              help='The category to subsample in (likely COUNTRY)')
@click.option('--output', required=True, type=click.Path(exists=False,
              writable=True, resolve_path=True), help='Output file')
def subsample_dm(distmat, mapping_file, max, category, output):
    """Subsample the distmat to max samples per category value"""
    mf = pd.read_csv(mapping_file, '\t', converters=defaultdict(str),
                     index_col='#SampleID')
    id_to_cat = dict(mf[category])

    def bin_f(x):
        return id_to_cat[x]

    dm = read(distmat, into=DistanceMatrix)
    dm = dm.filter([id for _, id in isubsample(dm.ids, max, bin_f=bin_f)])
    dm.to_file(output)

@mod2_pcoa.command()
@click.option('--coords', required=True, type=click.Path(resolve_path=True,
              readable=True, exists=True), help='Coordinates file')
@click.option('--mapping_file', required=True, type=click.Path(
              resolve_path=True, readable=True, exists=True),
              help='Mapping file')
@click.option('--output', required=True, type=click.Path(exists=True,
              writable=True, resolve_path=True), help='Output directory')
@click.option('--prefix', required=True, type=str, help='Output file prefix')
@click.option('--samples', required=False, type=str,
              help='Comma separated list of samples to print')
@click.option('--distmat', required=True, type=click.Path(resolve_path=True,
                                                          readable=True,
                                                          exists=True),
              help=('Input distance matrix to find nearest sample (if not '
                    'present in the coordinates'))
def country(coords, mapping_file, output, prefix, samples, distmat):
    """Generates as many figures as samples in the coordinates file"""
    o = read(coords, into=OrdinationResults)
    o_id_lookup = set(o.site_ids)

    dm = read(distmat, into=DistanceMatrix)
    dm_id_lookup = {i: idx for idx, i in enumerate(dm.ids)}
    coord_samples_in_dm = {idx for idx, i in enumerate(dm.ids)
                           if i in o_id_lookup}

    # we'll be computing min values, so we need to avoid catching the  diagonal
    np.fill_diagonal(dm._data, np.inf)

    x, y = o.site[:, 0], o.site[:, 1]

    # coordinates
    c_df = pd.DataFrame(o.site, o.site_ids)

    # mapping file
    mf = pd.read_csv(mapping_file, '\t', converters=defaultdict(str),
                     index_col='#SampleID')
    # mf = mf.loc[o.site_ids]

    if samples is None:
        samples = dm.ids[:]
    else:
        samples = set(samples.split(',')).intersection(set(dm.ids))
        samples = mf.loc[samples].index

    color_Venezuela = sns.color_palette('Paired', 12)[10]
    color_Malawi = sns.color_palette('Paired', 12)[1]
    color_Western = sns.color_palette('Paired', 12)[4]
    color_Highlight = sns.color_palette('Paired', 12)[5]
    color_no_data = (0.5, 0.5, 0.5)

    grp_colors = OrderedDict()
    grp_colors['no_data'] = color_no_data
    grp_colors['Australia'] = color_Western
    grp_colors['Belgium'] = color_Western
    grp_colors['Canada'] = color_Western
    grp_colors['China'] = color_Western
    grp_colors['Finland'] = color_Western
    grp_colors['France'] = color_Western
    grp_colors['Germany'] = color_Western
    grp_colors['Great Britain'] = color_Western
    grp_colors['Ireland'] = color_Western
    grp_colors['Japan'] = color_Western
    grp_colors['Netherlands'] = color_Western
    grp_colors['New Zealand'] = color_Western
    grp_colors['Norway'] = color_Western
    grp_colors['Scotland'] = color_Western
    grp_colors['Spain'] = color_Western
    grp_colors['Switzerland'] = color_Western
    grp_colors['Thailand'] = color_Western
    grp_colors['United Arab Emirates'] = color_Western
    grp_colors['United Kingdom'] = color_Western
    grp_colors['United States of America'] = color_Western
    grp_colors['Malawi'] = color_Malawi
    grp_colors['Venezuela'] = color_Venezuela

    for sample_to_plot in samples:
        if sample_to_plot in o_id_lookup:
            sample = sample_to_plot
        else:
            # find the closest sample in the distance matrix that is in the
            # coordinates data
            sample = None
            for i in dm[dm_id_lookup[sample_to_plot]].argsort():
                if i in coord_samples_in_dm:
                    sample = dm.ids[i]
                    break

            # this should not ever happen
            if sample is None:
                raise ValueError("Unable to find a similar sample?")

        # countour plot superimposed
        sns.kdeplot(x, y, cmap='bone')
        sns.set_context(rc={"lines.linewidth": 0.75})

        # change particapant's country's color to color_Highlight unless
        # country is Venezuela or Malawi
        if (mf.loc[sample_to_plot]['COUNTRY'] != 'Malawi') & (
                mf.loc[sample_to_plot]['COUNTRY'] != 'Venezuela'):
            grp_colors[mf.loc[sample_to_plot]['COUNTRY']] = color_Highlight

        # plot each country except participant's according to colors above
        for grp, color in grp_colors.iteritems():
            if grp == mf.loc[sample_to_plot]['COUNTRY']:
                continue
            sub_coords = c_df[mf.COUNTRY == grp]
            plt.scatter(sub_coords[0], sub_coords[1], color=color,
                        edgecolor=np.asarray(color)*0.6, lw=LINE_WIDTH,
                        alpha=ALPHA)

        # now plot participant's country
        grp = mf.loc[sample_to_plot]['COUNTRY']
        color = grp_colors[grp]
        sub_coords = c_df[mf.COUNTRY == grp]
        plt.scatter(sub_coords[0], sub_coords[1], color=color,
                    edgecolor=np.asarray(color)*0.6, lw=LINE_WIDTH,
                    alpha=ALPHA)

        # plot participant's dot
        plt.scatter(c_df.loc[sample][0], c_df.loc[sample][1],
                    color=grp_colors[mf.loc[sample_to_plot]['COUNTRY']],
                    s=270, edgecolor='w', zorder=1, lw=LINE_WIDTH_WHITE)
        plt.scatter(c_df.loc[sample][0], c_df.loc[sample][1],
                    color=grp_colors[mf.loc[sample_to_plot]['COUNTRY']],
                    s=250, edgecolor=np.asarray(grp_colors[
                        mf.loc[sample_to_plot]['COUNTRY']])*0.6,
                    zorder=2, lw=LINE_WIDTH_BLACK)

        # reset particapant's country's color to color_Western unless country
        # is Venezuela or Malawi
        if (mf.loc[sample_to_plot]['COUNTRY'] != 'Malawi') & (
                mf.loc[sample_to_plot]['COUNTRY'] != 'Venezuela'):
            grp_colors[mf.loc[sample_to_plot]['COUNTRY']] = color_Western

        plt.axis('off')
        my_dpi = 72
        figsize = (1000 / my_dpi, 1000 / my_dpi)
        out_file = os.path.join(output, '.'.join([prefix, sample, 'pdf']))
        plt.savefig(out_file, figsize=figsize, dpi=my_dpi)
        plt.close()

@mod2_pcoa.command()
@click.option('--coords', required=True, type=click.Path(resolve_path=True,
              readable=True, exists=True), help='Coordinates file')
@click.option('--mapping_file', required=True, type=click.Path(
              resolve_path=True, readable=True, exists=True),
              help='Mapping file')
@click.option('--color', required=True, type=str,
              help='Metadata category to set color by')
@click.option('--output', required=True, type=click.Path(exists=True,
              writable=True, resolve_path=True), help='Output directory')
@click.option('--prefix', required=True, type=str, help='Output file prefix')
@click.option('--samples', required=False, type=str,
              help='Comma separated list of samples to print')
def gradient(coords, mapping_file, color, output, prefix, samples):
    """Generates as many figures as samples in the coordinates file"""
    o = read(coords, into=OrdinationResults)

    # coordinates
    c_df = pd.DataFrame(o.site, o.site_ids)

    # mapping file
    mf = pd.read_csv(mapping_file, '\t', converters=defaultdict(str),
                     index_col='#SampleID')
    mf = mf.loc[o.site_ids]
    mf[color] = mf[color].convert_objects(convert_numeric=True)

    if samples is None:
        samples = mf.index
    else:
        samples = set(samples.split(',')).intersection(set(o.site_ids))
        samples = mf.loc[samples].index

    numeric = mf[~pd.isnull(mf[color])]
    non_numeric = mf[pd.isnull(mf[color])]

    color_array = plt.cm.RdBu(numeric[color]/max(numeric[color]))

    for sample in samples:

        # plot numeric metadata as colored gradient
        ids = numeric.index
        x, y = c_df.loc[ids][0], c_df.loc[ids][1]
        plt.scatter(x, y, c=numeric[color], cmap=plt.get_cmap('RdBu'),
                    alpha=ALPHA, lw=LINE_WIDTH, edgecolor=color_array*0.6)

        # plt.colorbar()

        # plot non-numeric metadata as gray
        ids = non_numeric.index
        x, y = c_df.loc[ids][0], c_df.loc[ids][1]
        plt.scatter(x, y, c='0.5', alpha=ALPHA, lw=LINE_WIDTH, edgecolor='0.3')

        # plot individual's dot
        try:
            color_index = numeric.index.tolist().index(sample)
        except ValueError:
            color_index = None

        if color_index is None:
            _color = (0.5, 0.5, 0.5)
        else:
            _color = color_array[color_index]

        plt.scatter(c_df.loc[sample][0], c_df.loc[sample][1],
                    color=_color, s=270, edgecolor='w', lw=LINE_WIDTH_WHITE)
        plt.scatter(c_df.loc[sample][0], c_df.loc[sample][1],
                    color=_color, s=250, edgecolor=np.asarray(_color)*0.6,
                    lw=LINE_WIDTH_BLACK)

        plt.axis('off')
        my_dpi = 72
        figsize = (1000 / my_dpi, 1000 / my_dpi)
        out_file = os.path.join(output, '.'.join([prefix, sample, 'pdf']))
        plt.savefig(out_file, figsize=figsize, dpi=my_dpi)
        plt.close()

if __name__ == '__main__':
    mod2_pcoa()


# -*- coding: utf-8 -*-
#   ramaplot.PDistDataset.py
#
#   Copyright (C) 2015 Karl T Debiec
#   All rights reserved.
#
#   This software may be modified and distributed under the terms of the
#   BSD license. See the LICENSE file for details.
"""
Manages probability distribution datasets.
"""
################################### MODULES ###################################
from __future__ import absolute_import,division,print_function,unicode_literals
from .myplotspec.Dataset import Dataset
################################### CLASSES ###################################
class PDistDataset(Dataset):
    """
    Manages probability distribution datasets.

    Generates probability distribution from series of Φ/Ψ values,
    representing either the probability of Φ/Ψ, or the expectation value
    of a selected measurement (e.g. energy) at that Φ/Ψ.

    Input data should be providied in a whitespace-delimited text file
    including columns for Φ, Ψ, and any additional data, such as this
    output from `cpptraj`'s `multidihedral` command::

      #Frame       phi:2     psi:2    chip:2 ...
             1  -62.1431  144.6768   72.2964 ...
             2  -63.2487  151.6551   71.9101 ...
           ...       ...       ...       ...

    """

    @classmethod
    def get_cache_key(cls, infile, phikey="phi", psikey="psi",
        zkey="free energy", mode="hist", bins=72, bandwidth=5, wrap=True,
        mask_cutoff=None,
        calc_populations=False, plot_populations=False,
        *args, **kwargs):
        """
        Generates tuple of arguments to be used as key for dataset
        cache.

        Arguments documented under :func:`__init__`.
        """
        from os.path import expandvars

        if zkey in ["free energy", "probability"]:
            x_bins, y_bins = cls.process_bins_arg(bins, dim=2)
            bins = (tuple(x_bins), tuple(y_bins))
        else:
            x_bins, y_bins, z_bins = cls.process_bins_arg(bins, dim=3)
            bins = (tuple(x_bins), tuple(y_bins), tuple(z_bins))

        if mode == "hist":
            return (cls, expandvars(infile), phikey, psikey, zkey, mode, bins,
                    wrap, mask_cutoff, calc_populations, plot_populations)
        elif mode == "kde":
            return (cls, expandvars(infile), phikey, psikey, zkey, mode, bins,
                    bandwidth, wrap, mask_cutoff, calc_populations,
                    plot_populations)

    @staticmethod
    def process_bins_arg(bins, dim=2):
        """
        Processes bin argument.

        Arguments:
          bins (int, list, ndarray): Bins to use for histogram or grid
            to use for kernel density estimate; if int, number of bins
            or gride points between -180° and 180° in Φ and Ψ, if list
            or ndarray, bins or grid directly

        Returns:
          out_bins (tuple): Processed bins
        """
        import numpy as np
        if dim == 2:
            if isinstance(bins, int):
                x_bins = y_bins = np.linspace(-180, 180, bins + 1)
            elif isinstance(bins, list):
                if len(bins) == 2:
                    if isinstance(bins[0], int):
                        x_bins = np.linspace(-180, 180, bins[0] + 1)
                    elif isinstance(bins[0], list):
                        x_bins = np.array(bins[0])
                    if isinstance(bins[1], int):
                        y_bins = np.linspace(-180, 180, bins[1] + 1)
                    elif isinstance(bins[1], list):
                        y_bins = np.array(bins[1])
                else:
                    x_bins = y_bins = np.array(bins)
            elif isinstance(bins, np.ndarray):
                x_bins = y_bins = bins
            return x_bins, y_bins
        elif dim == 3:
            if isinstance(bins, int):
                x_bins = y_bins = z_bins = np.linspace(-180, 180, bins + 1)
            elif isinstance(bins, list):
                if len(bins) == 2:
                    if isinstance(bins[0], int):
                        x_bins = y_bins = np.linspace(-180, 180, bins[0] + 1)
                    elif (isinstance(bins[0], list)
                    or    isinstance(bins[0], np.ndarray)):
                        x_bins = y_bins = np.array(bins[0])
                    if isinstance(bins[1], int):
                        z_bins = np.linspace(-180, 180, bins[1] + 1)
                    elif (isinstance(bins[1], list)
                    or    isinstance(bins[1], np.ndarray)):
                        z_bins = np.array(bins[1])
                elif len(bins) == 3:
                    if isinstance(bins[0], int):
                        x_bins = np.linspace(-180, 180, bins[0] + 1)
                    elif (isinstance(bins[0], list)
                    or    isinstance(bins[0], np.ndarray)):
                        x_bins = np.array(bins[0])
                    if isinstance(bins[1], int):
                        y_bins = np.linspace(-180, 180, bins[1] + 1)
                    elif (isinstance(bins[1], list)
                    or    isinstance(bins[1], np.ndarray)):
                        y_bins = np.array(bins[1])
                    if isinstance(bins[2], int):
                        z_bins = np.linspace(-180, 180, bins[2] + 1)
                    elif (isinstance(bins[2], list)
                    or    isinstance(bins[2], np.ndarray)):
                        z_bins = np.array(bins[2])
                else:
                    x_bins = y_bins = z_bins = np.array(bins)
            elif isinstance(bins, np.ndarray):
                x_bins = y_bins = z_bins = bins
            return x_bins, y_bins, z_bins
        else:
            raise TypeError()

    def __init__(self, phikey="phi", psikey="psi", zkey="free energy",
        mode="hist", bins=72, bandwidth=5, wrap=True, mask_cutoff=None,
        calc_populations=False, plot_populations=False,
        verbose=1, debug=0, **kwargs):
        """
        Arguments:
          infile (str): Path to text input file, may contain environment
            variables
          phikey (str): Key from which to load Φ
          psikey (str): Key from which to load Ψ
          zkey (str): Key from which to load distribution; if 'free
            energy' or 'probability', the 2D probability density of Φ
            and Ψ will be calculated and the selected representation
            returned; for other values a third dimension will be loaded
            from the `zkey` column of `infile`, the 3D probability
            density of Φ, Ψ, and `zkey` will be calculated, and the
            expectation value of `zkey` as a function of Φ and Ψ will be
            returned
          mode (str): Method of calculating probability distribution;
            may be either 'hist', to use a histogram, or 'kde', to use a
            kernel density estimate
          bins (int, list, ndarray): Bins to use for histogram or grid
            to use for kernel density estimate; if int, number of bins
            or gride points between -180° and 180° in Φ and Ψ, if list
            or ndarray, bins or grid directly
          bandwidth (float, optional): Bandwidth to use for kernel
            density estimate
          wrap (bool): Wrap x and y coordinates between 180° and 360° to
            between -180° and 0°
          wrap_z (bool): Wrap z coordinates between -180° and 0 to between 180°
            and 360°; probably only useful for plotting ω
          mask_cutoff (float): Cutoff beyond which distribution is
            masked, if `zkey` is 'free energy', this is a the maximum
            free energy above which the mask will be set, and if `zkey`
            is 'probability', this is the minimum probability below
            which the mask will be set
          hist_kw: Keyword arguments passed to numpy.histogram2d or
            numpy.histogramdd
          kde_kw: Keyword arguments passed to
            sklearn.neighbors.KernelDensity
          verbose (int): Level of verbose output
          debug (int): Level of debug output
          kwargs (dict): Additional keyword arguments

        .. todo:
          - Fix and validate 3D KDE
          - Auto-detect phikey and psikey
          - Support periodicic kernel density estimate
          - Support variable bandwidth kernel density estimate
        """
        import numpy as np
        import pandas as pd
        from .myplotspec import multi_get_copy

        # Manage arguments
        if str(mode.lower()) not in ["hist", "kde", "none"]:
            raise ValueError("Argument 'mode' does not support provided " +
              "value '{0}', may be 'hist', 'kde', or 'none'".format(mode))
        read_csv_kw = dict(delim_whitespace=True, index_col=0)
        read_csv_kw.update(kwargs.pop("read_csv_kw", {}))

        # Load data
        dataframe = self.load_dataset(verbose=verbose, debug=debug,
          read_csv_kw=read_csv_kw, **kwargs).dataframe
        if wrap:
            dataframe[phikey][dataframe[phikey] > 180] -= 360
            dataframe[psikey][dataframe[psikey] > 180] -= 360

        # Option 0: Store Φ, Ψ
        if mode == "none":

            # Store data in instance variable
            self.x = dataframe[phikey]
            self.y = dataframe[psikey]

        # Option 1: Calculate probability and free energy of Φ, Ψ
        elif zkey in ["free energy", "probability"]:
            x_bins, y_bins = self.process_bins_arg(bins, dim=2)
            x_centers = (x_bins[:-1] + x_bins[1:]) / 2
            y_centers = (y_bins[:-1] + y_bins[1:]) / 2
            x_width = np.mean(x_centers[1:] - x_centers[:-1])
            y_width = np.mean(y_centers[1:] - y_centers[:-1])

            # Option 1a: Use a histogram (fast but noisy)
            if mode == "hist":
                if verbose >= 1:
                    print("calculating probability distribution of " +
                          "'{0}' and '{1}' using a ".format(phikey, psikey) +
                          "histogram")
                hist_kw = dict(normed=False)
                hist_kw.update(kwargs.get("hist_kw", {}))
                hist_kw["bins"] = hist_kw.get("bins", [x_bins, y_bins])

                probability, _, _ = np.histogram2d(
                  dataframe[phikey], dataframe[psikey], **hist_kw)

            # Option 1b: Use a kernel density estimate (smooth but slow)
            elif mode == "kde":
                if verbose >= 1:
                    print("calculating probability distribution of " +
                          "'{0}' and '{1}' using a ".format(phikey, psikey) +
                          "kernel density estimate")
                from sklearn.neighbors import KernelDensity

                kde_kw = multi_get_copy("kde_kw", kwargs, {})
                kde_kw["bandwidth"] = kde_kw.get("bandwidth", bandwidth)
                xg, yg = np.meshgrid(x_centers, y_centers)
                xyg = np.vstack([yg.ravel(), xg.ravel()]).T
                samples = np.column_stack((dataframe[phikey],
                  dataframe[psikey]))

                kde = KernelDensity(**kde_kw)
                kde.fit(samples)
                probability_series = np.exp(kde.score_samples(xyg))
                probability = np.zeros((x_centers.size, y_centers.size))
                for phi, psi, p in np.column_stack((xyg, probability_series)):
                    x_index = np.where(x_centers == phi)[0][0]
                    y_index = np.where(y_centers == psi)[0][0]
                    probability[x_index, y_index] = p

            # Normalize and calculate free energy
            probability /= np.nansum(probability)
            free_energy = -1 * np.log(probability)
            free_energy[np.isinf(free_energy)] = np.nan
            free_energy -= np.nanmin(free_energy)

            # Store data in instance variable
            self.x_centers = x_centers
            self.y_centers = y_centers
            self.x_width = x_width
            self.y_width = y_width
            self.x_bins = x_bins
            self.y_bins = y_bins
            self.x = dataframe[phikey]
            self.y = dataframe[psikey]

        # Option 2: Calculate mean value of a third observable as a
        #   function of Φ, Ψ
        else:
            x_bins, y_bins, z_bins = self.process_bins_arg(bins, dim=3)
            x_centers = (x_bins[:-1] + x_bins[1:]) / 2
            y_centers = (y_bins[:-1] + y_bins[1:]) / 2
            z_centers = (z_bins[:-1] + z_bins[1:]) / 2
            x_width = np.mean(x_centers[1:] - x_centers[:-1])
            y_width = np.mean(y_centers[1:] - y_centers[:-1])

            if kwargs.get("wrap_z"):
                dataframe[zkey][dataframe[zkey] < 0] += 360

            # Option 2a: Use a histogram (fast but noisy)
            if mode == "hist":
                if verbose >= 1:
                    print("calculating mean value  of '{0}'".format(zkey) +
                          "as a function of '{0}' and ".format(phikey) +
                          "'{0}' using a histogram".format(psikey))
                hist_kw = dict(normed=True)
                hist_kw.update(kwargs.get("hist_kw", {}))
                hist_kw["bins"] = hist_kw.get("bins", [x_bins, y_bins, z_bins])

                prob_xyz, _ = np.histogramdd(np.column_stack(
                  (dataframe[phikey], dataframe[psikey], dataframe[zkey])),
                  **hist_kw)
                probability = np.sum(prob_xyz, axis=2)
                prob_z_given_xy = prob_xyz / probability[:,:,np.newaxis]
                weighted_z = prob_z_given_xy*z_centers[np.newaxis,np.newaxis,:]
                mean_z = np.sum(weighted_z, axis=2)

            # Option 2b: Use a kernel density estimate (smooth but slow)
            elif mode == "kde":
                raise()
                from copy import copy
                from sklearn.neighbors import KernelDensity

                kde_kw = multi_get_copy("kde_kw", kwargs, {})
                kde_kw["bandwidth"] = kde_kw.get("bandwidth", bandwidth)
                # Only a single bandwidth is supported; scale z
                #   dimension to span range of 120-240
#                scale_range = 340
                z = copy(dataframe[zkey])
#                z -= z.min()        # shift bottom to 0
#                z_range = z.max()     # save max

#                z *= (scale_range / z_range)
#                z += (360 - scale_range) / 2    # Give buffer on top and bottom

                xg, yg, zg = np.meshgrid(x_centers, y_centers, z_centers)
                xyzg = np.vstack([xg.ravel(), yg.ravel(), zg.ravel()]).T
                samples = np.column_stack((dataframe[phikey],
                  dataframe[psikey], z))

                kde = KernelDensity(**kde_kw)
                kde.fit(samples)
                probability_series = np.exp(kde.score_samples(xyzg))
                prob_xyz = np.zeros((x_centers.size, y_centers.size,
                  z_centers.size), np.float) * np.nan
                for phi,psi,z,p in np.column_stack((xyzg, probability_series)):
                    x_index = np.where(x_centers == phi)[0][0]
                    y_index = np.where(y_centers == psi)[0][0]
                    z_index = np.where(z_centers == z)[0][0]
                    prob_xyz[x_index, y_index, z_index] = p
                prob_xyz /= np.sum(prob_xyz)
                probability = np.sum(prob_xyz, axis=2)
                prob_z_given_xy = prob_xyz / probability[:,:,np.newaxis]
                weighted_z = prob_z_given_xy*z_centers[np.newaxis,np.newaxis,:]
                mean_z = np.sum(weighted_z, axis=2)

#                mean_z -= (360 - scale_range) / 2  # Shift back down
#                mean_z *= (z_range / scale_range) # Back from degrees to E
#                free_energy *= 627.503              # Convert to kcal/mol

            # Normalize and calculate free energy
            probability /= np.nansum(probability)
            free_energy = -1 * np.log(probability)
            free_energy[np.isinf(free_energy)] = np.nan
            free_energy -= np.nanmin(free_energy)

            # Store data in instance variable
            self.x_centers = x_centers
            self.y_centers = y_centers
            self.x_width = x_width
            self.y_width = y_width
            self.x_bins = x_bins
            self.y_bins = y_bins
            self.x = dataframe[phikey]
            self.y = dataframe[psikey]

        # Prepare mask
        if mode == "none":
            pass
        elif zkey == "probability":
            self.dist = probability
            if mask_cutoff is not None:
                self.mask = np.ma.masked_where(np.logical_and(
                  probability >= mask_cutoff,
                  np.logical_not(np.isnan(probability))),
                  np.ones_like(probability))
            else:
                self.mask = np.ma.masked_where(
                  np.logical_not(np.isnan(free_energy)),
                  np.ones_like(free_energy))
        elif zkey == "free energy":
            self.dist = free_energy
            if mask_cutoff is not None:
                self.mask = np.ma.masked_where(np.logical_and(
                  free_energy <= mask_cutoff,
                  np.logical_not(np.isnan(free_energy))),
                  np.ones_like(free_energy))
            else:
                self.mask = np.ma.masked_where(
                  np.logical_not(np.isnan(free_energy)),
                  np.ones_like(free_energy))
        else:
            self.dist = mean_z
            if mask_cutoff is not None:
                self.mask = np.ma.masked_where(np.logical_and(
                  free_energy <= mask_cutoff,
                  np.logical_not(np.isnan(free_energy))),
                  np.ones_like(free_energy))
            else:
                self.mask = np.ma.masked_where(
                  np.logical_not(np.isnan(free_energy)),
                  np.ones_like(free_energy))

        # Calculate state populations
        if calc_populations:
            states = kwargs.get("states",  [
              ("β",       -151,  151),
              ("PPII",     -66,  140),
              ("ξ",       -145,   55),
              ("γ'",       -81,   65),
              ("α",        -70,  -25),
              ("$L_α$",     55,   45),
              ("γ",         73,  -35),
              ("PPII'",     56, -124),
              ("plateau", -100, -130)])
            state_radius = kwargs.get("state_radius", 45)

            distances = np.zeros((len(states), len(x_centers), len(y_centers)))
            xs = []
            ys = []
            for i, (state, x, y) in enumerate(states):
                xs += [x]
                ys += [y]
                # There must be a better way to do this, but this works
                for j, xc in enumerate(x_centers):
                    for k, yc in enumerate(y_centers):
                        dx = (xc - x)
                        if dx <= -180 or dx >= 180:
                            dx = 360 - dx
                        else:
                            dx = dx
                        dy = (yc - y)
                        if dy <= -180 or dy >= 180:
                            dy = 360 - dy
                        else:
                            dy = dy
                        distances[i,j,k] = np.sqrt(dx**2 + dy**2)
            assignments = np.argmin(distances, axis=0)
            assignments[np.min(distances, axis=0) >= state_radius] = \
              len(states) + 1

            index, state_populations = [], []
            for i, (state, x, y) in enumerate(states):
                index += [state]
                state_populations += [(x, y,
                  np.nansum(probability[assignments==i]))]
            state_populations = pd.DataFrame(state_populations, index=index,
              columns=["Φ center", "Ψ center", "population"])
            self.state_populations = state_populations

            if verbose >= 1:
                print(state_populations)

            if plot_populations:
                self.dist = assignments
                self.mask = np.ma.masked_where(
                  np.logical_not(assignments == len(states) + 1),
                  np.ones_like(assignments))
                self.x = np.array(xs)
                self.y = np.array(ys)
                label, label_kw = [], []
                from .myplotspec import multi_get_copy
                default_label_kw = multi_get_copy(["default_label_kw",
                  "label_kw"], kwargs, {})
                for index, row in state_populations.iterrows():
                    label += ["{0}\n{1:2d}%".format(index,
                     int(row["population"]*100))]
                    label_kw += [default_label_kw.copy()]
                    label_kw[-1]["x"] = row["Φ center"]
                    label_kw[-1]["y"] = row["Ψ center"]
                self.label = label
                self.label_kw = label_kw


# Author: Vlad Niculae
# License: BSD 3 clause

import sys

import numpy as np

from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import SkipTest
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_false
from sklearn.utils.testing import if_safe_multiprocessing_with_blas

from sklearn.decomposition import SparsePCA, MiniBatchSparsePCA
from sklearn.utils import check_random_state

def generate_toy_data(n_components, n_samples, image_size, random_state=None):
    n_features = image_size[0] * image_size[1]

    rng = check_random_state(random_state)
    U = rng.randn(n_samples, n_components)
    V = rng.randn(n_components, n_features)

    centers = [(3, 3), (6, 7), (8, 1)]
    sz = [1, 2, 1]
    for k in range(n_components):
        img = np.zeros(image_size)
        xmin, xmax = centers[k][0] - sz[k], centers[k][0] + sz[k]
        ymin, ymax = centers[k][1] - sz[k], centers[k][1] + sz[k]
        img[xmin:xmax][:, ymin:ymax] = 1.0
        V[k, :] = img.ravel()

    # Y is defined by : Y = UV + noise
    Y = np.dot(U, V)
    Y += 0.1 * rng.randn(Y.shape[0], Y.shape[1])  # Add noise
    return Y, U, V

# SparsePCA can be a bit slow. To avoid having test times go up, we
# test different aspects of the code in the same test

def test_correct_shapes():
    rng = np.random.RandomState(0)
    X = rng.randn(12, 10)
    spca = SparsePCA(n_components=8, random_state=rng)
    U = spca.fit_transform(X)
    assert_equal(spca.components_.shape, (8, 10))
    assert_equal(U.shape, (12, 8))
    # test overcomplete decomposition
    spca = SparsePCA(n_components=13, random_state=rng)
    U = spca.fit_transform(X)
    assert_equal(spca.components_.shape, (13, 10))
    assert_equal(U.shape, (12, 13))

def test_fit_transform():
    alpha = 1
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 10, (8, 8), random_state=rng)  # wide array
    spca_lars = SparsePCA(n_components=3, method='lars', alpha=alpha,
                          random_state=0)
    spca_lars.fit(Y)

    # Test that CD gives similar results
    spca_lasso = SparsePCA(n_components=3, method='cd', random_state=0,
                           alpha=alpha)
    spca_lasso.fit(Y)
    assert_array_almost_equal(spca_lasso.components_, spca_lars.components_)

@if_safe_multiprocessing_with_blas
def test_fit_transform_parallel():
    alpha = 1
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 10, (8, 8), random_state=rng)  # wide array
    spca_lars = SparsePCA(n_components=3, method='lars', alpha=alpha,
                          random_state=0)
    spca_lars.fit(Y)
    U1 = spca_lars.transform(Y)
    # Test multiple CPUs
    spca = SparsePCA(n_components=3, n_jobs=2, method='lars', alpha=alpha,
                     random_state=0).fit(Y)
    U2 = spca.transform(Y)
    assert_true(not np.all(spca_lars.components_ == 0))
    assert_array_almost_equal(U1, U2)

def test_transform_nan():
    # Test that SparsePCA won't return NaN when there is 0 feature in all
    # samples.
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 10, (8, 8), random_state=rng)  # wide array
    Y[:, 0] = 0
    estimator = SparsePCA(n_components=8)
    assert_false(np.any(np.isnan(estimator.fit_transform(Y))))

def test_fit_transform_tall():
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 65, (8, 8), random_state=rng)  # tall array
    spca_lars = SparsePCA(n_components=3, method='lars',
                          random_state=rng)
    U1 = spca_lars.fit_transform(Y)
    spca_lasso = SparsePCA(n_components=3, method='cd', random_state=rng)
    U2 = spca_lasso.fit(Y).transform(Y)
    assert_array_almost_equal(U1, U2)

def test_initialization():
    rng = np.random.RandomState(0)
    U_init = rng.randn(5, 3)
    V_init = rng.randn(3, 4)
    model = SparsePCA(n_components=3, U_init=U_init, V_init=V_init, max_iter=0,
                      random_state=rng)
    model.fit(rng.randn(5, 4))
    assert_array_equal(model.components_, V_init)

def test_mini_batch_correct_shapes():
    rng = np.random.RandomState(0)
    X = rng.randn(12, 10)
    pca = MiniBatchSparsePCA(n_components=8, random_state=rng)
    U = pca.fit_transform(X)
    assert_equal(pca.components_.shape, (8, 10))
    assert_equal(U.shape, (12, 8))
    # test overcomplete decomposition
    pca = MiniBatchSparsePCA(n_components=13, random_state=rng)
    U = pca.fit_transform(X)
    assert_equal(pca.components_.shape, (13, 10))
    assert_equal(U.shape, (12, 13))

def test_mini_batch_fit_transform():
    raise SkipTest("skipping mini_batch_fit_transform.")
    alpha = 1
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 10, (8, 8), random_state=rng)  # wide array
    spca_lars = MiniBatchSparsePCA(n_components=3, random_state=0,
                                   alpha=alpha).fit(Y)
    U1 = spca_lars.transform(Y)
    # Test multiple CPUs
    if sys.platform == 'win32':  # fake parallelism for win32
        import sklearn.externals.joblib.parallel as joblib_par
        _mp = joblib_par.multiprocessing
        joblib_par.multiprocessing = None
        try:
            U2 = MiniBatchSparsePCA(n_components=3, n_jobs=2, alpha=alpha,
                                    random_state=0).fit(Y).transform(Y)
        finally:
            joblib_par.multiprocessing = _mp
    else:  # we can efficiently use parallelism
        U2 = MiniBatchSparsePCA(n_components=3, n_jobs=2, alpha=alpha,
                                random_state=0).fit(Y).transform(Y)
    assert_true(not np.all(spca_lars.components_ == 0))
    assert_array_almost_equal(U1, U2)
    # Test that CD gives similar results
    spca_lasso = MiniBatchSparsePCA(n_components=3, method='cd', alpha=alpha,
                                    random_state=0).fit(Y)
    assert_array_almost_equal(spca_lasso.components_, spca_lars.components_)


"""
选取发布时间为2018年的文章，并对其进行月份统计
"""
import numpy as np
import pandas as pd
from pyecharts import Bar

df = pd.read_csv('sg_articles.csv', header=None, names=["title", "article", "name", "date"])

list1 = []
list2 = []
for j in df['date']:
    # 获取文章发布年份及月份
    time_1 = j.split('-')[0]
    time_2 = j.split('-')[1]
    list1.append(time_1)
    list2.append(time_2)
df['year'] = list1
df['month'] = list2

# 选取发布时间为2018年的文章，并对其进行月份统计
df = df.loc[df['year'] == '2018']
month_message = df.groupby(['month'])
month_com = month_message['month'].agg(['count'])
month_com.reset_index(inplace=True)
month_com_last = month_com.sort_index()

attr = ["{}".format(str(i) + '月') for i in range(1, 12)]
v1 = np.array(month_com_last['count'])
v1 = ["{}".format(int(i)) for i in v1]
bar = Bar("微信文章发布时间分布", title_pos='center', title_top='18', width=800, height=400)
bar.add("", attr, v1, is_stack=True, is_label_show=True)
bar.render("微信文章发布时间分布.html")


# -*- coding: utf-8 -*-
from __future__ import division

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

import sys
sys.path.append('..')

from anc_field_py.ancfield import *
from anc_field_py.ancutil import *

def add_microphones(ancObject):
    # error_mic
    ancObject.AddMic([4,1.6,5.3])

    # reference_front
    ancObject.AddMic([4,1.6,4.5])

    # reference_back
    ancObject.AddMic([4,1.6,6.5])

    # reference_left
    ancObject.AddMic([3,1.6,5.5])

    # reference_right
    ancObject.AddMic([5,1.6,5.5])

    # reference_bottom
    ancObject.AddMic([4,0.6,5.5])

    # reference_top
    ancObject.AddMic([4,2.6,5.5])

    return ancObject

# =========================================================================
# SIMULATION 1
# Calculating noise to microphone paths
# =========================================================================
#
# Trying to run this simulation on CPU failed on an i7-3770, compiling the
# lrs_1.cl file fails. It maybe because the scene's size is too large for
# the CPU. Compiling it for the built in integrated GPU worked though.
#
# We create a simulation and immediately add microphones to it
anc = add_microphones(AncField('gpu', 'models/ie224'))

# noise_source
anc.AddSource([4,1.6,1.0], 5*impulse(2000, 6*32000, 32000))

anc.Visualize(1.6)
(x,y_noise) = anc.Run(4)

# Saving the impulse responses
np.savetxt('ie224-noise-to-error.dat', y_noise[0,:])
np.savetxt('ie224-noise-to-reference_front.dat', y_noise[1,:])
np.savetxt('ie224-noise-to-reference_back.dat', y_noise[2,:])
np.savetxt('ie224-noise-to-reference_left.dat', y_noise[3,:])
np.savetxt('ie224-noise-to-reference_right.dat', y_noise[4,:])
np.savetxt('ie224-noise-to-reference_bottom.dat', y_noise[5,:])
np.savetxt('ie224-noise-to-reference_top.dat', y_noise[6,:])

# =========================================================================
# SIMULATION 2
# Calculating actuator to microphone paths
# =========================================================================
#
# We create a simulation and immediately add microphones to it
anc = add_microphones(AncField('gpu', 'models/ie224'))

# actuator
anc.AddSource([4,1.6,5.5], 5*impulse(2000, 6*32000, 32000))

anc.Visualize(1.6)
(x,y_actuator) = anc.Run(4)

# Saving the impulse responses
np.savetxt('ie224-actuator-to-error.dat', y_actuator[0,:])
np.savetxt('ie224-actuator-to-reference_front.dat', y_actuator[1,:])
np.savetxt('ie224-actuator-to-reference_back.dat', y_actuator[2,:])
np.savetxt('ie224-actuator-to-reference_left.dat', y_actuator[3,:])
np.savetxt('ie224-actuator-to-reference_right.dat', y_actuator[4,:])
np.savetxt('ie224-actuator-to-reference_bottom.dat', y_actuator[5,:])
np.savetxt('ie224-actuator-to-reference_top.dat', y_actuator[6,:])

# =========================================================================
# GENERATING IMAGES FOR THE REPORT
# Calculating actuator to microphone paths
# =========================================================================
#
# Saving figures for the field simulation report
fig, ax = plt.subplots()
ax.plot(y_noise[0,:])
plt.title('ie224-noise-to-error')
fig.savefig('ie224-noise-to-error.png')

fig, ax = plt.subplots()
ax.plot(y_noise[1,:])
plt.title('ie224-noise-to-reference_front')
fig.savefig('ie224-noise-to-reference_front.png')

fig, ax = plt.subplots()
ax.plot(y_noise[2,:])
plt.title('ie224-noise-to-reference_back')
fig.savefig('ie224-noise-to-reference_back.png')

fig, ax = plt.subplots()
ax.plot(y_noise[3,:])
plt.title('ie224-noise-to-reference_left')
fig.savefig('ie224-noise-to-reference_left.png')

fig, ax = plt.subplots()
ax.plot(y_noise[4,:])
plt.title('ie224-noise-to-reference_right')
fig.savefig('ie224-noise-to-reference_right.png')

fig, ax = plt.subplots()
ax.plot(y_noise[5,:])
plt.title('ie224-noise-to-reference_bottom')
fig.savefig('ie224-noise-to-reference_bottom.png')

fig, ax = plt.subplots()
ax.plot(y_noise[6,:])
plt.title('ie224-noise-to-reference_top')
fig.savefig('ie224-noise-to-reference_top.png')

# Saving figures for the field simulation report
fig, ax = plt.subplots()
ax.plot(y_actuator[0,:])
plt.title('ie224-actuator-to-error')
fig.savefig('ie224-actuator-to-error.png')

fig, ax = plt.subplots()
ax.plot(y_actuator[1,:])
plt.title('ie224-actuator-to-reference_front')
fig.savefig('ie224-actuator-to-reference_front.png')

fig, ax = plt.subplots()
ax.plot(y_actuator[2,:])
plt.title('ie224-actuator-to-reference_back')
fig.savefig('ie224-actuator-to-reference_back.png')

fig, ax = plt.subplots()
ax.plot(y_actuator[3,:])
plt.title('ie224-actuator-to-reference_left')
fig.savefig('ie224-actuator-to-reference_left.png')

fig, ax = plt.subplots()
ax.plot(y_actuator[4,:])
plt.title('ie224-actuator-to-reference_right')
fig.savefig('ie224-actuator-to-reference_right.png')

fig, ax = plt.subplots()
ax.plot(y_actuator[5,:])
plt.title('ie224-actuator-to-reference_bottom')
fig.savefig('ie224-actuator-to-reference_bottom.png')

fig, ax = plt.subplots()
ax.plot(y_actuator[6,:])
plt.title('ie224-actuator-to-reference_top')
fig.savefig('ie224-actuator-to-reference_top.png')


# -*- coding: utf-8 -*-
# <nbformat>3.0</nbformat>

# <headingcell level=3>

# Example from pycse 1

# <codecell>

# copied from http://kitchingroup.cheme.cmu.edu/blog/tag/events/

from pycse import odelay
import matplotlib.pyplot as plt
import numpy as np

def ode(Y,x):
    y1, y2 = Y
    dy1dx = y2
    dy2dx = -y1
    return [dy1dx, dy2dx]

def event1(Y, x):
    y1, y2 = Y
    value = y2 - (-1.0)
    isterminal = True
    direction  = 0
    return value, isterminal, direction

def event2(Y, x):
    dy1dx, dy2dx = ode(Y,x)
    value = dy1dx - 0.0
    isterminal = False
    direction = -1  # derivative is decreasing towards a maximum
    return value, isterminal, direction

Y0 = [2.0, 1.0]
xspan = np.linspace(0, 5)
X, Y, XE, YE, IE = odelay(ode, Y0, xspan, events=[event1, event2])

plt.plot(X, Y)
for ie,xe,ye in zip(IE, XE, YE):
    if ie == 1: #this is the second event
        y1,y2 = ye
        plt.plot(xe, y1, 'ro')

plt.legend(['$y_1$', '$y_2$'], loc='best')
#plt.savefig('images/odelay-mult-eq.png')
plt.show()

# <headingcell level=3>

# Example from pycse 2

# <codecell>

# copied from: http://kitchingroup.cheme.cmu.edu/pycse/pycse.html#sec-10-1-8

# 10.1.8 Stopping the integration of an ODE at some condition

from pycse import *
import numpy as np

k = 0.23
Ca0 = 2.3

def dCadt(Ca, t):
    return -k * Ca**2

def stop(Ca, t):
    isterminal = True
    direction = 0
    value = 1.0 - Ca
    return value, isterminal, direction

tspan = np.linspace(0.0, 10.0)

t, CA, TE, YE, IE = odelay(dCadt, Ca0, tspan, events=[stop])

print 'At t = {0:1.2f} seconds the concentration of A is {1:1.2f} mol/L.'.format(t[-1], float(CA[-1]))

# <headingcell level=3>

# fsolve example

# <codecell>

from math import cos

def func(x):
    return x + 2*cos(x)  # finds where this is zero

def func2(x):
    out = [x[0]*cos(x[1]) - 4]
    out.append(x[1]*x[0] - x[1] - 5)
    return out  # finds where both elements of this array are zero

from scipy.optimize import fsolve
x0 = fsolve(func, 0.3)  # initial guess
print x0
print func(x0)
#-1.02986652932

x02 = fsolve(func2, [1, 1]) # initial guesses
print x02
print func2(x02)
#[ 6.50409711  0.90841421]


#!/usr/bin/env python3

import os
import os.path
import sys
import numpy as np
import matplotlib.pyplot as plt

from features_common import match_angle, base_plot

def outlier_frequency_plot(path, angles, threshold):
    f, ax = base_plot()
    ax.plot(100 * np.cumsum(np.abs(angles) > threshold) / angles.size)
    ax.set_xlabel("Match number")
    ax.set_ylabel("Outlier fraction (%)")
    ax.set_ylim([0, 100])
    f.savefig(path, bbox_inches='tight')
    plt.close(f)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        path = "."
    else:
        path = sys.argv[1]

    # Produce outlier plots for all directories containing outlier_threshold.txt
    for root, subdirs, files in os.walk(path):
        if "matches.txt" in files:
            shape = np.loadtxt(os.path.join(root, "shape.txt"))
            matches = np.loadtxt(os.path.join(root, "matches.txt"), comments="#")
            threshold = np.loadtxt(os.path.join(root, "outlier_threshold.txt"))

            if threshold.size == 1:
                print("outlier_analysis.py: " + root)
                # Compute matches angles
                angles = match_angle(matches, shape)
                outlier_frequency_plot(os.path.join(root, "plot_outliers.pdf"), angles, threshold)
            else:
                print("outlier_analysis.py: " + root + " --- empty outlier_threshold.txt")


import numpy as np
try:
    from matplotlib import lines
except ImportError:
    pass

__all__ = ['CanvasToolBase', 'ToolHandles']

def _pass(*args):
    pass

class CanvasToolBase(object):
    """Base canvas tool for matplotlib axes.

    Parameters
    ----------
    ax : :class:`matplotlib.axes.Axes`
        Matplotlib axes where tool is displayed.
    on_move : function
        Function called whenever a control handle is moved.
        This function must accept the end points of line as the only argument.
    on_release : function
        Function called whenever the control handle is released.
    on_enter : function
        Function called whenever the "enter" key is pressed.
    useblit : bool
        If True, update canvas by blitting, which is much faster than normal
        redrawing (turn off for debugging purposes).
    """
    def __init__(self, ax, on_move=None, on_enter=None, on_release=None,
                 useblit=True):
        self.ax = ax
        self.canvas = ax.figure.canvas
        self.img_background = None
        self.cids = []
        self._artists = []
        self.active = True

        if useblit:
            self.connect_event('draw_event', self._blit_on_draw_event)
        self.useblit = useblit

        self.callback_on_move = _pass if on_move is None else on_move
        self.callback_on_enter = _pass if on_enter is None else on_enter
        self.callback_on_release = _pass if on_release is None else on_release

        self.connect_event('key_press_event', self._on_key_press)

    def connect_event(self, event, callback):
        """Connect callback with an event.

        This should be used in lieu of `figure.canvas.mpl_connect` since this
        function stores call back ids for later clean up.
        """
        cid = self.canvas.mpl_connect(event, callback)
        self.cids.append(cid)

    def disconnect_events(self):
        """Disconnect all events created by this widget."""
        for c in self.cids:
            self.canvas.mpl_disconnect(c)

    def ignore(self, event):
        """Return True if event should be ignored.

        This method (or a version of it) should be called at the beginning
        of any event callback.
        """
        return not self.active

    def set_visible(self, val):
        for artist in self._artists:
            artist.set_visible(val)

    def _blit_on_draw_event(self, event=None):
        self.img_background = self.canvas.copy_from_bbox(self.ax.bbox)
        self._draw_artists()

    def _draw_artists(self):
        for artist in self._artists:
            self.ax.draw_artist(artist)

    def remove(self):
        """Remove artists and events from axes.

        Note that the naming here mimics the interface of Matplotlib artists.
        """
        #TODO: For some reason, RectangleTool doesn't get properly removed
        self.disconnect_events()
        for a in self._artists:
            a.remove()

    def redraw(self):
        """Redraw image and canvas artists.

        This method should be called by subclasses when artists are updated.
        """
        if self.useblit and self.img_background is not None:
            self.canvas.restore_region(self.img_background)
            self._draw_artists()
            self.canvas.blit(self.ax.bbox)
        else:
            self.canvas.draw_idle()

    def _on_key_press(self, event):
        if event.key == 'enter':
            self.callback_on_enter(self.geometry)
            self.set_visible(False)
            self.redraw()

    @property
    def geometry(self):
        """Geometry information that gets passed to callback functions."""
        return None

class ToolHandles(object):
    """Control handles for canvas tools.

    Parameters
    ----------
    ax : :class:`matplotlib.axes.Axes`
        Matplotlib axes where tool handles are displayed.
    x, y : 1D arrays
        Coordinates of control handles.
    marker : str
        Shape of marker used to display handle. See `matplotlib.pyplot.plot`.
    marker_props : dict
        Additional marker properties. See :class:`matplotlib.lines.Line2D`.
    """
    def __init__(self, ax, x, y, marker='o', marker_props=None):
        self.ax = ax

        props = dict(marker=marker, markersize=7, mfc='w', ls='none',
                     alpha=0.5, visible=False)
        props.update(marker_props if marker_props is not None else {})
        self._markers = lines.Line2D(x, y, animated=True, **props)
        self.ax.add_line(self._markers)
        self.artist = self._markers

    @property
    def x(self):
        return self._markers.get_xdata()

    @property
    def y(self):
        return self._markers.get_ydata()

    def set_data(self, pts, y=None):
        """Set x and y positions of handles"""
        if y is not None:
            x = pts
            pts = np.array([x, y])
        self._markers.set_data(pts)

    def set_visible(self, val):
        self._markers.set_visible(val)

    def set_animated(self, val):
        self._markers.set_animated(val)

    def draw(self):
        self.ax.draw_artist(self._markers)

    def closest(self, x, y):
        """Return index and pixel distance to closest index."""
        pts = np.transpose((self.x, self.y))
        # Transform data coordinates to pixel coordinates.
        pts = self.ax.transData.transform(pts)
        diff = pts - ((x, y))
        dist = np.sqrt(np.sum(diff**2, axis=1))
        return np.argmin(dist), np.min(dist)


""" This is the Bokeh charts testing interface.

"""
#-----------------------------------------------------------------------------
# Copyright (c) 2012 - 2014, Continuum Analytics, Inc. All rights reserved.
#
# Powered by the Bokeh Development Team.
#
# The full license is in the file LICENSE.txt, distributed with this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

from __future__ import absolute_import

from collections import OrderedDict
import unittest

import numpy as np
from numpy.testing import assert_array_equal
import pandas as pd

from bokeh.charts import DataAdapter

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

class TestDataAdapter(unittest.TestCase):
    def setUp(self):
        self._values = OrderedDict()
        self._values['first'] = [2., 5., 3.]
        self._values['second'] = [4., 1., 4.]
        self._values['third'] = [6., 4., 3.]

    def test_list(self):
        values = list(self._values.values())
        da = DataAdapter(values)

        self.assertEqual(da.values(), list(self._values.values()))
        self.assertEqual(da.columns, ['0', '1', '2'])
        self.assertEqual(da.keys(), ['0', '1', '2'])
        self.assertEqual(da.index, ['a', 'b', 'c'])

    def test_array(self):
        values = np.array(list(self._values.values()))
        da = DataAdapter(values)

        assert_array_equal(da.values(), list(self._values.values()))
        self.assertEqual(da.columns, ['0', '1', '2'])
        self.assertEqual(da.keys(), ['0', '1', '2'])
        self.assertEqual(da.index, ['a', 'b', 'c'])

    def test_pandas(self):
        values = pd.DataFrame(self._values)
        da = DataAdapter(values)

        # TODO: THIS SHOULD BE FIXED..
        #self.assertEqual(da.values(), list(self._values.values()))
        self.assertEqual(da.columns, ['first', 'second', 'third'])
        self.assertEqual(da.keys(), ['first', 'second', 'third'])
        # We expect data adapter index to be the same as the underlying pandas
        # object and not the default created by DataAdapter
        self.assertEqual(da.index, [0, 1, 2])

    def test_ordered_dict(self):
        da = DataAdapter(self._values)

        self.assertEqual(da.values(), list(self._values.values()))
        self.assertEqual(da.columns, ['first', 'second', 'third'])
        self.assertEqual(da.keys(), ['first', 'second', 'third'])
        self.assertEqual(da.index, ['a', 'b', 'c'])

    def test_blaze_data_no_fields(self):
        import blaze
        valuesdf = pd.DataFrame(self._values)
        values = blaze.Data(valuesdf)
        da = DataAdapter(values)

        assert_array_equal(da.values(), list(self._values.values()))
        self.assertEqual(da.columns, ['first', 'second', 'third'])
        self.assertEqual(da.keys(), ['first', 'second', 'third'])
        self.assertEqual(da.index, [0, 1, 2])

        xs, _values = DataAdapter.get_index_and_data(values, None)
        assert_array_equal([0,1,2], xs)


"""Dirac notation for states."""

from __future__ import print_function, division

from sympy import (cacheit, conjugate, Expr, Function, integrate, oo, sqrt,
                   Tuple)
from sympy.core.compatibility import u, range
from sympy.printing.pretty.stringpict import stringPict
from sympy.physics.quantum.qexpr import QExpr, dispatch_method

__all__ = [
    'KetBase',
    'BraBase',
    'StateBase',
    'State',
    'Ket',
    'Bra',
    'TimeDepState',
    'TimeDepBra',
    'TimeDepKet',
    'Wavefunction'
]

#-----------------------------------------------------------------------------
# States, bras and kets.
#-----------------------------------------------------------------------------

# ASCII brackets
_lbracket = "<"
_rbracket = ">"
_straight_bracket = "|"

# Unicode brackets
# MATHEMATICAL ANGLE BRACKETS
_lbracket_ucode = u("\N{MATHEMATICAL LEFT ANGLE BRACKET}")
_rbracket_ucode = u("\N{MATHEMATICAL RIGHT ANGLE BRACKET}")
# LIGHT VERTICAL BAR
_straight_bracket_ucode = u("\N{LIGHT VERTICAL BAR}")

# Other options for unicode printing of <, > and | for Dirac notation.

# LEFT-POINTING ANGLE BRACKET
# _lbracket = u"\u2329"
# _rbracket = u"\u232A"

# LEFT ANGLE BRACKET
# _lbracket = u"\u3008"
# _rbracket = u"\u3009"

# VERTICAL LINE
# _straight_bracket = u"\u007C"

class StateBase(QExpr):
    """Abstract base class for general abstract states in quantum mechanics.

    All other state classes defined will need to inherit from this class. It
    carries the basic structure for all other states such as dual, _eval_adjoint
    and label.

    This is an abstract base class and you should not instantiate it directly,
    instead use State.
    """

    @classmethod
    def _operators_to_state(self, ops, **options):
        """ Returns the eigenstate instance for the passed operators.

        This method should be overridden in subclasses. It will handle being
        passed either an Operator instance or set of Operator instances. It
        should return the corresponding state INSTANCE or simply raise a
        NotImplementedError. See cartesian.py for an example.
        """

        raise NotImplementedError("Cannot map operators to states in this class. Method not implemented!")

    def _state_to_operators(self, op_classes, **options):
        """ Returns the operators which this state instance is an eigenstate
        of.

        This method should be overridden in subclasses. It will be called on
        state instances and be passed the operator classes that we wish to make
        into instances. The state instance will then transform the classes
        appropriately, or raise a NotImplementedError if it cannot return
        operator instances. See cartesian.py for examples,
        """

        raise NotImplementedError(
            "Cannot map this state to operators. Method not implemented!")

    @property
    def operators(self):
        """Return the operator(s) that this state is an eigenstate of"""
        from .operatorset import state_to_operators  # import internally to avoid circular import errors
        return state_to_operators(self)

    def _enumerate_state(self, num_states, **options):
        raise NotImplementedError("Cannot enumerate this state!")

    def _represent_default_basis(self, **options):
        return self._represent(basis=self.operators)

    #-------------------------------------------------------------------------
    # Dagger/dual
    #-------------------------------------------------------------------------

    @property
    def dual(self):
        """Return the dual state of this one."""
        return self.dual_class()._new_rawargs(self.hilbert_space, *self.args)

    @classmethod
    def dual_class(self):
        """Return the class used to construt the dual."""
        raise NotImplementedError(
            'dual_class must be implemented in a subclass'
        )

    def _eval_adjoint(self):
        """Compute the dagger of this state using the dual."""
        return self.dual

    #-------------------------------------------------------------------------
    # Printing
    #-------------------------------------------------------------------------

    def _pretty_brackets(self, height, use_unicode=True):
        # Return pretty printed brackets for the state
        # Ideally, this could be done by pform.parens but it does not support the angled < and >

        # Setup for unicode vs ascii
        if use_unicode:
            lbracket, rbracket = self.lbracket_ucode, self.rbracket_ucode
            slash, bslash, vert = u('\N{BOX DRAWINGS LIGHT DIAGONAL UPPER RIGHT TO LOWER LEFT}'), \
                                  u('\N{BOX DRAWINGS LIGHT DIAGONAL UPPER LEFT TO LOWER RIGHT}'), \
                                  u('\N{BOX DRAWINGS LIGHT VERTICAL}')
        else:
            lbracket, rbracket = self.lbracket, self.rbracket
            slash, bslash, vert = '/', '\\', '|'

        # If height is 1, just return brackets
        if height == 1:
            return stringPict(lbracket), stringPict(rbracket)
        # Make height even
        height += (height % 2)

        brackets = []
        for bracket in lbracket, rbracket:
            # Create left bracket
            if bracket in set([_lbracket, _lbracket_ucode]):
                bracket_args = [ ' ' * (height//2 - i - 1) +
                                 slash for i in range(height // 2)]
                bracket_args.extend(
                    [ ' ' * i + bslash for i in range(height // 2)])
            # Create right bracket
            elif bracket in set([_rbracket, _rbracket_ucode]):
                bracket_args = [ ' ' * i + bslash for i in range(height // 2)]
                bracket_args.extend([ ' ' * (
                    height//2 - i - 1) + slash for i in range(height // 2)])
            # Create straight bracket
            elif bracket in set([_straight_bracket, _straight_bracket_ucode]):
                bracket_args = [vert for i in range(height)]
            else:
                raise ValueError(bracket)
            brackets.append(
                stringPict('\n'.join(bracket_args), baseline=height//2))
        return brackets

    def _sympystr(self, printer, *args):
        contents = self._print_contents(printer, *args)
        return '%s%s%s' % (self.lbracket, contents, self.rbracket)

    def _pretty(self, printer, *args):
        from sympy.printing.pretty.stringpict import prettyForm
        # Get brackets
        pform = self._print_contents_pretty(printer, *args)
        lbracket, rbracket = self._pretty_brackets(
            pform.height(), printer._use_unicode)
        # Put together state
        pform = prettyForm(*pform.left(lbracket))
        pform = prettyForm(*pform.right(rbracket))
        return pform

    def _latex(self, printer, *args):
        contents = self._print_contents_latex(printer, *args)
        # The extra {} brackets are needed to get matplotlib's latex
        # rendered to render this properly.
        return '{%s%s%s}' % (self.lbracket_latex, contents, self.rbracket_latex)

class KetBase(StateBase):
    """Base class for Kets.

    This class defines the dual property and the brackets for printing. This is
    an abstract base class and you should not instantiate it directly, instead
    use Ket.
    """

    lbracket = _straight_bracket
    rbracket = _rbracket
    lbracket_ucode = _straight_bracket_ucode
    rbracket_ucode = _rbracket_ucode
    lbracket_latex = r'\left|'
    rbracket_latex = r'\right\rangle '

    @classmethod
    def default_args(self):
        return ("psi",)

    @classmethod
    def dual_class(self):
        return BraBase

    def __mul__(self, other):
        """KetBase*other"""
        from sympy.physics.quantum.operator import OuterProduct
        if isinstance(other, BraBase):
            return OuterProduct(self, other)
        else:
            return Expr.__mul__(self, other)

    def __rmul__(self, other):
        """other*KetBase"""
        from sympy.physics.quantum.innerproduct import InnerProduct
        if isinstance(other, BraBase):
            return InnerProduct(other, self)
        else:
            return Expr.__rmul__(self, other)

    #-------------------------------------------------------------------------
    # _eval_* methods
    #-------------------------------------------------------------------------

    def _eval_innerproduct(self, bra, **hints):
        """Evaluate the inner product betweeen this ket and a bra.

        This is called to compute <bra|ket>, where the ket is ``self``.

        This method will dispatch to sub-methods having the format::

            ``def _eval_innerproduct_BraClass(self, **hints):``

        Subclasses should define these methods (one for each BraClass) to
        teach the ket how to take inner products with bras.
        """
        return dispatch_method(self, '_eval_innerproduct', bra, **hints)

    def _apply_operator(self, op, **options):
        """Apply an Operator to this Ket.

        This method will dispatch to methods having the format::

            ``def _apply_operator_OperatorName(op, **options):``

        Subclasses should define these methods (one for each OperatorName) to
        teach the Ket how operators act on it.

        Parameters
        ==========

        op : Operator
            The Operator that is acting on the Ket.
        options : dict
            A dict of key/value pairs that control how the operator is applied
            to the Ket.
        """
        return dispatch_method(self, '_apply_operator', op, **options)

class BraBase(StateBase):
    """Base class for Bras.

    This class defines the dual property and the brackets for printing. This
    is an abstract base class and you should not instantiate it directly,
    instead use Bra.
    """

    lbracket = _lbracket
    rbracket = _straight_bracket
    lbracket_ucode = _lbracket_ucode
    rbracket_ucode = _straight_bracket_ucode
    lbracket_latex = r'\left\langle '
    rbracket_latex = r'\right|'

    @classmethod
    def _operators_to_state(self, ops, **options):
        state = self.dual_class().operators_to_state(ops, **options)
        return state.dual

    def _state_to_operators(self, op_classes, **options):
        return self.dual._state_to_operators(op_classes, **options)

    def _enumerate_state(self, num_states, **options):
        dual_states = self.dual._enumerate_state(num_states, **options)
        return [x.dual for x in dual_states]

    @classmethod
    def default_args(self):
        return self.dual_class().default_args()

    @classmethod
    def dual_class(self):
        return KetBase

    def __mul__(self, other):
        """BraBase*other"""
        from sympy.physics.quantum.innerproduct import InnerProduct
        if isinstance(other, KetBase):
            return InnerProduct(self, other)
        else:
            return Expr.__mul__(self, other)

    def __rmul__(self, other):
        """other*BraBase"""
        from sympy.physics.quantum.operator import OuterProduct
        if isinstance(other, KetBase):
            return OuterProduct(other, self)
        else:
            return Expr.__rmul__(self, other)

    def _represent(self, **options):
        """A default represent that uses the Ket's version."""
        from sympy.physics.quantum.dagger import Dagger
        return Dagger(self.dual._represent(**options))

class State(StateBase):
    """General abstract quantum state used as a base class for Ket and Bra."""
    pass

class Ket(State, KetBase):
    """A general time-independent Ket in quantum mechanics.

    Inherits from State and KetBase. This class should be used as the base
    class for all physical, time-independent Kets in a system. This class
    and its subclasses will be the main classes that users will use for
    expressing Kets in Dirac notation [1]_.

    Parameters
    ==========

    args : tuple
        The list of numbers or parameters that uniquely specify the
        ket. This will usually be its symbol or its quantum numbers. For
        time-dependent state, this will include the time.

    Examples
    ========

    Create a simple Ket and looking at its properties::

        >>> from sympy.physics.quantum import Ket, Bra
        >>> from sympy import symbols, I
        >>> k = Ket('psi')
        >>> k
        |psi>
        >>> k.hilbert_space
        H
        >>> k.is_commutative
        False
        >>> k.label
        (psi,)

    Ket's know about their associated bra::

        >>> k.dual
        <psi|
        >>> k.dual_class()
        <class 'sympy.physics.quantum.state.Bra'>

    Take a linear combination of two kets::

        >>> k0 = Ket(0)
        >>> k1 = Ket(1)
        >>> 2*I*k0 - 4*k1
        2*I*|0> - 4*|1>

    Compound labels are passed as tuples::

        >>> n, m = symbols('n,m')
        >>> k = Ket(n,m)
        >>> k
        |nm>

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Bra-ket_notation
    """

    @classmethod
    def dual_class(self):
        return Bra

class Bra(State, BraBase):
    """A general time-independent Bra in quantum mechanics.

    Inherits from State and BraBase. A Bra is the dual of a Ket [1]_. This
    class and its subclasses will be the main classes that users will use for
    expressing Bras in Dirac notation.

    Parameters
    ==========

    args : tuple
        The list of numbers or parameters that uniquely specify the
        ket. This will usually be its symbol or its quantum numbers. For
        time-dependent state, this will include the time.

    Examples
    ========

    Create a simple Bra and look at its properties::

        >>> from sympy.physics.quantum import Ket, Bra
        >>> from sympy import symbols, I
        >>> b = Bra('psi')
        >>> b
        <psi|
        >>> b.hilbert_space
        H
        >>> b.is_commutative
        False

    Bra's know about their dual Ket's::

        >>> b.dual
        |psi>
        >>> b.dual_class()
        <class 'sympy.physics.quantum.state.Ket'>

    Like Kets, Bras can have compound labels and be manipulated in a similar
    manner::

        >>> n, m = symbols('n,m')
        >>> b = Bra(n,m) - I*Bra(m,n)
        >>> b
        -I*<mn| + <nm|

    Symbols in a Bra can be substituted using ``.subs``::

        >>> b.subs(n,m)
        <mm| - I*<mm|

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Bra-ket_notation
    """

    @classmethod
    def dual_class(self):
        return Ket

#-----------------------------------------------------------------------------
# Time dependent states, bras and kets.
#-----------------------------------------------------------------------------

class TimeDepState(StateBase):
    """Base class for a general time-dependent quantum state.

    This class is used as a base class for any time-dependent state. The main
    difference between this class and the time-independent state is that this
    class takes a second argument that is the time in addition to the usual
    label argument.

    Parameters
    ==========

    args : tuple
        The list of numbers or parameters that uniquely specify the ket. This
        will usually be its symbol or its quantum numbers. For time-dependent
        state, this will include the time as the final argument.
    """

    #-------------------------------------------------------------------------
    # Initialization
    #-------------------------------------------------------------------------

    @classmethod
    def default_args(self):
        return ("psi", "t")

    #-------------------------------------------------------------------------
    # Properties
    #-------------------------------------------------------------------------

    @property
    def label(self):
        """The label of the state."""
        return self.args[:-1]

    @property
    def time(self):
        """The time of the state."""
        return self.args[-1]

    #-------------------------------------------------------------------------
    # Printing
    #-------------------------------------------------------------------------

    def _print_time(self, printer, *args):
        return printer._print(self.time, *args)

    _print_time_repr = _print_time
    _print_time_latex = _print_time

    def _print_time_pretty(self, printer, *args):
        pform = printer._print(self.time, *args)
        return pform

    def _print_contents(self, printer, *args):
        label = self._print_label(printer, *args)
        time = self._print_time(printer, *args)
        return '%s;%s' % (label, time)

    def _print_label_repr(self, printer, *args):
        label = self._print_sequence(self.label, ',', printer, *args)
        time = self._print_time_repr(printer, *args)
        return '%s,%s' % (label, time)

    def _print_contents_pretty(self, printer, *args):
        label = self._print_label_pretty(printer, *args)
        time = self._print_time_pretty(printer, *args)
        return printer._print_seq((label, time), delimiter=';')

    def _print_contents_latex(self, printer, *args):
        label = self._print_sequence(
            self.label, self._label_separator, printer, *args)
        time = self._print_time_latex(printer, *args)
        return '%s;%s' % (label, time)

class TimeDepKet(TimeDepState, KetBase):
    """General time-dependent Ket in quantum mechanics.

    This inherits from ``TimeDepState`` and ``KetBase`` and is the main class
    that should be used for Kets that vary with time. Its dual is a
    ``TimeDepBra``.

    Parameters
    ==========

    args : tuple
        The list of numbers or parameters that uniquely specify the ket. This
        will usually be its symbol or its quantum numbers. For time-dependent
        state, this will include the time as the final argument.

    Examples
    ========

    Create a TimeDepKet and look at its attributes::

        >>> from sympy.physics.quantum import TimeDepKet
        >>> k = TimeDepKet('psi', 't')
        >>> k
        |psi;t>
        >>> k.time
        t
        >>> k.label
        (psi,)
        >>> k.hilbert_space
        H

    TimeDepKets know about their dual bra::

        >>> k.dual
        <psi;t|
        >>> k.dual_class()
        <class 'sympy.physics.quantum.state.TimeDepBra'>
    """

    @classmethod
    def dual_class(self):
        return TimeDepBra

class TimeDepBra(TimeDepState, BraBase):
    """General time-dependent Bra in quantum mechanics.

    This inherits from TimeDepState and BraBase and is the main class that
    should be used for Bras that vary with time. Its dual is a TimeDepBra.

    Parameters
    ==========

    args : tuple
        The list of numbers or parameters that uniquely specify the ket. This
        will usually be its symbol or its quantum numbers. For time-dependent
        state, this will include the time as the final argument.

    Examples
    ========

        >>> from sympy.physics.quantum import TimeDepBra
        >>> from sympy import symbols, I
        >>> b = TimeDepBra('psi', 't')
        >>> b
        <psi;t|
        >>> b.time
        t
        >>> b.label
        (psi,)
        >>> b.hilbert_space
        H
        >>> b.dual
        |psi;t>
    """

    @classmethod
    def dual_class(self):
        return TimeDepKet

class Wavefunction(Function):
    """Class for representations in continuous bases

    This class takes an expression and coordinates in its constructor. It can
    be used to easily calculate normalizations and probabilities.

    Parameters
    ==========

    expr : Expr
           The expression representing the functional form of the w.f.

    coords : Symbol or tuple
           The coordinates to be integrated over, and their bounds

    Examples
    ========

    Particle in a box, specifying bounds in the more primitive way of using
    Piecewise:

        >>> from sympy import Symbol, Piecewise, pi, N
        >>> from sympy.functions import sqrt, sin
        >>> from sympy.physics.quantum.state import Wavefunction
        >>> x = Symbol('x', real=True)
        >>> n = 1
        >>> L = 1
        >>> g = Piecewise((0, x < 0), (0, x > L), (sqrt(2//L)*sin(n*pi*x/L), True))
        >>> f = Wavefunction(g, x)
        >>> f.norm
        1
        >>> f.is_normalized
        True
        >>> p = f.prob()
        >>> p(0)
        0
        >>> p(L)
        0
        >>> p(0.5)
        2
        >>> p(0.85*L)
        2*sin(0.85*pi)**2
        >>> N(p(0.85*L))
        0.412214747707527

    Additionally, you can specify the bounds of the function and the indices in
    a more compact way:

        >>> from sympy import symbols, pi, diff
        >>> from sympy.functions import sqrt, sin
        >>> from sympy.physics.quantum.state import Wavefunction
        >>> x, L = symbols('x,L', positive=True)
        >>> n = symbols('n', integer=True, positive=True)
        >>> g = sqrt(2/L)*sin(n*pi*x/L)
        >>> f = Wavefunction(g, (x, 0, L))
        >>> f.norm
        1
        >>> f(L+1)
        0
        >>> f(L-1)
        sqrt(2)*sin(pi*n*(L - 1)/L)/sqrt(L)
        >>> f(-1)
        0
        >>> f(0.85)
        sqrt(2)*sin(0.85*pi*n/L)/sqrt(L)
        >>> f(0.85, n=1, L=1)
        sqrt(2)*sin(0.85*pi)
        >>> f.is_commutative
        False

    All arguments are automatically sympified, so you can define the variables
    as strings rather than symbols:

        >>> expr = x**2
        >>> f = Wavefunction(expr, 'x')
        >>> type(f.variables[0])
        <class 'sympy.core.symbol.Symbol'>

    Derivatives of Wavefunctions will return Wavefunctions:

        >>> diff(f, x)
        Wavefunction(2*x, x)

    """

    #Any passed tuples for coordinates and their bounds need to be
    #converted to Tuples before Function's constructor is called, to
    #avoid errors from calling is_Float in the constructor
    def __new__(cls, *args, **options):
        new_args = [None for i in args]
        ct = 0
        for arg in args:
            if isinstance(arg, tuple):
                new_args[ct] = Tuple(*arg)
            else:
                new_args[ct] = arg
            ct += 1

        return super(Function, cls).__new__(cls, *new_args, **options)

    def __call__(self, *args, **options):
        var = self.variables

        if len(args) != len(var):
            raise NotImplementedError(
                "Incorrect number of arguments to function!")

        ct = 0
        #If the passed value is outside the specified bounds, return 0
        for v in var:
            lower, upper = self.limits[v]

            #Do the comparison to limits only if the passed symbol is actually
            #a symbol present in the limits;
            #Had problems with a comparison of x > L
            if isinstance(args[ct], Expr) and \
                not (lower in args[ct].free_symbols
                     or upper in args[ct].free_symbols):
                continue

            if (args[ct] < lower) == True or (args[ct] > upper) == True:
                return 0

            ct += 1

        expr = self.expr

        #Allows user to make a call like f(2, 4, m=1, n=1)
        for symbol in list(expr.free_symbols):
            if str(symbol) in options.keys():
                val = options[str(symbol)]
                expr = expr.subs(symbol, val)

        return expr.subs(zip(var, args))

    def _eval_derivative(self, symbol):
        expr = self.expr
        deriv = expr._eval_derivative(symbol)

        return Wavefunction(deriv, *self.args[1:])

    def _eval_conjugate(self):
        return Wavefunction(conjugate(self.expr), *self.args[1:])

    def _eval_transpose(self):
        return self

    @property
    def free_symbols(self):
        return self.expr.free_symbols

    @property
    def is_commutative(self):
        """
        Override Function's is_commutative so that order is preserved in
        represented expressions
        """
        return False

    @classmethod
    def eval(self, *args):
        return None

    @property
    def variables(self):
        """
        Return the coordinates which the wavefunction depends on

        Examples
        ========

            >>> from sympy.physics.quantum.state import Wavefunction
            >>> from sympy import symbols
            >>> x,y = symbols('x,y')
            >>> f = Wavefunction(x*y, x, y)
            >>> f.variables
            (x, y)
            >>> g = Wavefunction(x*y, x)
            >>> g.variables
            (x,)

        """
        var = [g[0] if isinstance(g, Tuple) else g for g in self._args[1:]]
        return tuple(var)

    @property
    def limits(self):
        """
        Return the limits of the coordinates which the w.f. depends on If no
        limits are specified, defaults to ``(-oo, oo)``.

        Examples
        ========

            >>> from sympy.physics.quantum.state import Wavefunction
            >>> from sympy import symbols
            >>> x, y = symbols('x, y')
            >>> f = Wavefunction(x**2, (x, 0, 1))
            >>> f.limits
            {x: (0, 1)}
            >>> f = Wavefunction(x**2, x)
            >>> f.limits
            {x: (-oo, oo)}
            >>> f = Wavefunction(x**2 + y**2, x, (y, -1, 2))
            >>> f.limits
            {x: (-oo, oo), y: (-1, 2)}

        """
        limits = [(g[1], g[2]) if isinstance(g, Tuple) else (-oo, oo)
                  for g in self._args[1:]]
        return dict(zip(self.variables, tuple(limits)))

    @property
    def expr(self):
        """
        Return the expression which is the functional form of the Wavefunction

        Examples
        ========

            >>> from sympy.physics.quantum.state import Wavefunction
            >>> from sympy import symbols
            >>> x, y = symbols('x, y')
            >>> f = Wavefunction(x**2, x)
            >>> f.expr
            x**2

        """
        return self._args[0]

    @property
    def is_normalized(self):
        """
        Returns true if the Wavefunction is properly normalized

        Examples
        ========

            >>> from sympy import symbols, pi
            >>> from sympy.functions import sqrt, sin
            >>> from sympy.physics.quantum.state import Wavefunction
            >>> x, L = symbols('x,L', positive=True)
            >>> n = symbols('n', integer=True, positive=True)
            >>> g = sqrt(2/L)*sin(n*pi*x/L)
            >>> f = Wavefunction(g, (x, 0, L))
            >>> f.is_normalized
            True

        """

        return (self.norm == 1.0)

    @property
    @cacheit
    def norm(self):
        """
        Return the normalization of the specified functional form.

        This function integrates over the coordinates of the Wavefunction, with
        the bounds specified.

        Examples
        ========

            >>> from sympy import symbols, pi
            >>> from sympy.functions import sqrt, sin
            >>> from sympy.physics.quantum.state import Wavefunction
            >>> x, L = symbols('x,L', positive=True)
            >>> n = symbols('n', integer=True, positive=True)
            >>> g = sqrt(2/L)*sin(n*pi*x/L)
            >>> f = Wavefunction(g, (x, 0, L))
            >>> f.norm
            1
            >>> g = sin(n*pi*x/L)
            >>> f = Wavefunction(g, (x, 0, L))
            >>> f.norm
            sqrt(2)*sqrt(L)/2

        """

        exp = self.expr*conjugate(self.expr)
        var = self.variables
        limits = self.limits

        for v in var:
            curr_limits = limits[v]
            exp = integrate(exp, (v, curr_limits[0], curr_limits[1]))

        return sqrt(exp)

    def normalize(self):
        """
        Return a normalized version of the Wavefunction

        Examples
        ========

            >>> from sympy import symbols, pi
            >>> from sympy.functions import sqrt, sin
            >>> from sympy.physics.quantum.state import Wavefunction
            >>> x = symbols('x', real=True)
            >>> L = symbols('L', positive=True)
            >>> n = symbols('n', integer=True, positive=True)
            >>> g = sin(n*pi*x/L)
            >>> f = Wavefunction(g, (x, 0, L))
            >>> f.normalize()
            Wavefunction(sqrt(2)*sin(pi*n*x/L)/sqrt(L), (x, 0, L))

        """
        const = self.norm

        if const == oo:
            raise NotImplementedError("The function is not normalizable!")
        else:
            return Wavefunction((const)**(-1)*self.expr, *self.args[1:])

    def prob(self):
        """
        Return the absolute magnitude of the w.f., `|\psi(x)|^2`

        Examples
        ========

            >>> from sympy import symbols, pi
            >>> from sympy.functions import sqrt, sin
            >>> from sympy.physics.quantum.state import Wavefunction
            >>> x, L = symbols('x,L', real=True)
            >>> n = symbols('n', integer=True)
            >>> g = sin(n*pi*x/L)
            >>> f = Wavefunction(g, (x, 0, L))
            >>> f.prob()
            Wavefunction(sin(pi*n*x/L)**2, x)

        """

        return Wavefunction(self.expr*conjugate(self.expr), *self.variables)


from __future__ import print_function, division
import matplotlib.pyplot as plt
import numpy as np
import scipy
import pandas as pd
import seaborn as sns
from sklearn.datasets import load_boston
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import KNeighborsClassifier
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis, QuadraticDiscriminantAnalysis
from sklearn.metrics import confusion_matrix, accuracy_score, precision_score, recall_score
from pandas.tools.plotting import scatter_matrix
import statsmodels.formula.api as smf
import statsmodels.api as sm

#Load boston dataset from sklearn#
boston = load_boston()

#Columns#
#print(boston['feature_names'])
#Descriptio#
#print(boston['DESCR'])

rawdata = pd.DataFrame(boston.data, columns=boston.feature_names)
rawdata['MEDV'] = boston.target

#Convert to NaN#
data = rawdata.replace(to_replace='None', value=np.nan).copy()

#Create the binary variable CRIM01#
data['CRIM01'] = np.where(data['CRIM'] > data['CRIM'].median(), 1, 0)

#Columns#
numcols = list(data.columns)
numcols.remove('CRIM')

#Predictor without target columns#
xcols = list(numcols)
xcols.remove('CRIM01')

#Summary (mean, stdev, range, etc)#
print('\nFull data summary')
print(data.describe())

#Correlations#
print('\nData correlations')
dcorrs = data.corr()
print(dcorrs)

#Pair plot matrix#
#sns.set()
#sns.pairplot(data[numcols], hue='CRIM01')

print('\n\n### LOGISTIC REGRESSION###')

## Logistic regression with statsmodels ##
plr_form = 'CRIM01~' + '+'.join(xcols)
prelogreg = smf.glm(formula=plr_form, data=data, family=sm.families.Binomial()).fit()

#Remove predictors with high P-values from LogReg#
logit_pvals = prelogreg.pvalues
pred_keep = list(logit_pvals[logit_pvals < 0.05].index)
pred_keep.remove('Intercept')
print('\nAfter first LogReg iteration, keeping only: ', pred_keep)

# New LogReg with only low p-value predictors#
lr_form = 'CRIM01~' + '+'.join(pred_keep)
logreg = smf.glm(formula=lr_form, data=data, family=sm.families.Binomial()).fit()

print('\nLogistic regression fit summary')
print(logreg.summary())

#Splitting the data for train/test#
X_data = data[pred_keep]
Y_data = data['CRIM01']

X_train, X_test, Y_train, Y_test = train_test_split(X_data, Y_data, test_size=0.5, random_state=42, stratify=Y_data)

# Initiate logistic regression object
logit_clf = LogisticRegression()

# Fit model. Let X_train = matrix of predictors, Y_train = matrix of variables.
resLogit_clf = logit_clf.fit(X_train, Y_train)

#Predicted values for training set
Y_pred_logit = resLogit_clf.predict(X_test)

#Confusion matrix#
print("\nConfusion matrix logit:")
print(confusion_matrix(Y_test, Y_pred_logit))

#Accuracy, precision and recall#
print('\nAccuracy logit:', np.round(accuracy_score(Y_test, Y_pred_logit), 3))
print("Precision logit:", np.round(precision_score(Y_test, Y_pred_logit, pos_label=1), 3))
print("Recall logit:", np.round(recall_score(Y_test, Y_pred_logit, pos_label=1), 3))

print('\n\n### LINEAR DISCRIMINANT ANALYSIS ###')
# Initiate logistic regression object
lda_clf = LinearDiscriminantAnalysis()

# Fit model. Let X_train = matrix of new_pred, Y_train = matrix of variables.
reslda_clf = lda_clf.fit(X_train, Y_train)

#Predicted values for training set
Y_pred_lda = reslda_clf.predict(X_test)

#Prior probabilities#
print("\nPrior probabilities")
print(reslda_clf.classes_)
print(reslda_clf.priors_)

#Group means#
print("\nGroup means")
#print(reslda_clf.classes_)
print(reslda_clf.means_)

#Coefficients#
print("\nCoefficients")
#print(reslda_clf.classes_)
print(reslda_clf.intercept_)
print(reslda_clf.coef_)

#Confusion matrix#
print("\nConfusion matrix LDA:")
print(confusion_matrix(Y_test, Y_pred_lda))

#Accuracy, precision and recall#
print("\nAccuracy LDA:", np.round(accuracy_score(Y_test, Y_pred_lda), 3))
print("Precision LDA:", np.round(precision_score(Y_test, Y_pred_lda, pos_label=1), 3))
print("Recall LDA:", np.round(recall_score(Y_test, Y_pred_lda, pos_label=1), 3))

print('\n\n### QUADRATIC DISCRIMINANT ANALYSIS ###')
# Initiate QDA object
qda_clf = QuadraticDiscriminantAnalysis()

# Fit model. Let X_train = matrix of new_pred, Y_train = matrix of variables.
resqda_clf = qda_clf.fit(X_train, Y_train)

#Predicted values for training set
Y_pred_qda = resqda_clf.predict(X_test)

#Prior probabilities#
print("\nPrior probabilities")
print(resqda_clf.classes_)
print(resqda_clf.priors_)

#Group means#
print("\nGroup means")
#print(resqda_clf.classes_)
print(resqda_clf.means_)

#Confusion matrix#
print("\nConfusion matrix QDA:")
print(confusion_matrix(Y_test, Y_pred_qda))

#Accuracy, precision and recall#
print("\nAccuracy QDA:", np.round(accuracy_score(Y_test, Y_pred_qda), 3))
print("Precision QDA:", np.round(precision_score(Y_test, Y_pred_qda, pos_label=1), 3))
print("Recall QDA:", np.round(recall_score(Y_test, Y_pred_qda, pos_label=1), 3))

print('\n\n### K NEAREST NEIGHBORS ###')
#K value#
kval = 15
print('\nUsing k = ' + str(kval))

# Initiate KNN object
knn_clf = KNeighborsClassifier(n_neighbors=15)

# Fit model. Let X_train = matrix of new_pred, Y_train = matrix of variables.
resknn_clf = knn_clf.fit(X_train, Y_train)

#Predicted values for training set
Y_pred_knn = resknn_clf.predict(X_test)

#Confusion matrix#
print("\nConfusion matrix KNN:")
print(confusion_matrix(Y_test, Y_pred_knn))

#Accuracy, precision and recall#
print("\nAccuracy KNN:", np.round(accuracy_score(Y_test, Y_pred_knn), 3))
print("Precision KNN:", np.round(precision_score(Y_test, Y_pred_knn, pos_label=1), 3))
print("Recall KNN:", np.round(recall_score(Y_test, Y_pred_knn, pos_label=1), 3))

#plt.show()


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

Created on Sun Jan 06 09:50:54 2013

Author: Josef Perktold
"""

from __future__ import print_function

if __name__ == '__main__':

    import numpy as np
    import matplotlib.pyplot as plt
    from statsmodels.nonparametric.api import KernelReg
    import statsmodels.sandbox.nonparametric.dgp_examples as dgp

    seed = np.random.randint(999999)
    seed = 430973
    print(seed)
    np.random.seed(seed)

    funcs = [dgp.UnivariateFanGijbels1(),
             dgp.UnivariateFanGijbels2(),
             dgp.UnivariateFanGijbels1EU(),
             #dgp.UnivariateFanGijbels2(distr_x=stats.uniform(-2, 4))
             dgp.UnivariateFunc1()
             ]

    res = []
    fig = plt.figure()
    for i,func in enumerate(funcs):
        #f = func()
        f = func
        model = KernelReg(endog=[f.y], exog=[f.x], reg_type='ll',
                          var_type='c', bw='cv_ls')
        mean, mfx = model.fit()
        ax = fig.add_subplot(2, 2, i+1)
        f.plot(ax=ax)
        ax.plot(f.x, mean, color='r', lw=2, label='est. mean')
        ax.legend(loc='upper left')
        res.append((model, mean, mfx))

    fig.suptitle('Kernel Regression')
    fig.show()


import numpy as np

from bokeh.palettes import Spectral6

def process_data():
    from bokeh.sampledata.gapminder import fertility, life_expectancy, population, regions

    # Make the column names ints not strings for handling
    columns = list(fertility.columns)
    years = list(range(int(columns[0]), int(columns[-1])))
    rename_dict = dict(zip(columns, years))

    fertility = fertility.rename(columns=rename_dict)
    life_expectancy = life_expectancy.rename(columns=rename_dict)
    population = population.rename(columns=rename_dict)
    regions = regions.rename(columns=rename_dict)

    # Turn population into bubble sizes. Use min_size and factor to tweak.
    scale_factor = 200
    population_size = np.sqrt(population / np.pi) / scale_factor
    min_size = 3
    population_size = population_size.where(population_size >= min_size).fillna(min_size)

    # Use pandas categories and categorize & color the regions
    regions.Group = regions.Group.astype('category')
    regions_list = list(regions.Group.cat.categories)

    def get_color(r):
        return Spectral6[regions_list.index(r.Group)]
    regions['region_color'] = regions.apply(get_color, axis=1)

    return fertility, life_expectancy, population_size, regions, years, regions_list