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

import os
import sys
from tarfile import TarFile, TarInfo
from matplotlib import pylab as pl
import numpy as n
import getopt as opt
from python_util.util import *
from math import sqrt, ceil, floor
from python_util.gpumodel import IGPUModel
import random as r
import numpy.random as nr
from convnet import ConvNet
from python_util.options import *
from PIL import Image
from time import sleep

class ShowNetError(Exception):
    pass

class ShowConvNet(ConvNet):
    def __init__(self, op, load_dic):
        ConvNet.__init__(self, op, load_dic)

    def init_data_providers(self):
        self.need_gpu = self.op.get_value('show_preds')
        class Dummy:
            def advance_batch(self):
                pass
        if self.need_gpu:
            ConvNet.init_data_providers(self)
        else:
            self.train_data_provider = self.test_data_provider = Dummy()

    def import_model(self):
        if self.need_gpu:
            ConvNet.import_model(self)

    def init_model_state(self):
        if self.op.get_value('show_preds'):
            self.softmax_name = self.op.get_value('show_preds')

    def init_model_lib(self):
        if self.need_gpu:
            ConvNet.init_model_lib(self)

    def plot_cost(self):
        if self.show_cost not in self.train_outputs[0][0]:
            raise ShowNetError("Cost function with name '%s' not defined by given convnet." % self.show_cost)
#        print self.test_outputs
        train_errors = [eval(self.layers[self.show_cost]['outputFilter'])(o[0][self.show_cost], o[1])[self.cost_idx] for o in self.train_outputs]
        test_errors = [eval(self.layers[self.show_cost]['outputFilter'])(o[0][self.show_cost], o[1])[self.cost_idx] for o in self.test_outputs]
        if self.smooth_test_errors:
            test_errors = [sum(test_errors[max(0,i-len(self.test_batch_range)):i])/(i-max(0,i-len(self.test_batch_range))) for i in xrange(1,len(test_errors)+1)]
        numbatches = len(self.train_batch_range)
        test_errors = n.row_stack(test_errors)
        test_errors = n.tile(test_errors, (1, self.testing_freq))
        test_errors = list(test_errors.flatten())
        test_errors += [test_errors[-1]] * max(0,len(train_errors) - len(test_errors))
        test_errors = test_errors[:len(train_errors)]

        numepochs = len(train_errors) / float(numbatches)
        pl.figure(1)
        x = range(0, len(train_errors))
        pl.plot(x, train_errors, 'k-', label='Training set')
        pl.plot(x, test_errors, 'r-', label='Test set')
        pl.legend()
        ticklocs = range(numbatches, len(train_errors) - len(train_errors) % numbatches + 1, numbatches)
        epoch_label_gran = int(ceil(numepochs / 20.))
        epoch_label_gran = int(ceil(float(epoch_label_gran) / 10) * 10) if numepochs >= 10 else epoch_label_gran
        ticklabels = map(lambda x: str((x[1] / numbatches)) if x[0] % epoch_label_gran == epoch_label_gran-1 else '', enumerate(ticklocs))

        pl.xticks(ticklocs, ticklabels)
        pl.xlabel('Epoch')
#        pl.ylabel(self.show_cost)
        pl.title('%s[%d]' % (self.show_cost, self.cost_idx))
#        print "plotted cost"

    def make_filter_fig(self, filters, filter_start, fignum, _title, num_filters, combine_chans, FILTERS_PER_ROW=16):
        MAX_ROWS = 24
        MAX_FILTERS = FILTERS_PER_ROW * MAX_ROWS
        num_colors = filters.shape[0]
        f_per_row = int(ceil(FILTERS_PER_ROW / float(1 if combine_chans else num_colors)))
        filter_end = min(filter_start+MAX_FILTERS, num_filters)
        filter_rows = int(ceil(float(filter_end - filter_start) / f_per_row))

        filter_pixels = filters.shape[1]
        filter_size = int(sqrt(filters.shape[1]))
        fig = pl.figure(fignum)
        fig.text(.5, .95, '%s %dx%d filters %d-%d' % (_title, filter_size, filter_size, filter_start, filter_end-1), horizontalalignment='center')
        num_filters = filter_end - filter_start
        if not combine_chans:
            bigpic = n.zeros((filter_size * filter_rows + filter_rows + 1, filter_size*num_colors * f_per_row + f_per_row + 1), dtype=n.single)
        else:
            bigpic = n.zeros((3, filter_size * filter_rows + filter_rows + 1, filter_size * f_per_row + f_per_row + 1), dtype=n.single)

        for m in xrange(filter_start,filter_end ):
            filter = filters[:,:,m]
            y, x = (m - filter_start) / f_per_row, (m - filter_start) % f_per_row
            if not combine_chans:
                for c in xrange(num_colors):
                    filter_pic = filter[c,:].reshape((filter_size,filter_size))
                    bigpic[1 + (1 + filter_size) * y:1 + (1 + filter_size) * y + filter_size,
                           1 + (1 + filter_size*num_colors) * x + filter_size*c:1 + (1 + filter_size*num_colors) * x + filter_size*(c+1)] = filter_pic
            else:
                filter_pic = filter.reshape((3, filter_size,filter_size))
                bigpic[:,
                       1 + (1 + filter_size) * y:1 + (1 + filter_size) * y + filter_size,
                       1 + (1 + filter_size) * x:1 + (1 + filter_size) * x + filter_size] = filter_pic

        pl.xticks([])
        pl.yticks([])
        if not combine_chans:
            pl.imshow(bigpic, cmap=pl.cm.gray, interpolation='nearest')
        else:
            bigpic = bigpic.swapaxes(0,2).swapaxes(0,1)
            pl.imshow(bigpic, interpolation='nearest')

    def plot_filters(self):
        FILTERS_PER_ROW = 16
        filter_start = 0 # First filter to show
        if self.show_filters not in self.layers:
            raise ShowNetError("Layer with name '%s' not defined by given convnet." % self.show_filters)
        layer = self.layers[self.show_filters]
        filters = layer['weights'][self.input_idx]
#        filters = filters - filters.min()
#        filters = filters / filters.max()
        if layer['type'] == 'fc': # Fully-connected layer
            num_filters = layer['outputs']
            channels = self.channels
            filters = filters.reshape(channels, filters.shape[0]/channels, filters.shape[1])
        elif layer['type'] in ('conv', 'local'): # Conv layer
            num_filters = layer['filters']
            channels = layer['filterChannels'][self.input_idx]
            if layer['type'] == 'local':
                filters = filters.reshape((layer['modules'], channels, layer['filterPixels'][self.input_idx], num_filters))
                filters = filters[:, :, :, self.local_plane] # first map for now (modules, channels, pixels)
                filters = filters.swapaxes(0,2).swapaxes(0,1)
                num_filters = layer['modules']
#                filters = filters.swapaxes(0,1).reshape(channels * layer['filterPixels'][self.input_idx], num_filters * layer['modules'])
#                num_filters *= layer['modules']
                FILTERS_PER_ROW = layer['modulesX']
            else:
                filters = filters.reshape(channels, filters.shape[0]/channels, filters.shape[1])

        # Convert YUV filters to RGB
        if self.yuv_to_rgb and channels == 3:
            R = filters[0,:,:] + 1.28033 * filters[2,:,:]
            G = filters[0,:,:] + -0.21482 * filters[1,:,:] + -0.38059 * filters[2,:,:]
            B = filters[0,:,:] + 2.12798 * filters[1,:,:]
            filters[0,:,:], filters[1,:,:], filters[2,:,:] = R, G, B
        combine_chans = not self.no_rgb and channels == 3

        # Make sure you don't modify the backing array itself here -- so no -= or /=
        if self.norm_filters:
            #print filters.shape
            filters = filters - n.tile(filters.reshape((filters.shape[0] * filters.shape[1], filters.shape[2])).mean(axis=0).reshape(1, 1, filters.shape[2]), (filters.shape[0], filters.shape[1], 1))
            filters = filters / n.sqrt(n.tile(filters.reshape((filters.shape[0] * filters.shape[1], filters.shape[2])).var(axis=0).reshape(1, 1, filters.shape[2]), (filters.shape[0], filters.shape[1], 1)))
            #filters = filters - n.tile(filters.min(axis=0).min(axis=0), (3, filters.shape[1], 1))
            #filters = filters / n.tile(filters.max(axis=0).max(axis=0), (3, filters.shape[1], 1))
        #else:
        filters = filters - filters.min()
        filters = filters / filters.max()

        self.make_filter_fig(filters, filter_start, 2, 'Layer %s' % self.show_filters, num_filters, combine_chans, FILTERS_PER_ROW=FILTERS_PER_ROW)

    def plot_predictions(self):
        epoch, batch, data = self.get_next_batch(train=False) # get a test batch
        num_classes = self.test_data_provider.get_num_classes()
        NUM_ROWS = 2
        NUM_COLS = 4
        NUM_IMGS = NUM_ROWS * NUM_COLS if not self.save_preds else data[0].shape[1]
        NUM_TOP_CLASSES = min(num_classes, 5) # show this many top labels
        NUM_OUTPUTS = self.model_state['layers'][self.softmax_name]['outputs']
        PRED_IDX = 1

        label_names = [lab.split(',')[0] for lab in self.test_data_provider.batch_meta['label_names']]
        if self.only_errors:
            preds = n.zeros((data[0].shape[1], NUM_OUTPUTS), dtype=n.single)
        else:
            preds = n.zeros((NUM_IMGS, NUM_OUTPUTS), dtype=n.single)
            #rand_idx = nr.permutation(n.r_[n.arange(1), n.where(data[1] == 552)[1], n.where(data[1] == 795)[1], n.where(data[1] == 449)[1], n.where(data[1] == 274)[1]])[:NUM_IMGS]
            rand_idx = nr.randint(0, data[0].shape[1], NUM_IMGS)
            if NUM_IMGS < data[0].shape[1]:
                data = [n.require(d[:,rand_idx], requirements='C') for d in data]
#        data += [preds]
        # Run the model
        print  [d.shape for d in data], preds.shape
        self.libmodel.startFeatureWriter(data, [preds], [self.softmax_name])
        IGPUModel.finish_batch(self)
        print preds
        data[0] = self.test_data_provider.get_plottable_data(data[0])

        if self.save_preds:
            if not gfile.Exists(self.save_preds):
                gfile.MakeDirs(self.save_preds)
            preds_thresh = preds > 0.5 # Binarize predictions
            data[0] = data[0] * 255.0
            data[0][data[0]<0] = 0
            data[0][data[0]>255] = 255
            data[0] = n.require(data[0], dtype=n.uint8)
            dir_name = '%s_predictions_batch_%d' % (os.path.basename(self.save_file), batch)
            tar_name = os.path.join(self.save_preds, '%s.tar' % dir_name)
            tfo = gfile.GFile(tar_name, "w")
            tf = TarFile(fileobj=tfo, mode='w')
            for img_idx in xrange(NUM_IMGS):
                img = data[0][img_idx,:,:,:]
                imsave = Image.fromarray(img)
                prefix = "CORRECT" if data[1][0,img_idx] == preds_thresh[img_idx,PRED_IDX] else "FALSE_POS" if preds_thresh[img_idx,PRED_IDX] == 1 else "FALSE_NEG"
                file_name = "%s_%.2f_%d_%05d_%d.png" % (prefix, preds[img_idx,PRED_IDX], batch, img_idx, data[1][0,img_idx])
#                gf = gfile.GFile(file_name, "w")
                file_string = StringIO()
                imsave.save(file_string, "PNG")
                tarinf = TarInfo(os.path.join(dir_name, file_name))
                tarinf.size = file_string.tell()
                file_string.seek(0)
                tf.addfile(tarinf, file_string)
            tf.close()
            tfo.close()
#                gf.close()
            print "Wrote %d prediction PNGs to %s" % (preds.shape[0], tar_name)
        else:
            fig = pl.figure(3, figsize=(12,9))
            fig.text(.4, .95, '%s test samples' % ('Mistaken' if self.only_errors else 'Random'))
            if self.only_errors:
                # what the net got wrong
                if NUM_OUTPUTS > 1:
                    err_idx = [i for i,p in enumerate(preds.argmax(axis=1)) if p not in n.where(data[2][:,i] > 0)[0]]
                else:
                    err_idx = n.where(data[1][0,:] != preds[:,0].T)[0]
                    print err_idx
                err_idx = r.sample(err_idx, min(len(err_idx), NUM_IMGS))
                data[0], data[1], preds = data[0][:,err_idx], data[1][:,err_idx], preds[err_idx,:]

            import matplotlib.gridspec as gridspec
            import matplotlib.colors as colors
            cconv = colors.ColorConverter()
            gs = gridspec.GridSpec(NUM_ROWS*2, NUM_COLS,
                                   width_ratios=[1]*NUM_COLS, height_ratios=[2,1]*NUM_ROWS )
            #print data[1]
            for row in xrange(NUM_ROWS):
                for col in xrange(NUM_COLS):
                    img_idx = row * NUM_COLS + col
                    if data[0].shape[0] <= img_idx:
                        break
                    pl.subplot(gs[(row * 2) * NUM_COLS + col])
                    #pl.subplot(NUM_ROWS*2, NUM_COLS, row * 2 * NUM_COLS + col + 1)
                    pl.xticks([])
                    pl.yticks([])
                    img = data[0][img_idx,:,:,:]
                    pl.imshow(img, interpolation='lanczos')
                    show_title = data[1].shape[0] == 1
                    true_label = [int(data[1][0,img_idx])] if show_title else n.where(data[1][:,img_idx]==1)[0]
                    #print true_label
                    #print preds[img_idx,:].shape
                    #print preds[img_idx,:].max()
                    true_label_names = [label_names[i] for i in true_label]
                    img_labels = sorted(zip(preds[img_idx,:], label_names), key=lambda x: x[0])[-NUM_TOP_CLASSES:]
                    #print img_labels
                    axes = pl.subplot(gs[(row * 2 + 1) * NUM_COLS + col])
                    height = 0.5
                    ylocs = n.array(range(NUM_TOP_CLASSES))*height
                    pl.barh(ylocs, [l[0] for l in img_labels], height=height, \
                            color=['#ffaaaa' if l[1] in true_label_names else '#aaaaff' for l in img_labels])
                    #pl.title(", ".join(true_labels))
                    if show_title:
                        pl.title(", ".join(true_label_names), fontsize=15, fontweight='bold')
                    else:
                        print true_label_names
                    pl.yticks(ylocs + height/2, [l[1] for l in img_labels], x=1, backgroundcolor=cconv.to_rgba('0.65', alpha=0.5), weight='bold')
                    for line in enumerate(axes.get_yticklines()):
                        line[1].set_visible(False)
                    #pl.xticks([width], [''])
                    #pl.yticks([])
                    pl.xticks([])
                    pl.ylim(0, ylocs[-1] + height)
                    pl.xlim(0, 1)

    def start(self):
        self.op.print_values()
#        print self.show_cost
        if self.show_cost:
            self.plot_cost()
        if self.show_filters:
            self.plot_filters()
        if self.show_preds:
            self.plot_predictions()

        if pl:
            pl.show()
        sys.exit(0)

    @classmethod
    def get_options_parser(cls):
        op = ConvNet.get_options_parser()
        for option in list(op.options):
            if option not in ('gpu', 'load_file', 'inner_size', 'train_batch_range', 'test_batch_range', 'multiview_test', 'data_path', 'pca_noise', 'scalar_mean'):
                op.delete_option(option)
        op.add_option("show-cost", "show_cost", StringOptionParser, "Show specified objective function", default="")
        op.add_option("show-filters", "show_filters", StringOptionParser, "Show learned filters in specified layer", default="")
        op.add_option("norm-filters", "norm_filters", BooleanOptionParser, "Individually normalize filters shown with --show-filters", default=0)
        op.add_option("input-idx", "input_idx", IntegerOptionParser, "Input index for layer given to --show-filters", default=0)
        op.add_option("cost-idx", "cost_idx", IntegerOptionParser, "Cost function return value index for --show-cost", default=0)
        op.add_option("no-rgb", "no_rgb", BooleanOptionParser, "Don't combine filter channels into RGB in layer given to --show-filters", default=False)
        op.add_option("yuv-to-rgb", "yuv_to_rgb", BooleanOptionParser, "Convert RGB filters to YUV in layer given to --show-filters", default=False)
        op.add_option("channels", "channels", IntegerOptionParser, "Number of channels in layer given to --show-filters (fully-connected layers only)", default=0)
        op.add_option("show-preds", "show_preds", StringOptionParser, "Show predictions made by given softmax on test set", default="")
        op.add_option("save-preds", "save_preds", StringOptionParser, "Save predictions to given path instead of showing them", default="")
        op.add_option("only-errors", "only_errors", BooleanOptionParser, "Show only mistaken predictions (to be used with --show-preds)", default=False, requires=['show_preds'])
        op.add_option("local-plane", "local_plane", IntegerOptionParser, "Local plane to show", default=0)
        op.add_option("smooth-test-errors", "smooth_test_errors", BooleanOptionParser, "Use running average for test error plot?", default=1)

        op.options['load_file'].default = None
        return op

if __name__ == "__main__":
    #nr.seed(6)
    try:
        op = ShowConvNet.get_options_parser()
        op, load_dic = IGPUModel.parse_options(op)
        model = ShowConvNet(op, load_dic)
        model.start()
    except (UnpickleError, ShowNetError, opt.GetoptError), e:
        print "----------------"
        print "Error:"
        print e


"""
================================
Gaussian Mixture Model Selection
================================

This example shows that model selection can be performed with
Gaussian Mixture Models using information-theoretic criteria (BIC).
Model selection concerns both the covariance type
and the number of components in the model.
In that case, AIC also provides the right result (not shown to save time),
but BIC is better suited if the problem is to identify the right model.
Unlike Bayesian procedures, such inferences are prior-free.

In that case, the model with 2 components and full covariance
(which corresponds to the true generative model) is selected.
"""

import numpy as np
import itertools

from scipy import linalg
import matplotlib.pyplot as plt
import matplotlib as mpl

from sklearn import mixture

print(__doc__)

# Number of samples per component
n_samples = 500

# Generate random sample, two components
np.random.seed(0)
C = np.array([[0., -0.1], [1.7, .4]])
X = np.r_[np.dot(np.random.randn(n_samples, 2), C),
          .7 * np.random.randn(n_samples, 2) + np.array([-6, 3])]

lowest_bic = np.infty
bic = []
n_components_range = range(1, 7)
cv_types = ['spherical', 'tied', 'diag', 'full']
for cv_type in cv_types:
    for n_components in n_components_range:
        # Fit a Gaussian mixture with EM
        gmm = mixture.GaussianMixture(n_components=n_components,
                                      covariance_type=cv_type)
        gmm.fit(X)
        bic.append(gmm.bic(X))
        if bic[-1] < lowest_bic:
            lowest_bic = bic[-1]
            best_gmm = gmm

bic = np.array(bic)
color_iter = itertools.cycle(['navy', 'turquoise', 'cornflowerblue',
                              'darkorange'])
clf = best_gmm
bars = []

# Plot the BIC scores
spl = plt.subplot(2, 1, 1)
for i, (cv_type, color) in enumerate(zip(cv_types, color_iter)):
    xpos = np.array(n_components_range) + .2 * (i - 2)
    bars.append(plt.bar(xpos, bic[i * len(n_components_range):
                                  (i + 1) * len(n_components_range)],
                        width=.2, color=color))
plt.xticks(n_components_range)
plt.ylim([bic.min() * 1.01 - .01 * bic.max(), bic.max()])
plt.title('BIC score per model')
xpos = np.mod(bic.argmin(), len(n_components_range)) + .65 +\
    .2 * np.floor(bic.argmin() / len(n_components_range))
plt.text(xpos, bic.min() * 0.97 + .03 * bic.max(), '*', fontsize=14)
spl.set_xlabel('Number of components')
spl.legend([b[0] for b in bars], cv_types)

# Plot the winner
splot = plt.subplot(2, 1, 2)
Y_ = clf.predict(X)
for i, (mean, cov, color) in enumerate(zip(clf.means_, clf.covariances_,
                                           color_iter)):
    v, w = linalg.eigh(cov)
    if not np.any(Y_ == i):
        continue
    plt.scatter(X[Y_ == i, 0], X[Y_ == i, 1], .8, color=color)

    # Plot an ellipse to show the Gaussian component
    angle = np.arctan2(w[0][1], w[0][0])
    angle = 180. * angle / np.pi  # convert to degrees
    v = 2. * np.sqrt(2.) * np.sqrt(v)
    ell = mpl.patches.Ellipse(mean, v[0], v[1], 180. + angle, color=color)
    ell.set_clip_box(splot.bbox)
    ell.set_alpha(.5)
    splot.add_artist(ell)

plt.xticks(())
plt.yticks(())
plt.title('Selected GMM: full model, 2 components')
plt.subplots_adjust(hspace=.35, bottom=.02)
plt.show()


# -*- coding: utf-8 -*-
"""
==================================
Color Quantization using K-Means
==================================

Performs a pixel-wise Vector Quantization (VQ) of an image of the summer palace
(China), reducing the number of colors required to show the image from 96,615
unique colors to 64, while preserving the overall appearance quality.

In this example, pixels are represented in a 3D-space and K-means is used to
find 64 color clusters. In the image processing literature, the codebook
obtained from K-means (the cluster centers) is called the color palette. Using
a single byte, up to 256 colors can be addressed, whereas an RGB encoding
requires 3 bytes per pixel. The GIF file format, for example, uses such a
palette.

For comparison, a quantized image using a random codebook (colors picked up
randomly) is also shown.
"""
# Authors: Robert Layton <robertlayton@gmail.com>
#          Olivier Grisel <olivier.grisel@ensta.org>
#          Mathieu Blondel <mathieu@mblondel.org>
#
# License: BSD 3 clause

print(__doc__)
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.metrics import pairwise_distances_argmin
from sklearn.datasets import load_sample_image
from sklearn.utils import shuffle
from time import time

n_colors = 64

# Load the Summer Palace photo
china = load_sample_image("china.jpg")

# Convert to floats instead of the default 8 bits integer coding. Dividing by
# 255 is important so that plt.imshow behaves works well on float data (need to
# be in the range [0-1])
china = np.array(china, dtype=np.float64) / 255

# Load Image and transform to a 2D numpy array.
w, h, d = original_shape = tuple(china.shape)
assert d == 3
image_array = np.reshape(china, (w * h, d))

print("Fitting model on a small sub-sample of the data")
t0 = time()
image_array_sample = shuffle(image_array, random_state=0)[:1000]
kmeans = KMeans(n_clusters=n_colors, random_state=0).fit(image_array_sample)
print("done in %0.3fs." % (time() - t0))

# Get labels for all points
print("Predicting color indices on the full image (k-means)")
t0 = time()
labels = kmeans.predict(image_array)
print("done in %0.3fs." % (time() - t0))

codebook_random = shuffle(image_array, random_state=0)[:n_colors + 1]
print("Predicting color indices on the full image (random)")
t0 = time()
labels_random = pairwise_distances_argmin(codebook_random,
                                          image_array,
                                          axis=0)
print("done in %0.3fs." % (time() - t0))

def recreate_image(codebook, labels, w, h):
    """Recreate the (compressed) image from the code book & labels"""
    d = codebook.shape[1]
    image = np.zeros((w, h, d))
    label_idx = 0
    for i in range(w):
        for j in range(h):
            image[i][j] = codebook[labels[label_idx]]
            label_idx += 1
    return image

# Display all results, alongside original image
plt.figure(1)
plt.clf()
ax = plt.axes([0, 0, 1, 1])
plt.axis('off')
plt.title('Original image (96,615 colors)')
plt.imshow(china)

plt.figure(2)
plt.clf()
ax = plt.axes([0, 0, 1, 1])
plt.axis('off')
plt.title('Quantized image (64 colors, K-Means)')
plt.imshow(recreate_image(kmeans.cluster_centers_, labels, w, h))

plt.figure(3)
plt.clf()
ax = plt.axes([0, 0, 1, 1])
plt.axis('off')
plt.title('Quantized image (64 colors, Random)')
plt.imshow(recreate_image(codebook_random, labels_random, w, h))
plt.show()


#!/usr/bin/env python
#
# Copyright 2009 Free Software Foundation, Inc.
#
# This file is part of GNU Radio
#
# GNU Radio is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3, or (at your option)
# any later version.
#
# GNU Radio is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with GNU Radio; see the file COPYING.  If not, write to
# the Free Software Foundation, Inc., 51 Franklin Street,
# Boston, MA 02110-1301, USA.
#

from gnuradio import gr, blks2
import sys, math, time

try:
    import scipy
    from scipy import fftpack
except ImportError:
    print "Error: Program requires scipy (see: www.scipy.org)."
    sys.exit(1)

try:
    import pylab
except ImportError:
    print "Error: Program requires matplotlib (see: matplotlib.sourceforge.net)."
    sys.exit(1)

class fmtx(gr.hier_block2):
    def __init__(self, lo_freq, audio_rate, if_rate):

        gr.hier_block2.__init__(self, "build_fm",
                                gr.io_signature(1, 1, gr.sizeof_float),      # Input signature
                                gr.io_signature(1, 1, gr.sizeof_gr_complex)) # Output signature

        fmtx = blks2.nbfm_tx (audio_rate, if_rate, max_dev=5e3, tau=75e-6)

        # Local oscillator
        lo = gr.sig_source_c (if_rate,        # sample rate
                              gr.GR_SIN_WAVE, # waveform type
                              lo_freq,        #frequency
                              1.0,            # amplitude
                              0)              # DC Offset
        mixer = gr.multiply_cc ()

        self.connect (self, fmtx, (mixer, 0))
        self.connect (lo, (mixer, 1))
        self.connect (mixer, self)

class fmtest(gr.top_block):
    def __init__(self):
        gr.top_block.__init__(self)

        self._nsamples = 1000000
        self._audio_rate = 8000

        # Set up N channels with their own baseband and IF frequencies
        self._N = 5
        chspacing = 16000
        freq = [10, 20, 30, 40, 50]
        f_lo = [0, 1*chspacing, -1*chspacing, 2*chspacing, -2*chspacing]

        self._if_rate = 4*self._N*self._audio_rate

        # Create a signal source and frequency modulate it
        self.sum = gr.add_cc ()
        for n in xrange(self._N):
            sig = gr.sig_source_f(self._audio_rate, gr.GR_SIN_WAVE, freq[n], 0.5)
            fm = fmtx(f_lo[n], self._audio_rate, self._if_rate)
            self.connect(sig, fm)
            self.connect(fm, (self.sum, n))

        self.head = gr.head(gr.sizeof_gr_complex, self._nsamples)
        self.snk_tx = gr.vector_sink_c()
        self.channel = blks2.channel_model(0.1)

        self.connect(self.sum, self.head, self.channel, self.snk_tx)

        # Design the channlizer
        self._M = 10
        bw = chspacing/2.0
        t_bw = chspacing/10.0
        self._chan_rate = self._if_rate / self._M
        self._taps = gr.firdes.low_pass_2(1, self._if_rate, bw, t_bw,
                                          attenuation_dB=100,
                                          window=gr.firdes.WIN_BLACKMAN_hARRIS)
        tpc = math.ceil(float(len(self._taps)) /  float(self._M))

        print "Number of taps:     ", len(self._taps)
        print "Number of channels: ", self._M
        print "Taps per channel:   ", tpc

        self.pfb = blks2.pfb_channelizer_ccf(self._M, self._taps)

        self.connect(self.channel, self.pfb)

        # Create a file sink for each of M output channels of the filter and connect it
        self.fmdet = list()
        self.squelch = list()
        self.snks = list()
        for i in xrange(self._M):
            self.fmdet.append(blks2.nbfm_rx(self._audio_rate, self._chan_rate))
            self.squelch.append(blks2.standard_squelch(self._audio_rate*10))
            self.snks.append(gr.vector_sink_f())
            self.connect((self.pfb, i), self.fmdet[i], self.squelch[i], self.snks[i])

    def num_tx_channels(self):
        return self._N

    def num_rx_channels(self):
        return self._M

def main():

    fm = fmtest()

    tstart = time.time()
    fm.run()
    tend = time.time()

    if 1:
        fig1 = pylab.figure(1, figsize=(12,10), facecolor="w")
        fig2 = pylab.figure(2, figsize=(12,10), facecolor="w")
        fig3 = pylab.figure(3, figsize=(12,10), facecolor="w")

        Ns = 10000
        Ne = 100000

        fftlen = 8192
        winfunc = scipy.blackman

        # Plot transmitted signal
        fs = fm._if_rate

        d = fm.snk_tx.data()[Ns:Ns+Ne]
        sp1_f = fig1.add_subplot(2, 1, 1)

        X,freq = sp1_f.psd(d, NFFT=fftlen, noverlap=fftlen/4, Fs=fs,
                           window = lambda d: d*winfunc(fftlen),
                           visible=False)
        X_in = 10.0*scipy.log10(abs(fftpack.fftshift(X)))
        f_in = scipy.arange(-fs/2.0, fs/2.0, fs/float(X_in.size))
        p1_f = sp1_f.plot(f_in, X_in, "b")
        sp1_f.set_xlim([min(f_in), max(f_in)+1])
        sp1_f.set_ylim([-120.0, 20.0])

        sp1_f.set_title("Input Signal", weight="bold")
        sp1_f.set_xlabel("Frequency (Hz)")
        sp1_f.set_ylabel("Power (dBW)")

        Ts = 1.0/fs
        Tmax = len(d)*Ts

        t_in = scipy.arange(0, Tmax, Ts)
        x_in = scipy.array(d)
        sp1_t = fig1.add_subplot(2, 1, 2)
        p1_t = sp1_t.plot(t_in, x_in.real, "b-o")
        #p1_t = sp1_t.plot(t_in, x_in.imag, "r-o")
        sp1_t.set_ylim([-5, 5])

        # Set up the number of rows and columns for plotting the subfigures
        Ncols = int(scipy.floor(scipy.sqrt(fm.num_rx_channels())))
        Nrows = int(scipy.floor(fm.num_rx_channels() / Ncols))
        if(fm.num_rx_channels() % Ncols != 0):
            Nrows += 1

        # Plot each of the channels outputs. Frequencies on Figure 2 and
        # time signals on Figure 3
        fs_o = fm._audio_rate
        for i in xrange(len(fm.snks)):
            # remove issues with the transients at the beginning
            # also remove some corruption at the end of the stream
            #    this is a bug, probably due to the corner cases
            d = fm.snks[i].data()[Ns:Ne]

            sp2_f = fig2.add_subplot(Nrows, Ncols, 1+i)
            X,freq = sp2_f.psd(d, NFFT=fftlen, noverlap=fftlen/4, Fs=fs_o,
                               window = lambda d: d*winfunc(fftlen),
                               visible=False)
            #X_o = 10.0*scipy.log10(abs(fftpack.fftshift(X)))
            X_o = 10.0*scipy.log10(abs(X))
            #f_o = scipy.arange(-fs_o/2.0, fs_o/2.0, fs_o/float(X_o.size))
            f_o = scipy.arange(0, fs_o/2.0, fs_o/2.0/float(X_o.size))
            p2_f = sp2_f.plot(f_o, X_o, "b")
            sp2_f.set_xlim([min(f_o), max(f_o)+0.1])
            sp2_f.set_ylim([-120.0, 20.0])
            sp2_f.grid(True)

            sp2_f.set_title(("Channel %d" % i), weight="bold")
            sp2_f.set_xlabel("Frequency (kHz)")
            sp2_f.set_ylabel("Power (dBW)")

            Ts = 1.0/fs_o
            Tmax = len(d)*Ts
            t_o = scipy.arange(0, Tmax, Ts)

            x_t = scipy.array(d)
            sp2_t = fig3.add_subplot(Nrows, Ncols, 1+i)
            p2_t = sp2_t.plot(t_o, x_t.real, "b")
            p2_t = sp2_t.plot(t_o, x_t.imag, "r")
            sp2_t.set_xlim([min(t_o), max(t_o)+1])
            sp2_t.set_ylim([-1, 1])

            sp2_t.set_xlabel("Time (s)")
            sp2_t.set_ylabel("Amplitude")

        pylab.show()

if __name__ == "__main__":
    main()


from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import matplotlib as mpl
import matplotlib.pyplot as plt
from nose.tools import with_setup, make_decorator, assert_true
import warnings

figsize_orig = mpl.rcParams["figure.figsize"]
def setup_package():
    mpl.rcParams["figure.figsize"] = (11.0, 8.0)

def teardown_package():
    mpl.rcParams["figure.figsize"] = figsize_orig

import os

# Testing framework shamelessly stolen from matplotlib...

# Tests which should be run with 'python tests.py' or via 'must be
# included here.
default_test_modules = [
    'ggplot.tests.test_basic',
    'ggplot.tests.test_readme_examples',
    'ggplot.tests.test_ggplot_internals',
    'ggplot.tests.test_geom',
    'ggplot.tests.test_stat',
    'ggplot.tests.test_stat_calculate_methods',
    'ggplot.tests.test_stat_summary',
    'ggplot.tests.test_geom_rect',
    'ggplot.tests.test_geom_dotplot',
    'ggplot.tests.test_geom_bar',
    'ggplot.tests.test_qplot',
    'ggplot.tests.test_geom_lines',
    'ggplot.tests.test_geom_linerange',
    'ggplot.tests.test_geom_pointrange',
    'ggplot.tests.test_faceting',
    'ggplot.tests.test_stat_function',
    'ggplot.tests.test_scale_facet_wrap',
    'ggplot.tests.test_scale_log',
    'ggplot.tests.test_reverse',
    'ggplot.tests.test_ggsave',
    'ggplot.tests.test_theme_mpl',
    'ggplot.tests.test_colors',
    'ggplot.tests.test_chart_components',
    'ggplot.tests.test_legend',
    'ggplot.tests.test_element_target',
    'ggplot.tests.test_element_text',
    'ggplot.tests.test_theme',
    'ggplot.tests.test_theme_bw',
    'ggplot.tests.test_theme_gray',
    'ggplot.tests.test_theme_mpl',
    'ggplot.tests.test_theme_seaborn'
]

_multiprocess_can_split_ = True

# Check that the test directories exist
if not os.path.exists(os.path.join(
        os.path.dirname(__file__), 'baseline_images')):
    raise IOError(
        'The baseline image directory does not exist. '
        'This is most likely because the test data is not installed. '
        'You may need to install ggplot from source to get the '
        'test data.')

def _assert_same_ggplot_image(gg, name, test_file, tol=17):
    """Asserts that the ggplot object produces the right image"""
    fig = gg.draw()
    return _assert_same_figure_images(fig, name, test_file, tol=tol)

class ImagesComparisonFailure(Exception):
    pass

def _assert_same_figure_images(fig, name, test_file, tol=17):
    """Asserts that the figure object produces the right image"""
    import os
    import shutil
    from matplotlib import cbook
    from matplotlib.testing.compare import compare_images
    from nose.tools import assert_is_not_none

    if not ".png" in name:
        name = name+".png"

    basedir = os.path.abspath(os.path.dirname(test_file))
    basename = os.path.basename(test_file)
    subdir = os.path.splitext(basename)[0]

    baseline_dir = os.path.join(basedir, 'baseline_images', subdir)
    result_dir = os.path.abspath(os.path.join('result_images', subdir))

    if not os.path.exists(result_dir):
        cbook.mkdirs(result_dir)

    orig_expected_fname = os.path.join(baseline_dir, name)
    actual_fname = os.path.join(result_dir, name)

    def make_test_fn(fname, purpose):
        base, ext = os.path.splitext(fname)
        return '%s-%s%s' % (base, purpose, ext)
    expected_fname = make_test_fn(actual_fname, 'expected')
    # Save the figure before testing whether the original image
    # actually exists. This make creating new tests much easier,
    # as the result image can afterwards just be copied.
    fig.savefig(actual_fname)
    if os.path.exists(orig_expected_fname):
        shutil.copyfile(orig_expected_fname, expected_fname)
    else:
        raise Exception("Baseline image %s is missing" % orig_expected_fname)
    err = compare_images(expected_fname, actual_fname,
                         tol, in_decorator=True)
    if err:
        msg = 'images not close: {actual:s} vs. {expected:s} (RMS {rms:.2f})'.format(**err)
        raise ImagesComparisonFailure(msg)
    return err

def get_assert_same_ggplot(test_file):
    """Returns a "assert_same_ggplot" function for these test file

    call it like `assert_same_ggplot = get_assert_same_ggplot(__file__)`
    """
    def curried(*args, **kwargs):
        kwargs["test_file"] = test_file
        return _assert_same_ggplot_image(*args, **kwargs)
    curried.__doc__ = _assert_same_ggplot_image.__doc__
    return curried

def assert_same_elements(first,second, msg=None):
    assert_true(len(first) == len(second), "different length")
    assert_true(all([a==b for a,b in zip(first,second)]), "Unequal: %s vs %s" % (first, second))

def image_comparison(baseline_images=None, tol=17, extensions=None):
    """
    call signature::

      image_comparison(baseline_images=['my_figure'], tol=17)

    Compare images generated by the test with those specified in
    *baseline_images*, which must correspond else an
    ImagesComparisonFailure exception will be raised.

    Keyword arguments:

      *baseline_images*: list
        A list of strings specifying the names of the images generated
        by calls to :meth:`matplotlib.figure.savefig`.

      *tol*: (default 13)
        The RMS threshold above which the test is considered failed.

    """

    if baseline_images is None:
        raise ValueError('baseline_images must be specified')

    if extensions:
        # ignored, only for compatibility with matplotlibs decorator!
        pass

    def compare_images_decorator(func):
        import inspect
        _file = inspect.getfile(func)
        def decorated():
            # make sure we don't carry over bad images from former tests.
            assert len(plt.get_fignums()) == 0, "no of open figs: %s -> find the last test with ' " \
                                        "python tests.py -v' and add a '@cleanup' decorator." % \
                                        str(plt.get_fignums())
            func()
            assert len(plt.get_fignums()) == len(baseline_images), "different number of " \
                                                                   "baseline_images and actuall " \
                                                                   "plots."
            for fignum, baseline in zip(plt.get_fignums(), baseline_images):
                figure = plt.figure(fignum)
                _assert_same_figure_images(figure, baseline, _file, tol=tol)
        # also use the cleanup decorator to close any open figures!
        return make_decorator(cleanup(func))(decorated)
    return compare_images_decorator

def cleanup(func):
    """Decorator to add cleanup to the testing function

      @cleanup
      def test_something():
          " ... "

    Note that `@cleanup` is useful *only* for test functions, not for test
    methods or inside of TestCase subclasses.
    """

    def _teardown():
        plt.close('all')
        warnings.resetwarnings() #reset any warning filters set in tests

    return with_setup(setup=_setup, teardown=_teardown)(func)

# This is called from the cleanup decorator
def _setup():
    # The baseline images are created in this locale, so we should use
    # it during all of the tests.
    import locale
    import warnings
    from matplotlib.backends import backend_agg, backend_pdf, backend_svg

    try:
        locale.setlocale(locale.LC_ALL, str('en_US.UTF-8'))
    except locale.Error:
        try:
            locale.setlocale(locale.LC_ALL, str('English_United States.1252'))
        except locale.Error:
            warnings.warn(
                "Could not set locale to English/United States. "
                "Some date-related tests may fail")

    mpl.use('Agg', warn=False)  # use Agg backend for these tests
    if mpl.get_backend().lower() != "agg" and mpl.get_backend().lower() != "qt4agg":
        raise Exception(("Using a wrong matplotlib backend ({0}), which will not produce proper "
                        "images").format(mpl.get_backend()))

    # These settings *must* be hardcoded for running the comparison
    # tests
    mpl.rcdefaults()  # Start with all defaults
    mpl.rcParams['text.hinting'] = True
    mpl.rcParams['text.antialiased'] = True
    #mpl.rcParams['text.hinting_factor'] = 8

    # Clear the font caches.  Otherwise, the hinting mode can travel
    # from one test to another.
    backend_agg.RendererAgg._fontd.clear()
    backend_pdf.RendererPdf.truetype_font_cache.clear()
    backend_svg.RendererSVG.fontd.clear()
    # make sure we don't carry over bad plots from former tests
    assert len(plt.get_fignums()) == 0, "no of open figs: %s -> find the last test with ' " \
                                        "python tests.py -v' and add a '@cleanup' decorator." % \
                                        str(plt.get_fignums())

# This is here to run it like "from ggplot.tests import test; test()"
def test(verbosity=1):
    """run the ggplot test suite"""
    old_backend = mpl.rcParams['backend']
    try:
        mpl.use('agg')
        import nose
        import nose.plugins.builtin
        from matplotlib.testing.noseclasses import KnownFailure
        from nose.plugins.manager import PluginManager
        from nose.plugins import multiprocess

        # store the old values before overriding
        plugins = []
        plugins.append( KnownFailure() )
        plugins.extend( [plugin() for plugin in nose.plugins.builtin.plugins] )

        manager = PluginManager(plugins=plugins)
        config = nose.config.Config(verbosity=verbosity, plugins=manager)

        # Nose doesn't automatically instantiate all of the plugins in the
        # child processes, so we have to provide the multiprocess plugin with
        # a list.
        multiprocess._instantiate_plugins = [KnownFailure]

        success = nose.run( defaultTest=default_test_modules,
                            config=config,
                            )
    finally:
        if old_backend.lower() != 'agg':
            mpl.use(old_backend)

    return success

test.__test__ = False # nose: this function is not a test


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

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

import numpy as np

from tensorflow.contrib.learn.python.learn.learn_io import pandas_io
from tensorflow.python.framework import errors
from tensorflow.python.platform import test
from tensorflow.python.training import coordinator
from tensorflow.python.training import queue_runner_impl

# pylint: disable=g-import-not-at-top
try:
  import pandas as pd
  HAS_PANDAS = True
except ImportError:
  HAS_PANDAS = False

class PandasIoTest(test.TestCase):

  def makeTestDataFrame(self):
    index = np.arange(100, 104)
    a = np.arange(4)
    b = np.arange(32, 36)
    x = pd.DataFrame({'a': a, 'b': b}, index=index)
    y = pd.Series(np.arange(-32, -28), index=index)
    return x, y

  def callInputFnOnce(self, input_fn, session):
    results = input_fn()
    coord = coordinator.Coordinator()
    threads = queue_runner_impl.start_queue_runners(session, coord=coord)
    result_values = session.run(results)
    coord.request_stop()
    coord.join(threads)
    return result_values

  def testPandasInputFn_IndexMismatch(self):
    if not HAS_PANDAS:
      return
    x, _ = self.makeTestDataFrame()
    y_noindex = pd.Series(np.arange(-32, -28))
    with self.assertRaises(ValueError):
      pandas_io.pandas_input_fn(
          x, y_noindex, batch_size=2, shuffle=False, num_epochs=1)

  def testPandasInputFn_ProducesExpectedOutputs(self):
    if not HAS_PANDAS:
      return
    with self.cached_session() as session:
      x, y = self.makeTestDataFrame()
      input_fn = pandas_io.pandas_input_fn(
          x, y, batch_size=2, shuffle=False, num_epochs=1)

      features, target = self.callInputFnOnce(input_fn, session)

      self.assertAllEqual(features['a'], [0, 1])
      self.assertAllEqual(features['b'], [32, 33])
      self.assertAllEqual(target, [-32, -31])

  def testPandasInputFn_ProducesOutputsForLargeBatchAndMultipleEpochs(self):
    if not HAS_PANDAS:
      return
    with self.cached_session() as session:
      index = np.arange(100, 102)
      a = np.arange(2)
      b = np.arange(32, 34)
      x = pd.DataFrame({'a': a, 'b': b}, index=index)
      y = pd.Series(np.arange(-32, -30), index=index)
      input_fn = pandas_io.pandas_input_fn(
          x, y, batch_size=128, shuffle=False, num_epochs=2)

      results = input_fn()

      coord = coordinator.Coordinator()
      threads = queue_runner_impl.start_queue_runners(session, coord=coord)

      features, target = session.run(results)
      self.assertAllEqual(features['a'], [0, 1, 0, 1])
      self.assertAllEqual(features['b'], [32, 33, 32, 33])
      self.assertAllEqual(target, [-32, -31, -32, -31])

      with self.assertRaises(errors.OutOfRangeError):
        session.run(results)

      coord.request_stop()
      coord.join(threads)

  def testPandasInputFn_ProducesOutputsWhenDataSizeNotDividedByBatchSize(self):
    if not HAS_PANDAS:
      return
    with self.cached_session() as session:
      index = np.arange(100, 105)
      a = np.arange(5)
      b = np.arange(32, 37)
      x = pd.DataFrame({'a': a, 'b': b}, index=index)
      y = pd.Series(np.arange(-32, -27), index=index)

      input_fn = pandas_io.pandas_input_fn(
          x, y, batch_size=2, shuffle=False, num_epochs=1)

      results = input_fn()

      coord = coordinator.Coordinator()
      threads = queue_runner_impl.start_queue_runners(session, coord=coord)

      features, target = session.run(results)
      self.assertAllEqual(features['a'], [0, 1])
      self.assertAllEqual(features['b'], [32, 33])
      self.assertAllEqual(target, [-32, -31])

      features, target = session.run(results)
      self.assertAllEqual(features['a'], [2, 3])
      self.assertAllEqual(features['b'], [34, 35])
      self.assertAllEqual(target, [-30, -29])

      features, target = session.run(results)
      self.assertAllEqual(features['a'], [4])
      self.assertAllEqual(features['b'], [36])
      self.assertAllEqual(target, [-28])

      with self.assertRaises(errors.OutOfRangeError):
        session.run(results)

      coord.request_stop()
      coord.join(threads)

  def testPandasInputFn_OnlyX(self):
    if not HAS_PANDAS:
      return
    with self.cached_session() as session:
      x, _ = self.makeTestDataFrame()
      input_fn = pandas_io.pandas_input_fn(
          x, y=None, batch_size=2, shuffle=False, num_epochs=1)

      features = self.callInputFnOnce(input_fn, session)

      self.assertAllEqual(features['a'], [0, 1])
      self.assertAllEqual(features['b'], [32, 33])

  def testPandasInputFn_ExcludesIndex(self):
    if not HAS_PANDAS:
      return
    with self.cached_session() as session:
      x, y = self.makeTestDataFrame()
      input_fn = pandas_io.pandas_input_fn(
          x, y, batch_size=2, shuffle=False, num_epochs=1)

      features, _ = self.callInputFnOnce(input_fn, session)

      self.assertFalse('index' in features)

  def assertInputsCallableNTimes(self, input_fn, session, n):
    inputs = input_fn()
    coord = coordinator.Coordinator()
    threads = queue_runner_impl.start_queue_runners(session, coord=coord)
    for _ in range(n):
      session.run(inputs)
    with self.assertRaises(errors.OutOfRangeError):
      session.run(inputs)
    coord.request_stop()
    coord.join(threads)

  def testPandasInputFn_RespectsEpoch_NoShuffle(self):
    if not HAS_PANDAS:
      return
    with self.cached_session() as session:
      x, y = self.makeTestDataFrame()
      input_fn = pandas_io.pandas_input_fn(
          x, y, batch_size=4, shuffle=False, num_epochs=1)

      self.assertInputsCallableNTimes(input_fn, session, 1)

  def testPandasInputFn_RespectsEpoch_WithShuffle(self):
    if not HAS_PANDAS:
      return
    with self.cached_session() as session:
      x, y = self.makeTestDataFrame()
      input_fn = pandas_io.pandas_input_fn(
          x, y, batch_size=4, shuffle=True, num_epochs=1)

      self.assertInputsCallableNTimes(input_fn, session, 1)

  def testPandasInputFn_RespectsEpoch_WithShuffleAutosize(self):
    if not HAS_PANDAS:
      return
    with self.cached_session() as session:
      x, y = self.makeTestDataFrame()
      input_fn = pandas_io.pandas_input_fn(
          x, y, batch_size=2, shuffle=True, queue_capacity=None, num_epochs=2)

      self.assertInputsCallableNTimes(input_fn, session, 4)

  def testPandasInputFn_RespectsEpochUnevenBatches(self):
    if not HAS_PANDAS:
      return
    x, y = self.makeTestDataFrame()
    with self.cached_session() as session:
      input_fn = pandas_io.pandas_input_fn(
          x, y, batch_size=3, shuffle=False, num_epochs=1)

      # Before the last batch, only one element of the epoch should remain.
      self.assertInputsCallableNTimes(input_fn, session, 2)

  def testPandasInputFn_Idempotent(self):
    if not HAS_PANDAS:
      return
    x, y = self.makeTestDataFrame()
    for _ in range(2):
      pandas_io.pandas_input_fn(
          x, y, batch_size=2, shuffle=False, num_epochs=1)()
    for _ in range(2):
      pandas_io.pandas_input_fn(
          x, y, batch_size=2, shuffle=True, num_epochs=1)()

if __name__ == '__main__':
  test.main()


#!/usr/bin/python3
# -*- coding:utf-8 -*-

import __future__
import parser
import sys
import matplotlib.pyplot as plt
#plt.style.use('ggplot')
import numpy as np
import operator
from collections import *

caseSize = (8192, 8192)
if parser.args.res:
    maxAvailableNode = parser.args.res
else:
    maxAvailableNode = 8

sizeDataDict = []
for p in range(0, int(np.log2(maxAvailableNode)) + 1):
    filterDict = {'nSizeX' : caseSize[0], 'nSizeY' : caseSize[1], 'R' : 64 * 2**p}
    print filterDict
    data = parser.getData(filterDict)
    if len(data):
        sizeDataDict.append(data)

if len(sizeDataDict) == 0:
    print("No data found.")
    sys.exit(1)

loopTimeDict = dict()
for data in sizeDataDict:
    for key, value in data.items():
        keyDict = parser.extractKey(key)

        Nt = keyDict['Nt']
        R = keyDict['R']

        if keyDict['Ny'] != caseSize[0] or keyDict['Nx'] != caseSize[1]:
            print("Error in collected data")
            sys.exit(1)

        for run in value:
            nSDD = run['point'][0] * run['point'][1]

            # On several nodes, select only pure SDD, which is the best result.
            if R > 64 and nSDD < R:
                continue

            # Don't remove HyperThreading.
            # We assume that hyperthreading with SDD leads to same results as with SDS.
            #if R > 64 and nSDD == R and Nt > 1.0:
            #    continue

            # On a single node, select only pure SDS
            if R == 64 and nSDD > 1:
                continue

            loopT = run['loopTime'] * caseSize[0] * caseSize[1] * keyDict['Ni'] / 1000.
            if R not in loopTimeDict.keys():
                loopTimeDict[R] = list()
            loopTimeDict[R].append(loopT)

# And now, we must plot that
fig = plt.figure(0, figsize=(9, 6))
ax = fig.add_subplot(111)
#ax = fig.add_subplot(211)
#ax.set_xscale('log', basex=2)
#ax.set_yscale('log')

maxSimulationNumber = 42
xArray = range(1, maxSimulationNumber + 1)

'''
#Perfect Scale

loopTimeDict[128] = [k / 2. for k in loopTimeDict[64]]
loopTimeDict[256] = [k / 4. for k in loopTimeDict[64]]
loopTimeDict[512] = [k / 8. for k in loopTimeDict[64]]
'''

for r in sorted(loopTimeDict):
    nodeNeeded = r // 64
    minT = np.min(loopTimeDict[r])
    print("Min Time %s node(s) = %s" % (nodeNeeded, minT))
    totalTimeArray = np.zeros(maxSimulationNumber)
    for i in xArray:
        totalTimeArray[i-1] = minT * (1 + (i * nodeNeeded - 1) // maxAvailableNode)
    ax.plot(xArray, totalTimeArray, '-', label="Batch Size %s" % (r // 64))
    parser.outputCurve("ergodicity_scaling-%s.dat" % (r//64), xArray, totalTimeArray)

'''
minSize = int(np.sqrt(np.min(syncTimeDict.keys())))
maxSize = int(np.sqrt(np.max(syncTimeDict.keys())))
nodeNumber = (caseSize[0] * caseSize[1] / (maxSize * maxSize))
'''

plt.title('%sx%s batch time with %s node(s) available at the same time.' % (caseSize[0], caseSize[1], maxAvailableNode))

plt.xlabel('Total number of simulation to run')
plt.ylabel('Loop Time')
plt.legend()

'''
bx = fig.add_subplot(212)
bx.set_xscale('log', basex=2)
bx.plot(sorted(sdsWeakDict), [np.min(v) for k, v in sorted(sdsWeakDict.items(), key=operator.itemgetter(0))], 'g+-', label="SDS scaling")
bx.plot(sorted(sddWeakDict), [np.min(v) for k, v in sorted(sddWeakDict.items())], 'b+-', label="SDD scaling")
#bx.plot(sorted(hybridWeakDict), [np.min(v) for k, v in sorted(hybridWeakDict.items())], 'y+-', label="Hybrid scaling")
bx.plot(sorted(sddWeakDict), [firstValueSDD for k in sorted(sddWeakDict.keys())], 'b--', label="SDD ideal")
bx.plot(sorted(sdsWeakDict), [firstValueSDS for k in sorted(sdsWeakDict.keys())], 'g--', label="SDS ideal")

for k in sdsWeakDict:
    bx.plot(np.full(len(sdsWeakDict[k]), k), sdsWeakDict[k], 'g+')
for k in sddWeakDict:
    bx.plot(np.full(len(sddWeakDict[k]), k), sddWeakDict[k], 'b+')

plt.title('Weak Scaling from %sx%s to %sx%s' % (initSize, initSize, initSize * 2**((maxPower-1) / 2), initSize * 2**((maxPower-1) / 2)) )
plt.xlabel('Core(s)')
plt.ylabel('Loop Time / iteration')
plt.legend()
'''

plt.show()


"""Loader for the Labeled Faces in the Wild (LFW) dataset

This dataset is a collection of JPEG pictures of famous people collected
over the internet, all details are available on the official website:

    http://vis-www.cs.umass.edu/lfw/

Each picture is centered on a single face. The typical task is called
Face Verification: given a pair of two pictures, a binary classifier
must predict whether the two images are from the same person.

An alternative task, Face Recognition or Face Identification is:
given the picture of the face of an unknown person, identify the name
of the person by refering to a gallery of previously seen pictures of
identified persons.

Both Face Verification and Face Recognition are tasks that are typically
performed on the output of a model trained to perform Face Detection. The
most popular model for Face Detection is called Viola-Johns and is
implemented in the OpenCV library. The LFW faces were extracted by this face
detector from various online websites.
"""
# Copyright (c) 2011 Olivier Grisel <olivier.grisel@ensta.org>
# License: Simplified BSD

from os import listdir, makedirs, remove
from os.path import join, exists, isdir

import logging
import numpy as np
import urllib

from .base import get_data_home, Bunch
from ..externals.joblib import Memory

logger = logging.getLogger(__name__)

BASE_URL = "http://vis-www.cs.umass.edu/lfw/"
ARCHIVE_NAME = "lfw.tgz"
FUNNELED_ARCHIVE_NAME = "lfw-funneled.tgz"
TARGET_FILENAMES = [
    'pairsDevTrain.txt',
    'pairsDevTest.txt',
    'pairs.txt',
]

def scale_face(face):
    """Scale back to 0-1 range in case of normalization for plotting"""
    scaled = face - face.min()
    scaled /= scaled.max()
    return scaled

#
# Common private utilities for data fetching from the original LFW website
# local disk caching, and image decoding.
#

def check_fetch_lfw(data_home=None, funneled=True, download_if_missing=True):
    """Helper function to download any missing LFW data"""
    data_home = get_data_home(data_home=data_home)
    lfw_home = join(data_home, "lfw_home")

    if funneled:
        archive_path = join(lfw_home, FUNNELED_ARCHIVE_NAME)
        data_folder_path = join(lfw_home, "lfw_funneled")
        archive_url = BASE_URL + FUNNELED_ARCHIVE_NAME
    else:
        archive_path = join(lfw_home, ARCHIVE_NAME)
        data_folder_path = join(lfw_home, "lfw")
        archive_url = BASE_URL + ARCHIVE_NAME

    if not exists(lfw_home):
        makedirs(lfw_home)

    for target_filename in TARGET_FILENAMES:
        target_filepath = join(lfw_home, target_filename)
        if not exists(target_filepath):
            if download_if_missing:
                url = BASE_URL + target_filename
                logger.warn("Downloading LFW metadata: %s", url)
                urllib.urlretrieve(url, target_filepath)
            else:
                raise IOError("%s is missing" % target_filepath)

    if not exists(data_folder_path):

        if not exists(archive_path):
            if download_if_missing:
                logger.warn("Downloading LFW data (~200MB): %s", archive_url)
                urllib.urlretrieve(archive_url, archive_path)
            else:
                raise IOError("%s is missing" % target_filepath)

        import tarfile
        logger.info("Decompressing the data archive to %s", data_folder_path)
        tarfile.open(archive_path, "r:gz").extractall(path=lfw_home)
        remove(archive_path)

    return lfw_home, data_folder_path

def _load_imgs(file_paths, slice_, color, resize):
    """Internally used to load images"""

    # Try to import imread and imresize from PIL. We do this here to prevent
    # the whole sklearn.datasets module from depending on PIL.
    try:
        try:
            from scipy.misc import imread
        except ImportError:
            from scipy.misc.pilutil import imread
        from scipy.misc import imresize
    except ImportError:
        raise ImportError("The Python Imaging Library (PIL)"
                          "is required to load data from jpeg files")

    # compute the portion of the images to load to respect the slice_ parameter
    # given by the caller
    default_slice = (slice(0, 250), slice(0, 250))
    if slice_ is None:
        slice_ = default_slice
    else:
        slice_ = tuple(s or ds for s, ds in zip(slice_, default_slice))

    h_slice, w_slice = slice_
    h = (h_slice.stop - h_slice.start) / (h_slice.step or 1)
    w = (w_slice.stop - w_slice.start) / (w_slice.step or 1)

    if resize is not None:
        resize = float(resize)
        h = int(resize * h)
        w = int(resize * w)

    # allocate some contiguous memory to host the decoded image slices
    n_faces = len(file_paths)
    if not color:
        faces = np.zeros((n_faces, h, w), dtype=np.float32)
    else:
        faces = np.zeros((n_faces, h, w, 3), dtype=np.float32)

    # iterate over the collected file path to load the jpeg files as numpy
    # arrays
    for i, file_path in enumerate(file_paths):
        if i % 1000 == 0:
            logger.info("Loading face #%05d / %05d", i + 1, n_faces)
        face = np.asarray(imread(file_path)[slice_], dtype=np.float32)
        face /= 255.0  # scale uint8 coded colors to the [0.0, 1.0] floats
        if resize is not None:
            face = imresize(face, resize)
        if not color:
            # average the color channels to compute a gray levels
            # representaion
            face = face.mean(axis=2)

        faces[i, ...] = face

    return faces

#
# Task #1:  Face Identification on picture with names
#

def _fetch_lfw_people(data_folder_path, slice_=None, color=False, resize=None,
                     min_faces_per_person=0):
    """Perform the actual data loading for the lfw people dataset

    This operation is meant to be cached by a joblib wrapper.
    """
    # scan the data folder content to retain people with more that
    # `min_faces_per_person` face pictures
    person_names, file_paths = [], []
    for person_name in sorted(listdir(data_folder_path)):
        folder_path = join(data_folder_path, person_name)
        if not isdir(folder_path):
            continue
        paths = [join(folder_path, f) for f in listdir(folder_path)]
        n_pictures = len(paths)
        if n_pictures >= min_faces_per_person:
            person_name = person_name.replace('_', ' ')
            person_names.extend([person_name] * n_pictures)
            file_paths.extend(paths)

    n_faces = len(file_paths)
    if n_faces == 0:
        raise ValueError("min_faces_per_person=%d is too restrictive" %
                         min_faces_per_person)

    target_names = np.unique(person_names)
    target = np.searchsorted(target_names, person_names)

    faces = _load_imgs(file_paths, slice_, color, resize)

    # shuffle the faces with a deterministic RNG scheme to avoid having
    # all faces of the same person in a row, as it would break some
    # cross validation and learning algorithms such as SGD and online
    # k-means that make an IID assumption

    indices = np.arange(n_faces)
    np.random.RandomState(42).shuffle(indices)
    faces, target = faces[indices], target[indices]
    return faces, target, target_names

def fetch_lfw_people(data_home=None, funneled=True, resize=0.5,
                    min_faces_per_person=None, color=False,
                    slice_=(slice(70, 195), slice(78, 172)),
                    download_if_missing=True):
    """Loader for the Labeled Faces in the Wild (LFW) people dataset

    This dataset is a collection of JPEG pictures of famous people
    collected on the internet, all details are available on the
    official website:

        http://vis-www.cs.umass.edu/lfw/

    Each picture is centered on a single face. Each pixel of each channel
    (color in RGB) is encoded by a float in range 0.0 - 1.0.

    The task is called Face Recognition (or Identification): given the
    picture of a face, find the name of the person given a training set
    (gallery).

    Parameters
    ----------
    data_home: optional, default: None
        Specify another download and cache folder for the datasets. By default
        all scikit learn data is stored in '~/scikit_learn_data' subfolders.

    funneled: boolean, optional, default: True
        Download and use the funneled variant of the dataset.

    resize: float, optional, default 0.5
        Ratio used to resize the each face picture.

    min_faces_per_person: int, optional, default None
        The extracted dataset will only retain pictures of people that have at
        least `min_faces_per_person` different pictures.

    color: boolean, optional, default False
        Keep the 3 RGB channels instead of averaging them to a single
        gray level channel. If color is True the shape of the data has
        one more dimension than than the shape with color = False.

    slice_: optional
        Provide a custom 2D slice (height, width) to extract the
        'interesting' part of the jpeg files and avoid use statistical
        correlation from the background

    download_if_missing: optional, True by default
        If False, raise a IOError if the data is not locally available
        instead of trying to download the data from the source site.
    """
    lfw_home, data_folder_path = check_fetch_lfw(
        data_home=data_home, funneled=funneled,
        download_if_missing=download_if_missing)
    logger.info('Loading LFW people faces from %s', lfw_home)

    # wrap the loader in a memoizing function that will return memmaped data
    # arrays for optimal memory usage
    m = Memory(cachedir=lfw_home, compress=6, verbose=0)
    load_func = m.cache(_fetch_lfw_people)

    # load and memoize the pairs as np arrays
    faces, target, target_names = load_func(
        data_folder_path, resize=resize,
        min_faces_per_person=min_faces_per_person, color=color, slice_=slice_)

    # pack the results as a Bunch instance
    return Bunch(data=faces.reshape(len(faces), -1), images=faces,
                 target=target, target_names=target_names,
                 DESCR="LFW faces dataset")

#
# Task #2:  Face Verification on pairs of face pictures
#

def _fetch_lfw_pairs(index_file_path, data_folder_path, slice_=None,
                    color=False, resize=None):
    """Perform the actual data loading for the LFW pairs dataset

    This operation is meant to be cached by a joblib wrapper.
    """
    # parse the index file to find the number of pairs to be able to allocate
    # the right amount of memory before starting to decode the jpeg files
    with open(index_file_path, 'rb') as index_file:
        split_lines = [ln.strip().split('\t') for ln in index_file]
    pair_specs = [sl for sl in split_lines if len(sl) > 2]
    n_pairs = len(pair_specs)

    # interating over the metadata lines for each pair to find the filename to
    # decode and load in memory
    target = np.zeros(n_pairs, dtype=np.int)
    file_paths = list()
    for i, components in enumerate(pair_specs):
        if len(components) == 3:
            target[i] = 1
            pair = (
                (components[0], int(components[1]) - 1),
                (components[0], int(components[2]) - 1),
            )
        elif len(components) == 4:
            target[i] = 0
            pair = (
                (components[0], int(components[1]) - 1),
                (components[2], int(components[3]) - 1),
            )
        else:
            raise ValueError("invalid line %d: %r" % (i + 1, components))
        for j, (name, idx) in enumerate(pair):
            person_folder = join(data_folder_path, name)
            filenames = list(sorted(listdir(person_folder)))
            file_path = join(person_folder, filenames[idx])
            file_paths.append(file_path)

    pairs = _load_imgs(file_paths, slice_, color, resize)
    shape = list(pairs.shape)
    n_faces = shape.pop(0)
    shape.insert(0, 2)
    shape.insert(0, n_faces // 2)
    pairs.shape = shape

    return pairs, target, np.array(['Different persons', 'Same person'])

def load_lfw_people(download_if_missing=False, **kwargs):
    """Alias for fetch_lfw_people(download_if_missing=False)

    Check fetch_lfw_people.__doc__ for the documentation and parameter list.
    """
    return fetch_lfw_people(download_if_missing=download_if_missing, **kwargs)

def fetch_lfw_pairs(subset='train', data_home=None, funneled=True, resize=0.5,
                   color=False, slice_=(slice(70, 195), slice(78, 172)),
                    download_if_missing=True):
    """Loader for the Labeled Faces in the Wild (LFW) pairs dataset

    This dataset is a collection of JPEG pictures of famous people
    collected on the internet, all details are available on the
    official website:

        http://vis-www.cs.umass.edu/lfw/

    Each picture is centered on a single face. Each pixel of each channel
    (color in RGB) is encoded by a float in range 0.0 - 1.0.

    The task is called Face Verification: given a pair of two pictures,
    a binary classifier must predict whether the two images are from
    the same person.

    In the official `README.txt`_ this task is described as the
    "Restricted" task.  As I am not sure as to implement the
    "Unrestricted" variant correctly, I left it as unsupported for now.

      .. _`README.txt`: http://vis-www.cs.umass.edu/lfw/README.txt

    Parameters
    ----------
    subset: optional, default: 'train'
        Select the dataset to load: 'train' for the development training
        set, 'test' for the development test set, and '10_folds' for the
        official evaluation set that is meant to be used with a 10-folds
        cross validation.

    data_home: optional, default: None
        Specify another download and cache folder for the datasets. By
        default all scikit learn data is stored in '~/scikit_learn_data'
        subfolders.

    funneled: boolean, optional, default: True
        Download and use the funneled variant of the dataset.

    resize: float, optional, default 0.5
        Ratio used to resize the each face picture.

    color: boolean, optional, default False
        Keep the 3 RGB channels instead of averaging them to a single
        gray level channel. If color is True the shape of the data has
        one more dimension than than the shape with color = False.

    slice_: optional
        Provide a custom 2D slice (height, width) to extract the
        'interesting' part of the jpeg files and avoid use statistical
        correlation from the background

    download_if_missing: optional, True by default
        If False, raise a IOError if the data is not locally available
        instead of trying to download the data from the source site.
    """
    lfw_home, data_folder_path = check_fetch_lfw(
        data_home=data_home, funneled=funneled,
        download_if_missing=download_if_missing)
    logger.info('Loading %s LFW pairs from %s', subset, lfw_home)

    # wrap the loader in a memoizing function that will return memmaped data
    # arrays for optimal memory usage
    m = Memory(cachedir=lfw_home, compress=6, verbose=0)
    load_func = m.cache(_fetch_lfw_pairs)

    # select the right metadata file according to the requested subset
    label_filenames = {
        'train': 'pairsDevTrain.txt',
        'test': 'pairsDevTest.txt',
        '10_folds': 'pairs.txt',
    }
    if subset not in label_filenames:
        raise ValueError("subset='%s' is invalid: should be one of %r" % (
            subset, list(sorted(label_filenames.keys()))))
    index_file_path = join(lfw_home, label_filenames[subset])

    # load and memoize the pairs as np arrays
    pairs, target, target_names = load_func(
        index_file_path, data_folder_path, resize=resize, color=color,
        slice_=slice_)

    # pack the results as a Bunch instance
    return Bunch(data=pairs.reshape(len(pairs), -1), pairs=pairs,
                 target=target, target_names=target_names,
                 DESCR="'%s' segment of the LFW pairs dataset" % subset)

def load_lfw_pairs(download_if_missing=False, **kwargs):
    """Alias for fetch_lfw_pairs(download_if_missing=False)

    Check fetch_lfw_pairs.__doc__ for the documentation and parameter list.
    """
    return fetch_lfw_pairs(download_if_missing=download_if_missing, **kwargs)


"""
===============================
Nearest Centroid Classification
===============================

Sample usage of Nearest Centroid classification.
It will plot the decision boundaries for each class.
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from sklearn import datasets
from sklearn.neighbors import NearestCentroid

n_neighbors = 15

# import some data to play with
iris = datasets.load_iris()
X = iris.data[:, :2]  # we only take the first two features. We could
                      # avoid this ugly slicing by using a two-dim dataset
y = iris.target

h = .02  # step size in the mesh

# Create color maps
cmap_light = ListedColormap(['#FFAAAA', '#AAFFAA', '#AAAAFF'])
cmap_bold = ListedColormap(['#FF0000', '#00FF00', '#0000FF'])

for shrinkage in [None, 0.1]:
    # we create an instance of Neighbours Classifier and fit the data.
    clf = NearestCentroid(shrink_threshold=shrinkage)
    clf.fit(X, y)
    y_pred = clf.predict(X)
    print(shrinkage, np.mean(y == y_pred))
    # Plot the decision boundary. For that, we will assign a color to each
    # point in the mesh [x_min, m_max]x[y_min, y_max].
    x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
    y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                         np.arange(y_min, y_max, h))
    Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])

    # Put the result into a color plot
    Z = Z.reshape(xx.shape)
    plt.figure()
    plt.pcolormesh(xx, yy, Z, cmap=cmap_light)

    # Plot also the training points
    plt.scatter(X[:, 0], X[:, 1], c=y, cmap=cmap_bold)
    plt.title("3-Class classification (shrink_threshold=%r)"
              % shrinkage)
    plt.axis('tight')

plt.show()


"""
===========================================
Robust linear model estimation using RANSAC
===========================================

In this example we see how to robustly fit a linear model to faulty data using
the RANSAC algorithm.

"""
import numpy as np
from matplotlib import pyplot as plt

from sklearn import linear_model, datasets

n_samples = 1000
n_outliers = 50

X, y, coef = datasets.make_regression(n_samples=n_samples, n_features=1,
                                      n_informative=1, noise=10,
                                      coef=True, random_state=0)

# Add outlier data
np.random.seed(0)
X[:n_outliers] = 3 + 0.5 * np.random.normal(size=(n_outliers, 1))
y[:n_outliers] = -3 + 10 * np.random.normal(size=n_outliers)

# Fit line using all data
model = linear_model.LinearRegression()
model.fit(X, y)

# Robustly fit linear model with RANSAC algorithm
model_ransac = linear_model.RANSACRegressor(linear_model.LinearRegression())
model_ransac.fit(X, y)
inlier_mask = model_ransac.inlier_mask_
outlier_mask = np.logical_not(inlier_mask)

# Predict data of estimated models
line_X = np.arange(-5, 5)
line_y = model.predict(line_X[:, np.newaxis])
line_y_ransac = model_ransac.predict(line_X[:, np.newaxis])

# Compare estimated coefficients
print("Estimated coefficients (true, normal, RANSAC):")
print(coef, model.coef_, model_ransac.estimator_.coef_)

plt.plot(X[inlier_mask], y[inlier_mask], '.g', label='Inliers')
plt.plot(X[outlier_mask], y[outlier_mask], '.r', label='Outliers')
plt.plot(line_X, line_y, '-k', label='Linear regressor')
plt.plot(line_X, line_y_ransac, '-b', label='RANSAC regressor')
plt.legend(loc='lower right')
plt.show()


"""
==================
Two-class AdaBoost
==================

This example fits an AdaBoosted decision stump on a non-linearly separable
classification dataset composed of two "Gaussian quantiles" clusters
(see :func:`sklearn.datasets.make_gaussian_quantiles`) and plots the decision
boundary and decision scores. The distributions of decision scores are shown
separately for samples of class A and B. The predicted class label for each
sample is determined by the sign of the decision score. Samples with decision
scores greater than zero are classified as B, and are otherwise classified
as A. The magnitude of a decision score determines the degree of likeness with
the predicted class label. Additionally, a new dataset could be constructed
containing a desired purity of class B, for example, by only selecting samples
with a decision score above some value.

"""
print(__doc__)

# Author: Noel Dawe <noel.dawe@gmail.com>
#
# License: BSD 3 clause

import numpy as np
import matplotlib.pyplot as plt

from sklearn.ensemble import AdaBoostClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.datasets import make_gaussian_quantiles

# Construct dataset
X1, y1 = make_gaussian_quantiles(cov=2.,
                                 n_samples=200, n_features=2,
                                 n_classes=2, random_state=1)
X2, y2 = make_gaussian_quantiles(mean=(3, 3), cov=1.5,
                                 n_samples=300, n_features=2,
                                 n_classes=2, random_state=1)
X = np.concatenate((X1, X2))
y = np.concatenate((y1, - y2 + 1))

# Create and fit an AdaBoosted decision tree
bdt = AdaBoostClassifier(DecisionTreeClassifier(max_depth=1),
                         algorithm="SAMME",
                         n_estimators=200)

bdt.fit(X, y)

plot_colors = "br"
plot_step = 0.02
class_names = "AB"

plt.figure(figsize=(10, 5))

# Plot the decision boundaries
plt.subplot(121)
x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, plot_step),
                     np.arange(y_min, y_max, plot_step))

Z = bdt.predict(np.c_[xx.ravel(), yy.ravel()])
Z = Z.reshape(xx.shape)
cs = plt.contourf(xx, yy, Z, cmap=plt.cm.Paired)
plt.axis("tight")

# Plot the training points
for i, n, c in zip(range(2), class_names, plot_colors):
    idx = np.where(y == i)
    plt.scatter(X[idx, 0], X[idx, 1],
                c=c, cmap=plt.cm.Paired,
                label="Class %s" % n)
plt.xlim(x_min, x_max)
plt.ylim(y_min, y_max)
plt.legend(loc='upper right')
plt.xlabel('x')
plt.ylabel('y')
plt.title('Decision Boundary')

# Plot the two-class decision scores
twoclass_output = bdt.decision_function(X)
plot_range = (twoclass_output.min(), twoclass_output.max())
plt.subplot(122)
for i, n, c in zip(range(2), class_names, plot_colors):
    plt.hist(twoclass_output[y == i],
             bins=10,
             range=plot_range,
             facecolor=c,
             label='Class %s' % n,
             alpha=.5)
x1, x2, y1, y2 = plt.axis()
plt.axis((x1, x2, y1, y2 * 1.2))
plt.legend(loc='upper right')
plt.ylabel('Samples')
plt.xlabel('Score')
plt.title('Decision Scores')

plt.tight_layout()
plt.subplots_adjust(wspace=0.35)
plt.show()


import logging
import os

import pandas as pd
import multiprocessing
import numpy as np

import dask
from dask.diagnostics import ProgressBar
from dask import delayed

from cloudpickle import cloudpickle
import tempfile

import py_entitymatching.catalog.catalog_manager as cm
import py_entitymatching.utils.catalog_helper as ch
import py_entitymatching.utils.generic_helper as gh

from py_entitymatching.utils.validation_helper import validate_object_type
from py_entitymatching.feature.extractfeatures import get_feature_vals_by_cand_split

from py_entitymatching.utils.validation_helper import validate_object_type
from py_entitymatching.dask.utils import validate_chunks, get_num_partitions, \
    get_num_cores, wrap

logger = logging.getLogger(__name__)

def dask_extract_feature_vecs(candset, attrs_before=None, feature_table=None,
                              attrs_after=None, verbose=False,
                              show_progress=True, n_chunks=1):
    """
    WARNING THIS COMMAND IS EXPERIMENTAL AND NOT TESTED. USE AT YOUR OWN RISK

    This function extracts feature vectors from a DataFrame (typically a
    labeled candidate set).

    Specifically, this function uses feature
    table, ltable and rtable (that is present in the `candset`'s
    metadata) to extract feature vectors.

    Args:
        candset (DataFrame): The input candidate set for which the features
            vectors should be extracted.

        attrs_before (list): The list of attributes from the input candset,
            that should be added before the feature vectors (defaults to None).

        feature_table (DataFrame): A DataFrame containing a list of
            features that should be used to compute the feature vectors (
            defaults to None).

        attrs_after (list): The list of attributes from the input candset
            that should be added after the feature vectors (defaults to None).

        verbose (boolean): A flag to indicate whether the debug information
            should be displayed (defaults to False).

        show_progress (boolean): A flag to indicate whether the progress of
            extracting feature vectors must be displayed (defaults to True).

        n_chunks (int): The number of partitions to split the candidate set. If it
            is set to -1, the number of partitions will be set to the
            number of cores in the machine.

    Returns:
        A pandas DataFrame containing feature vectors.

        The DataFrame will have metadata ltable and rtable, pointing
        to the same ltable and rtable as the input candset.

        Also, the output
        DataFrame will have three columns: key, foreign key ltable, foreign
        key rtable copied from input candset to the output DataFrame. These
        three columns precede the columns mentioned in `attrs_before`.

    Raises:
        AssertionError: If `candset` is not of type pandas
            DataFrame.
        AssertionError: If `attrs_before` has attributes that
            are not present in the input candset.
        AssertionError: If `attrs_after` has attribtues that
            are not present in the input candset.
        AssertionError: If `feature_table` is set to None.
        AssertionError: If `n_chunks` is not of type
                int.

    Examples:
        >>> import py_entitymatching as em
        >>> from py_entitymatching.dask.dask_extract_features import dask_extract_feature_vecs
        >>> A = em.read_csv_metadata('path_to_csv_dir/table_A.csv', key='ID')
        >>> B = em.read_csv_metadata('path_to_csv_dir/table_B.csv', key='ID')
        >>> match_f = em.get_features_for_matching(A, B)
        >>> # G is the labeled dataframe which should be converted into feature vectors
        >>> H = dask_extract_feature_vecs(G, features=match_f, attrs_before=['title'], attrs_after=['gold_labels'])

    """
    logger.warning(
        "WARNING THIS COMMAND IS EXPERIMENTAL AND NOT TESTED. USE AT YOUR OWN RISK.")

    # Validate input parameters

    # # We expect the input candset to be of type pandas DataFrame.
    validate_object_type(candset, pd.DataFrame, error_prefix='Input cand.set')

    # # If the attrs_before is given, Check if the attrs_before are present in
    # the input candset
    if attrs_before != None:
        if not ch.check_attrs_present(candset, attrs_before):
            logger.error(
                'The attributes mentioned in attrs_before is not present '
                'in the input table')
            raise AssertionError(
                'The attributes mentioned in attrs_before is not present '
                'in the input table')

    # # If the attrs_after is given, Check if the attrs_after are present in
    # the input candset
    if attrs_after != None:
        if not ch.check_attrs_present(candset, attrs_after):
            logger.error(
                'The attributes mentioned in attrs_after is not present '
                'in the input table')
            raise AssertionError(
                'The attributes mentioned in attrs_after is not present '
                'in the input table')

    # We expect the feature table to be a valid object
    if feature_table is None:
        logger.error('Feature table cannot be null')
        raise AssertionError('The feature table cannot be null')

    # Do metadata checking
    # # Mention what metadata is required to the user
    ch.log_info(logger, 'Required metadata: cand.set key, fk ltable, '
                        'fk rtable, '
                        'ltable, rtable, ltable key, rtable key', verbose)

    # # Get metadata
    ch.log_info(logger, 'Getting metadata from catalog', verbose)

    key, fk_ltable, fk_rtable, ltable, rtable, l_key, r_key = \
        cm.get_metadata_for_candset(
            candset, logger, verbose)

    # # Validate metadata
    ch.log_info(logger, 'Validating metadata', verbose)
    cm._validate_metadata_for_candset(candset, key, fk_ltable, fk_rtable,
                                      ltable, rtable, l_key, r_key,
                                      logger, verbose)

    # Extract features

    # id_list = [(row[fk_ltable], row[fk_rtable]) for i, row in
    #            candset.iterrows()]
    # id_list = [tuple(tup) for tup in candset[[fk_ltable, fk_rtable]].values]

    # # Set index for convenience
    l_df = ltable.set_index(l_key, drop=False)
    r_df = rtable.set_index(r_key, drop=False)

    # # Apply feature functions
    ch.log_info(logger, 'Applying feature functions', verbose)
    col_names = list(candset.columns)
    fk_ltable_idx = col_names.index(fk_ltable)
    fk_rtable_idx = col_names.index(fk_rtable)

    validate_object_type(n_chunks, int, 'Parameter n_chunks')
    validate_chunks(n_chunks)

    n_chunks = get_num_partitions(n_chunks, len(candset))

    c_splits = np.array_split(candset, n_chunks)

    pickled_obj = cloudpickle.dumps(feature_table)

    feat_vals_by_splits = []

    for i in range(len(c_splits)):
        partial_result = delayed(get_feature_vals_by_cand_split)(pickled_obj,
                                                                 fk_ltable_idx,
                                                                 fk_rtable_idx, l_df,
                                                                 r_df, c_splits[i],
                                                                 False)
        feat_vals_by_splits.append(partial_result)

    feat_vals_by_splits = delayed(wrap)(feat_vals_by_splits)
    if show_progress:
        with ProgressBar():
            feat_vals_by_splits = feat_vals_by_splits.compute(scheduler="processes",
                                                              num_workers=get_num_cores())
    else:
        feat_vals_by_splits = feat_vals_by_splits.compute(scheduler="processes",
                                                          num_workers=get_num_cores())

    feat_vals = sum(feat_vals_by_splits, [])

    # Construct output table
    feature_vectors = pd.DataFrame(feat_vals, index=candset.index.values)
    # # Rearrange the feature names in the input feature table order
    feature_names = list(feature_table['feature_name'])
    feature_vectors = feature_vectors[feature_names]

    ch.log_info(logger, 'Constructing output table', verbose)
    # print(feature_vectors)
    # # Insert attrs_before
    if attrs_before:
        if not isinstance(attrs_before, list):
            attrs_before = [attrs_before]
        attrs_before = gh.list_diff(attrs_before, [key, fk_ltable, fk_rtable])
        attrs_before.reverse()
        for a in attrs_before:
            feature_vectors.insert(0, a, candset[a])

    # # Insert keys
    feature_vectors.insert(0, fk_rtable, candset[fk_rtable])
    feature_vectors.insert(0, fk_ltable, candset[fk_ltable])
    feature_vectors.insert(0, key, candset[key])

    # # insert attrs after
    if attrs_after:
        if not isinstance(attrs_after, list):
            attrs_after = [attrs_after]
        attrs_after = gh.list_diff(attrs_after, [key, fk_ltable, fk_rtable])
        attrs_after.reverse()
        col_pos = len(feature_vectors.columns)
        for a in attrs_after:
            feature_vectors.insert(col_pos, a, candset[a])
            col_pos += 1

    # Reset the index
    # feature_vectors.reset_index(inplace=True, drop=True)

    # # Update the catalog
    cm.init_properties(feature_vectors)
    cm.copy_properties(candset, feature_vectors)

    # Finally, return the feature vectors
    return feature_vectors