# -*- coding: utf-8 -*-
"""
Created on Mon Apr  2 15:24:59 2018

@author: root
"""

import cPickle as pkl
import numpy

import cv2
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import skimage
import skimage.transform
import skimage.io
from PIL import Image, ImageEnhance
import scipy.misc

import tensorflow as tf
import numpy as np
import os
import struct
import scipy.io as sio
from array import array as pyarray
from numpy import array, int8, uint8, zeros
import collections
import pickle

import functools
import sets

from tensorflow.python.ops import rnn, array_ops
from tensorflow.contrib.rnn import GRUCell, DropoutWrapper, MultiRNNCell
from tensorflow.python import debug as tf_debug

from attn_7_1_ex import VariableSequenceClassification

from temp_createStruct5 import ConstructLookupTable
from time import gmtime, strftime

from logging_util import makelog

logfile=makelog()

class DataSet(object):

    def __init__(self, layername, numMap):
        """Construct a DataSet."""

        mat_contents = sio.loadmat('/home/titanz/Documents/SueHan/matlab/PlantClefVGG_net/RNN_plantclef/train_obs_list.mat')
        self._trainList = mat_contents['train_obs_list']

        mat_contents = sio.loadmat('/home/titanz/Documents/SueHan/matlab/PlantClefVGG_net/RNN_plantclef/train_obs_class.mat')
        self._trainLabels = mat_contents['train_obs_class']

        mat_contents = sio.loadmat('/home/titanz/Documents/SueHan/matlab/PlantClefVGG_net/RNN_plantclef/test_obs_list.mat')
        self._testList = mat_contents['test_obs_list']

        mat_contents = sio.loadmat('/home/titanz/Documents/SueHan/matlab/PlantClefVGG_net/RNN_plantclef/test_obs_class.mat')
        self._testLabels = mat_contents['test_obs_class']

        self.layerextract = layername
        self.numMap = numMap
        self._num_examples = self._trainLabels.shape[0]
        self._perm_list = np.arange(self._num_examples)
        np.random.shuffle(self._perm_list)

        self._trainLabelsPerm = self._trainLabels[self._perm_list]

        self._num_testexamples = self._testLabels.shape[0]
        self._perm_list_test = np.arange(self._num_testexamples)

        self._batch_seq = 0
        self._epochs_completed = 0
        self._index_in_epoch = 0
        self._index_in_epoch_test = 0
        self._max_seq = 0

        self.Batch_Up_model = ConstructLookupTable()
        self.mydict2_test256 = self.Batch_Up_model.main(self._testList,2) # for train_testID ! = 1
        self.feature_size_conv = self.numMap*14*14
        self.feature_size_fc = 4096

    def trainList(self):
        return self._trainList

    def trainLabels(self):
        return self._trainLabels

    def trainLabelsPerm(self):
        return self._trainLabelsPerm

    def testList(self):
        return self._testList

    def testLabels(self):
        return self._testLabels

    def num_examples(self):
        return self._num_examples

    def num_testexamples(self):
        return self._num_testexamples

    def epochs_completed(self):
        return self._epochs_completed

    def index_in_epoch(self):
        return self._index_in_epoch

    def max_seq(self):
        return self._max_seq

    def batch_seq(self):
        return self._batch_seq

    def PrepareTrainingBatch(self,Newpermbatch, batch_size, indicator):
        if indicator == 1:
           mydictG = self.Batch_Up_model.main(self._trainList,1) # for train_testID == 1
        else:
           mydictG = self.mydict2_test256

        i = 0
        temp = np.zeros(batch_size)
        while i < batch_size:
              temp[i] = len(mydictG[Newpermbatch[i]][1])
              i = i + 1
        self._max_seq = int(np.amax(temp))
        self._batch_seq = temp

        batch_conv = np.zeros([batch_size,self._max_seq,self.feature_size_conv])
        batch_fc = np.zeros([batch_size,self._max_seq,self.feature_size_fc])

        i = 0
        while i < batch_size:

              media_length = len(mydictG[Newpermbatch[i]][1])
              j = 0
              while j < media_length:
                    ### for 256 image size for testing
#                    pkl_file1 = open(mydictG[Newpermbatch[i]][1][j][0], 'rb')
#                    output = pickle.load(pkl_file1)
#                    pkl_file1.close()
#
#                    pkl_file2 = open(mydictG[Newpermbatch[i]][1][j][1], 'rb')
#                    output2 = pickle.load(pkl_file2)
#                    pkl_file2.close()
#
#                    pkl_file3 = open(mydictG[Newpermbatch[i]][1][j][2], 'rb')
#                    output3 = pickle.load(pkl_file3)
#                    pkl_file3.close()
#
#                    output.update(output2)
#                    output.update(output3)
#                    mat_contents = output[self.layerextract[0]]
#                    batch_conv[i][j][:] = mat_contents.reshape(self.feature_size_conv)  #'conv5_3'
#
#                    mat_contents = output[self.layerextract[1]]
##                    batch_fc[i][j][:] = mat_contents.reshape(self.feature_size_conv)    #'conv5_3_O'
#                    batch_fc[i][j][:] = mat_contents  #'convfc7'
#
#                    j = j + 1

######################################################################

                  ## for 384,512 image size for testing
                  if indicator == 1:   # training    ###################
                    pkl_file1 = open(mydictG[Newpermbatch[i]][1][j][0], 'rb')
                    output = pickle.load(pkl_file1)
                    pkl_file1.close()

                    pkl_file2 = open(mydictG[Newpermbatch[i]][1][j][1], 'rb')
                    output2 = pickle.load(pkl_file2)
                    pkl_file2.close()

                    pkl_file3 = open(mydictG[Newpermbatch[i]][1][j][2], 'rb')
                    output3 = pickle.load(pkl_file3)
                    pkl_file3.close()

                    output.update(output2)
                    output.update(output3)
                    mat_contents = output[self.layerextract[0]]
                    batch_conv[i][j][:] = mat_contents.reshape(self.feature_size_conv)   #'conv5_3'

                    mat_contents = output[self.layerextract[1]]
                    batch_fc[i][j][:] = mat_contents.reshape(self.feature_size_fc)   #'conv5_3_O'

                    j = j + 1

                  else:   # testing

                    pkl_file1 = open(mydictG[Newpermbatch[i]][1][j][0], 'rb')
                    output = pickle.load(pkl_file1)
                    pkl_file1.close()

                    pkl_file2 = open(mydictG[Newpermbatch[i]][1][j][1], 'rb')
                    output2 = pickle.load(pkl_file2)
                    pkl_file2.close()

                    output.update(output2)
                    mat_contents = output[self.layerextract[0]]
                    batch_conv[i][j][:] = mat_contents.reshape(self.feature_size_conv)    #'conv5_3'

                    mat_contents = output[self.layerextract[1]]
                    batch_fc[i][j][:] = mat_contents.reshape(self.feature_size_fc)   #'conv5_3_O'

                    j = j + 1

#########################################################
              if indicator == 1:
                 J = np.arange(media_length)
                 np.random.shuffle(J)

                 temp_arr = batch_conv[i,:media_length,:]
                 temp_arr = temp_arr[J,:]
                 batch_conv[i,:media_length,:] = temp_arr

                 temp_arr = batch_fc[i,:media_length,:]
                 temp_arr = temp_arr[J,:]
                 batch_fc[i,:media_length,:] = temp_arr

              i = i + 1

        return batch_fc, batch_conv

    def dense_to_one_hot(self,labels_dense, num_classes=1000):
        labels_dense = labels_dense.astype(int)
        num_labels = labels_dense.shape[0]
        index_offset = np.arange(num_labels) * num_classes
        labels_one_hot = np.zeros((num_labels, num_classes))
        labels_one_hot.flat[index_offset + labels_dense.ravel()] = 1
        labels_one_hot = labels_one_hot.astype(np.float32)

        temp = zeros((labels_one_hot.shape[0],self._max_seq,num_classes))
        i=0
        while i < labels_one_hot.shape[0]:
              temp[i][0:int(self._batch_seq[i])] = labels_one_hot[i]
              i=i+1

        return temp

    def next_batch(self,batch_size):
        start = self._index_in_epoch
        self._index_in_epoch += batch_size
        if self._index_in_epoch > self._num_examples:
            # Finished epoch
            self._epochs_completed += 1
            # Shuffle the data
            self._perm_list = np.arange(self._num_examples)
            np.random.shuffle(self._perm_list)
            self._trainLabelsPerm = self._trainLabels[self._perm_list]
            # Start next epoch
            start = 0
            self._index_in_epoch = batch_size
            assert batch_size <= self._num_examples
        end = self._index_in_epoch

        return self.PrepareTrainingBatch(self._perm_list[start:end], batch_size, 1), self.dense_to_one_hot(self._trainLabelsPerm[start:end])

    def PrepareTestingBatch(self,test_total):
        start = self._index_in_epoch_test
        self._index_in_epoch_test += test_total
        if self._index_in_epoch_test > self._num_testexamples:
            start = 0
            self._index_in_epoch_test = test_total
            assert test_total <= self._num_testexamples
        end = self._index_in_epoch_test

        return self.PrepareTrainingBatch(self._perm_list_test[start:end], test_total, 0), self.dense_to_one_hot(self._testLabels[start:end])

     ## Testing
    def Reset_index_in_epoch_test(self, init_v = 0):
        self._index_in_epoch_test = init_v

    def crop_image(self, x, target_height=224, target_width=224):
        image = skimage.img_as_float(skimage.io.imread(x)).astype(np.float32)

        if len(image.shape) == 2:
            image = np.tile(image[:,:,None], 3)
        elif len(image.shape) == 4:
            image = image[:,:,:,0]

        height, width, rgb = image.shape
        if width == height:
            resized_image = cv2.resize(image, (target_height,target_width))

        elif height < width:
            resized_image = cv2.resize(image, (int(width * float(target_height)/height), target_width))
            cropping_length = int((resized_image.shape[1] - target_height) / 2)
            resized_image = resized_image[:,cropping_length:resized_image.shape[1] - cropping_length]

        else:
            resized_image = cv2.resize(image, (target_height, int(height * float(target_width) / width)))
            cropping_length = int((resized_image.shape[0] - target_width) / 2)
            resized_image = resized_image[cropping_length:resized_image.shape[0] - cropping_length,:]

        return cv2.resize(resized_image, (target_height, target_width))

####### Network Parameters ########
training_iters = 10000000 # run 10000 epoch
batch_size =  15
display_step = 280  #280
test_num_total = 15
layername_conv = 'conv5_3'
layername_fc = 'fc7_final'
layername = [layername_conv, layername_fc]
numMap = 512
featMap = 14*14
num_classes = 1000
dropTrain = 0.5
dropTest = 1
prob_path = '/media/titanz/Data3TB/tensorboard_log/model_20180211/prob_256/'
savefigfile ='/media/titanz/Data3TB/tensorboard_log/model_20180309/attn_visual_imsz384/'
#################################

plantclefdata = DataSet(layername,numMap)

# tf Graph input
x = tf.placeholder("float", [None, None, 4096])
data = tf.placeholder("float", [None, None, numMap*14*14])
target = tf.placeholder("float", [None, None, num_classes])
dropout = tf.placeholder(tf.float32)
batch_size2 = tf.placeholder(tf.int32)

model = VariableSequenceClassification(x, data, target, dropout, batch_size2)
sess = tf.InteractiveSession()
#sess = tf.Session()
#sess = tf_debug.LocalCLIDebugWrapperSession(sess)

saver = tf.train.Saver(max_to_keep = None)
saver.restore(sess, "/media/titanz/Data3TB/tensorboard_log/model_20180309/model_55160")

#################################################################################
mat_contents = sio.loadmat('/home/titanz/Documents/SueHan/matlab/PlantClefVGG_net/RNN_plantclef/test_obs_list.mat')
testList = mat_contents['test_obs_list']

mat_contents = sio.loadmat('/media/titanz/Data3TB/tensorboard_log/model_20180309/obs_tp_media_re.mat')
obs_tp_media_re = mat_contents['obs_tp_media_re']

imagesz = 1   # choose image size 1 = 256 , 2 = 384, 3 = 512
if imagesz == 1:  #256
    path_folder = '/media/titanz/Data3TB/PlantclefVGA2/PlantClefImageTest/resize_species/species/'
elif imagesz == 2:  #384
    path_folder = '/media/titanz/data3TB_02/PlantClefolddata_augmented/PlantClefImageTest_SR/resize_species_384/'
else: #512
    path_folder = '/media/titanz/data3TB_02/PlantClefolddata_augmented/PlantClefImageTest_SR/resize_species_512/'

smooth = True

# read python dict back from the testing_256 file
pkl_file_test = open('/home/titanz/tensor_flow/tensorflow-master/tensorflow/examples/RNN/myfile_test_256.pkl', 'rb')
mydict2_test = pickle.load(pkl_file_test)
pkl_file_test.close()

mat_contents = sio.loadmat('/media/titanz/Data3TB/tensorflowlist/species/ClassID_CNNID.mat')
classIDList = mat_contents['ClassID_CNNID']

mediafolderTest_content = '/media/titanz/Data3TB/tensorflowlist/VGG_multipath_res_bn_lastconv/test_obs_media_content_256/'

mat_contents = sio.loadmat('/home/titanz/Documents/SueHan/matlab/PlantClefVGG_net/RNN_plantclef/test_obs_class.mat')
testLabels = mat_contents['test_obs_class']
################################################################################

for count in xrange(obs_tp_media_re.shape[0]):

    print('Obs num {:7.0f}'.format(count))

    ObsIDchosen = obs_tp_media_re[count][0].astype(int) - 1#8258#12# chose 0 <= x < 13887

    Obs_name = testList[ObsIDchosen].astype(int)
    Obs_name = str(Obs_name).split('[')
    Obs_name = Obs_name[1].split(']')
    directory = savefigfile + str(Obs_name[0])
    if not os.path.exists(directory):
        os.makedirs(directory)

    plantclefdata.Reset_index_in_epoch_test(init_v = ObsIDchosen)

    (test_data_x, test_data_conv), test_label = plantclefdata.PrepareTestingBatch(test_num_total)

    pred, alpha_forward1, alpha_forward2, alpha_forward3, alpha_backward1, alpha_backward2, alpha_backward3 = sess.run(model.alpha_list_com, feed_dict={x: test_data_x, data: test_data_conv, batch_size2: test_num_total, target: test_label, dropout: dropTest})#, batch_size: batch_size})
    pred_re = pred[0,:,:]
    B = np.argmax(pred_re,axis=1)
    alpha_forward1 = np.array(alpha_forward1).swapaxes(1,0)  # alpha(max_seq, batch, 196)
    alpha_forward2 = np.array(alpha_forward2).swapaxes(1,0)  # alpha(max_seq, batch, 196)
    alpha_forward3 = np.array(alpha_forward3).swapaxes(1,0)  # alpha(max_seq, batch, 196)
    mat_contents2 = sio.loadmat(mediafolderTest_content + str(mydict2_test[ObsIDchosen][0]) + '.mat',mat_dtype=True)
    used = mat_contents2['v']
    alphas1 = np.array(alpha_forward1).swapaxes(1,0)  # alpha(max_seq, batch, 196)
    alphas1 = alphas1[0:used.shape[1],:,:]
    alphas2 = np.array(alpha_forward2).swapaxes(1,0)  # alpha(max_seq, batch, 196)
    alphas2 = alphas2[0:used.shape[1],:,:]
    alphas3 = np.array(alpha_forward3).swapaxes(1,0)  # alpha(max_seq, batch, 196)
    alphas3 = alphas3[0:used.shape[1],:,:]
    pred_re = pred[0,:,:]
    pred_re = pred_re[0:used.shape[1],:]
    B = np.argmax(pred_re,axis=1)
    B  = B[0:used.shape[1]]

    class_picken = testLabels[ObsIDchosen]
    class_picken = class_picken.astype(int)
    index_plot = 1;
    index_plotnc = 1;
    for ii in xrange(alphas1.shape[0]):    # eg: 0,1,2  #list(range(0,alphas.shape[0]*2,2)):
       organlabel = int(used[0,ii])
       if organlabel == 0:
            organlabelD = 'Branch'
       elif organlabel == 1:
            organlabelD = 'Entire'
       elif organlabel == 2:
            organlabelD = 'Flower'
       elif organlabel == 3:
            organlabelD = 'Fruit'
       elif organlabel == 4:
            organlabelD = 'Leaf'
       elif organlabel == 5:
            organlabelD = 'LeafScan'
       else:
            organlabelD = 'Stem'

       plt.figure(1)
       L =  mydict2_test[ObsIDchosen][1][ii].split('/')
       name_str = L[len(L)-1]
       name_str2 = name_str.split('.mat')
       name_str3 = name_str2[0]
       path = path_folder + '{:04d}'.format(class_picken[0]) + '-' + str(classIDList[class_picken[0]][0]) +'/' + name_str3
       img = plantclefdata.crop_image(path)
       plt.imshow(img)
       if smooth:
           alpha_img = skimage.transform.pyramid_expand(alphas1[ii,0,:].reshape(14,14), upscale=16, sigma=20)
       else:
           alpha_img = skimage.transform.resize(alphas1[ii,0,:].reshape(14,14), [img.shape[0], img.shape[1]])

       plt.imshow(alpha_img, alpha=0.7)
       plt.set_cmap(cm.Greys_r)
       plt.axis('off')
       plt.savefig(directory + '/' + "course" + str(ii) + ".png")
       plt.imshow(img)
       plt.axis('off')
       lab2 = organlabelD + '_'+ str(class_picken[0]) + '-' + str(B[ii])
       plt.savefig(directory + '/'  + lab2 + '_' + str(ii) + ".png")
       plt.figure(2)
       plt.imshow(img)

       if smooth:
           alpha_img = skimage.transform.pyramid_expand(alphas2[ii,0,:].reshape(14,14), upscale=16, sigma=20)
       else:
           alpha_img = skimage.transform.resize(alphas2[ii,0,:].reshape(14,14), [img.shape[0], img.shape[1]])

       plt.imshow(alpha_img, alpha=0.7)    # show attn
       plt.set_cmap(cm.Greys_r)
       plt.axis('off')
       plt.savefig(directory + '/'  + "fine" + str(ii) + ".png")


"""
The :mod:`sklearn.datasets` module includes utilities to load datasets,
including methods to load and fetch popular reference datasets. It also
features some artificial data generators.
"""

from .base import load_diabetes
from .base import load_digits
from .base import load_files
from .base import load_iris
from .base import load_breast_cancer
from .base import load_linnerud
from .base import load_boston
from .base import get_data_home
from .base import clear_data_home
from .base import load_sample_images
from .base import load_sample_image
from .covtype import fetch_covtype
from .kddcup99 import fetch_kddcup99
from .mlcomp import load_mlcomp
from .lfw import load_lfw_pairs
from .lfw import load_lfw_people
from .lfw import fetch_lfw_pairs
from .lfw import fetch_lfw_people
from .twenty_newsgroups import fetch_20newsgroups
from .twenty_newsgroups import fetch_20newsgroups_vectorized
from .mldata import fetch_mldata, mldata_filename
from .samples_generator import make_classification
from .samples_generator import make_multilabel_classification
from .samples_generator import make_hastie_10_2
from .samples_generator import make_regression
from .samples_generator import make_blobs
from .samples_generator import make_moons
from .samples_generator import make_circles
from .samples_generator import make_friedman1
from .samples_generator import make_friedman2
from .samples_generator import make_friedman3
from .samples_generator import make_low_rank_matrix
from .samples_generator import make_sparse_coded_signal
from .samples_generator import make_sparse_uncorrelated
from .samples_generator import make_spd_matrix
from .samples_generator import make_swiss_roll
from .samples_generator import make_s_curve
from .samples_generator import make_sparse_spd_matrix
from .samples_generator import make_gaussian_quantiles
from .samples_generator import make_biclusters
from .samples_generator import make_checkerboard
from .svmlight_format import load_svmlight_file
from .svmlight_format import load_svmlight_files
from .svmlight_format import dump_svmlight_file
from .olivetti_faces import fetch_olivetti_faces
from .species_distributions import fetch_species_distributions
from .california_housing import fetch_california_housing
from .rcv1 import fetch_rcv1

__all__ = ['clear_data_home',
           'dump_svmlight_file',
           'fetch_20newsgroups',
           'fetch_20newsgroups_vectorized',
           'fetch_lfw_pairs',
           'fetch_lfw_people',
           'fetch_mldata',
           'fetch_olivetti_faces',
           'fetch_species_distributions',
           'fetch_california_housing',
           'fetch_covtype',
           'fetch_rcv1',
           'fetch_kddcup99',
           'get_data_home',
           'load_boston',
           'load_diabetes',
           'load_digits',
           'load_files',
           'load_iris',
           'load_breast_cancer',
           'load_lfw_pairs',
           'load_lfw_people',
           'load_linnerud',
           'load_mlcomp',
           'load_sample_image',
           'load_sample_images',
           'load_svmlight_file',
           'load_svmlight_files',
           'make_biclusters',
           'make_blobs',
           'make_circles',
           'make_classification',
           'make_checkerboard',
           'make_friedman1',
           'make_friedman2',
           'make_friedman3',
           'make_gaussian_quantiles',
           'make_hastie_10_2',
           'make_low_rank_matrix',
           'make_moons',
           'make_multilabel_classification',
           'make_regression',
           'make_s_curve',
           'make_sparse_coded_signal',
           'make_sparse_spd_matrix',
           'make_sparse_uncorrelated',
           'make_spd_matrix',
           'make_swiss_roll',
           'mldata_filename']


"""
================================================================
Plot the decision surface of a decision tree on the iris dataset
================================================================

Plot the decision surface of a decision tree trained on pairs
of features of the iris dataset.

See :ref:`decision tree <tree>` for more information on the estimator.

For each pair of iris features, the decision tree learns decision
boundaries made of combinations of simple thresholding rules inferred from
the training samples.
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier

# Parameters
n_classes = 3
plot_colors = "bry"
plot_step = 0.02

# Load data
iris = load_iris()

for pairidx, pair in enumerate([[0, 1], [0, 2], [0, 3],
                                [1, 2], [1, 3], [2, 3]]):
    # We only take the two corresponding features
    X = iris.data[:, pair]
    y = iris.target

    # Shuffle
    idx = np.arange(X.shape[0])
    np.random.seed(13)
    np.random.shuffle(idx)
    X = X[idx]
    y = y[idx]

    # Standardize
    mean = X.mean(axis=0)
    std = X.std(axis=0)
    X = (X - mean) / std

    # Train
    clf = DecisionTreeClassifier().fit(X, y)

    # Plot the decision boundary
    plt.subplot(2, 3, pairidx + 1)

    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 = clf.predict(np.c_[xx.ravel(), yy.ravel()])
    Z = Z.reshape(xx.shape)
    cs = plt.contourf(xx, yy, Z, cmap=plt.cm.Paired)

    plt.xlabel(iris.feature_names[pair[0]])
    plt.ylabel(iris.feature_names[pair[1]])
    plt.axis("tight")

    # Plot the training points
    for i, color in zip(range(n_classes), plot_colors):
        idx = np.where(y == i)
        plt.scatter(X[idx, 0], X[idx, 1], c=color, label=iris.target_names[i],
                    cmap=plt.cm.Paired)

    plt.axis("tight")

plt.suptitle("Decision surface of a decision tree using paired features")
plt.legend()
plt.show()


# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2013-2015, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import csv
from optparse import OptionParser

from matplotlib import pyplot as plt
import numpy as np

from hpelm import ELM

from swarm_runner import SwarmRunner
from scipy import random

import pandas as pd
from htmresearch.support.sequence_learning_utils import *

from htmresearch.algorithms.online_extreme_learning_machine import OSELM

plt.ion()

def initializeELMnet(nDimInput, nDimOutput, numNeurons=10):
  # Build ELM network with nDim input units,
  # numNeurons hidden units (LSTM cells) and nDimOutput cells

  # net = ELM(nDimInput, nDimOutput)
  # net.add_neurons(numNeurons, "sigm")

  net = OSELM(nDimInput, nDimOutput,
            numHiddenNeurons=numNeurons, activationFunction='sig')

  return net

def readDataSet(dataSet):
  filePath = 'data/'+dataSet+'.csv'
  # df = pd.read_csv(filePath, header=0, skiprows=[1, 2], names=['time', 'data'])
  # sequence = df['data']

  if dataSet=='nyc_taxi':
    df = pd.read_csv(filePath, header=0, skiprows=[1,2],
                     names=['time', 'data', 'timeofday', 'dayofweek'])
    sequence = df['data']
    dayofweek = df['dayofweek']
    timeofday = df['timeofday']

    seq = pd.DataFrame(np.array(pd.concat([sequence, timeofday, dayofweek], axis=1)),
                        columns=['data', 'timeofday', 'dayofweek'])
  elif dataSet=='sine':
    df = pd.read_csv(filePath, header=0, skiprows=[1, 2], names=['time', 'data'])
    sequence = df['data']
    seq = pd.DataFrame(np.array(sequence), columns=['data'])
  else:
    raise(' unrecognized dataset type ')

  return seq

def getTimeEmbeddedMatrix(sequence, numLags=100, predictionStep=1,
                      useTimeOfDay=True, useDayOfWeek=True):
  print "generate time embedded matrix "
  print "the training data contains ", str(nTrain-predictionStep), "records"

  inDim = numLags + int(useTimeOfDay) + int(useDayOfWeek)

  if useTimeOfDay:
    print "include time of day as input field"
  if useDayOfWeek:
    print "include day of week as input field"

  X = np.zeros(shape=(len(sequence), inDim))
  T = np.zeros(shape=(len(sequence), 1))
  for i in xrange(numLags-1, len(sequence)-predictionStep):
    if useTimeOfDay and useDayOfWeek:
      sample = np.concatenate([np.array(sequence['data'][(i-numLags+1):(i+1)]),
                               np.array([sequence['timeofday'][i]]),
                               np.array([sequence['dayofweek'][i]])])
    elif useTimeOfDay:
      sample = np.concatenate([np.array(sequence['data'][(i-numLags+1):(i+1)]),
                               np.array([sequence['timeofday'][i]])])
    elif useDayOfWeek:
      sample = np.concatenate([np.array(sequence['data'][(i-numLags+1):(i+1)]),
                               np.array([sequence['dayofweek'][i]])])
    else:
      sample = np.array(sequence['data'][(i-numLags+1):(i+1)])

    X[i, :] = sample
    T[i, :] = sequence['data'][i+predictionStep]

  return (X, T)

def _getArgs():
  parser = OptionParser(usage="%prog PARAMS_DIR OUTPUT_DIR [options]"
                              "\n\nCompare TM performance with trivial predictor using "
                              "model outputs in prediction directory "
                              "and outputting results to result directory.")
  parser.add_option("-d",
                    "--dataSet",
                    type=str,
                    default='nyc_taxi',
                    dest="dataSet",
                    help="DataSet Name, choose from sine, SantaFe_A, MackeyGlass")

  # parser.add_option("-n",
  #                   "--predictionstep",
  #                   type=int,
  #                   default=1,
  #                   dest="predictionstep",
  #                   help="number of steps ahead to be predicted")

  (options, remainder) = parser.parse_args()
  print options

  return options, remainder

def saveResultToFile(dataSet, predictedInput, algorithmName):
  inputFileName = 'data/' + dataSet + '.csv'
  inputFile = open(inputFileName, "rb")

  csvReader = csv.reader(inputFile)

  # skip header rows
  csvReader.next()
  csvReader.next()
  csvReader.next()

  outputFileName = './prediction/' + dataSet + '_' + algorithmName + '_pred.csv'
  outputFile = open(outputFileName, "w")
  csvWriter = csv.writer(outputFile)
  csvWriter.writerow(
    ['timestamp', 'data', 'prediction-' + str(predictionStep) + 'step'])
  csvWriter.writerow(['datetime', 'float', 'float'])
  csvWriter.writerow(['', '', ''])

  for i in xrange(len(sequence)):
    row = csvReader.next()
    csvWriter.writerow([row[0], row[1], predictedInput[i]])

  inputFile.close()
  outputFile.close()

if __name__ == "__main__":

  (_options, _args) = _getArgs()
  dataSet = _options.dataSet

  print "run ELM on ", dataSet
  SWARM_CONFIG = SwarmRunner.importSwarmDescription(dataSet)
  predictedField = SWARM_CONFIG['inferenceArgs']['predictedField']
  nTrain = SWARM_CONFIG["streamDef"]['streams'][0]['last_record']
  predictionStep = SWARM_CONFIG['inferenceArgs']['predictionSteps'][0]

  useTimeOfDay = True
  useDayOfWeek = True

  nTrain = 500
  numLags = 100

  # prepare dataset as pyBrain sequential dataset
  sequence = readDataSet(dataSet)

  # standardize data by subtracting mean and dividing by std
  meanSeq = np.mean(sequence['data'])
  stdSeq = np.std(sequence['data'])
  sequence['data'] = (sequence['data'] - meanSeq)/stdSeq

  meanTimeOfDay = np.mean(sequence['timeofday'])
  stdTimeOfDay = np.std(sequence['timeofday'])
  sequence['timeofday'] = (sequence['timeofday'] - meanTimeOfDay)/stdTimeOfDay

  meanDayOfWeek = np.mean(sequence['dayofweek'])
  stdDayOfWeek = np.std(sequence['dayofweek'])
  sequence['dayofweek'] = (sequence['dayofweek'] - meanDayOfWeek)/stdDayOfWeek

  (X, T) = getTimeEmbeddedMatrix(sequence, numLags, predictionStep,
                                 useTimeOfDay, useDayOfWeek)

  random.seed(6)
  net = initializeELMnet(nDimInput=X.shape[1],
                         nDimOutput=1, numNeurons=50)

  net.initializePhase(X[:nTrain, :], T[:nTrain, :])

  predictedInput = np.zeros((len(sequence),))
  targetInput = np.zeros((len(sequence),))
  trueData = np.zeros((len(sequence),))
  for i in xrange(nTrain, len(sequence)-predictionStep):
    net.train(X[[i], :], T[[i], :])
    Y = net.predict(X[[i], :])

    predictedInput[i] = Y[-1]
    targetInput[i] = sequence['data'][i+predictionStep]
    trueData[i] = sequence['data'][i]
    print "Iteration {} target input {:2.2f} predicted Input {:2.2f} ".format(
      i, targetInput[i], predictedInput[i])

  predictedInput = (predictedInput * stdSeq) + meanSeq
  targetInput = (targetInput * stdSeq) + meanSeq
  trueData = (trueData * stdSeq) + meanSeq

  saveResultToFile(dataSet, predictedInput, 'elm')

  plt.figure()
  plt.plot(targetInput)
  plt.plot(predictedInput)
  plt.xlim([12800, 13500])
  plt.ylim([0, 30000])

  skipTrain = 6000
  from plot import computeSquareDeviation
  squareDeviation = computeSquareDeviation(predictedInput, targetInput)
  squareDeviation[:skipTrain] = None
  nrmse = np.sqrt(np.nanmean(squareDeviation)) / np.nanstd(targetInput)
  print "NRMSE {}".format(nrmse)


# Author: Brian M. Clapper, G Varoquaux
# License: BSD

import numpy as np

# XXX we should be testing the public API here
from sklearn.utils.linear_assignment_ import _hungarian

def test_hungarian():
    matrices = [
        # Square
        ([[400, 150, 400],
          [400, 450, 600],
          [300, 225, 300]],
         850  # expected cost
         ),

        # Rectangular variant
        ([[400, 150, 400, 1],
          [400, 450, 600, 2],
          [300, 225, 300, 3]],
         452  # expected cost
         ),

        # Square
        ([[10, 10,  8],
          [9,  8,  1],
          [9,  7,  4]],
         18
         ),

        # Rectangular variant
        ([[10, 10,  8, 11],
          [9, 8, 1, 1],
          [9, 7, 4, 10]],
         15
         ),

        # n == 2, m == 0 matrix
        ([[], []],
         0
         ),
    ]

    for cost_matrix, expected_total in matrices:
        cost_matrix = np.array(cost_matrix)
        indexes = _hungarian(cost_matrix)
        total_cost = 0
        for r, c in indexes:
            x = cost_matrix[r, c]
            total_cost += x
        assert expected_total == total_cost

        indexes = _hungarian(cost_matrix.T)
        total_cost = 0
        for c, r in indexes:
            x = cost_matrix[r, c]
            total_cost += x
        assert expected_total == total_cost


"""
Visualize Regularization Path
=============================

Plot the edge level coefficients (inverse covariance entries)
as a function of the regularization parameter.

"""

import sys
import numpy as np
from sklearn.datasets import make_sparse_spd_matrix

sys.path.append("..")
from inverse_covariance import QuicGraphicalLasso
from inverse_covariance.plot_util import trace_plot
from inverse_covariance.profiling import LatticeGraph

def make_data(n_samples, n_features):
    prng = np.random.RandomState(1)
    prec = make_sparse_spd_matrix(
        n_features, alpha=.98, smallest_coef=.4, largest_coef=.7, random_state=prng
    )
    cov = np.linalg.inv(prec)
    d = np.sqrt(np.diag(cov))
    cov /= d
    cov /= d[:, np.newaxis]
    prec *= d
    prec *= d[:, np.newaxis]
    X = prng.multivariate_normal(np.zeros(n_features), cov, size=n_samples)
    X -= X.mean(axis=0)
    X /= X.std(axis=0)
    return X, cov, prec

def make_data_banded(n_samples, n_features):
    alpha = 0.1
    cov, prec, adj = LatticeGraph(
        n_blocks=2, random_sign=True, chain_blocks=True, seed=1
    ).create(n_features, alpha)
    prng = np.random.RandomState(2)
    X = prng.multivariate_normal(np.zeros(n_features), cov, size=n_samples)
    return X, cov, prec

def show_quic_coefficient_trace(X):
    path = np.logspace(np.log10(0.01), np.log10(1.0), num=50, endpoint=True)[::-1]
    estimator = QuicGraphicalLasso(lam=1.0, path=path, mode="path")
    estimator.fit(X)
    trace_plot(estimator.precision_, estimator.path_, n_edges=20)

def show_quic_coefficient_trace_truth(X, truth):
    path = np.logspace(np.log10(0.01), np.log10(1.0), num=50, endpoint=True)[::-1]
    estimator = QuicGraphicalLasso(lam=1.0, path=path, mode="path")
    estimator.fit(X)
    trace_plot(estimator.precision_, estimator.path_, n_edges=6, ground_truth=truth)

if __name__ == "__main__":
    # example 1
    n_samples = 10
    n_features = 5
    X, cov, prec = make_data(n_samples, n_features)

    print("Showing basic Erdos-Renyi example with ")
    print("   n_samples=10")
    print("   n_features=5")
    print("   n_edges=20")
    show_quic_coefficient_trace(X)

    # use ground truth for display
    print("Showing basic Erdos-Renyi example with ")
    print("   n_samples=100")
    print("   n_features=5")
    print("   n_edges=6")
    print("   ground_truth (shows only false pos and negatives)")
    show_quic_coefficient_trace_truth(X, prec)

    # example 2
    n_samples = 110
    n_features = 100
    X, cov, prec = make_data_banded(n_samples, n_features)

    print("Showing basic Lattice example with ")
    print("   n_samples=110")
    print("   n_features=100")
    print("   n_blocks=2")
    print("   random_sign=True")
    print("   n_edges=20")
    show_quic_coefficient_trace(X)

    # use ground truth for display
    print("Showing basic Lattice example with ")
    print("   n_samples=110")
    print("   n_features=100")
    print("   n_blocks=2")
    print("   random_sign=True")
    print("   n_edges=6")
    print("   ground_truth (shows only false pos and negatives)")
    show_quic_coefficient_trace_truth(X, prec)


#!/usr/bin/env python
#
# PyUSBtmc
# display_channel.py
#
# Copyright (c) 2011 Mike Hadmack
# Copyright (c) 2010 Matt Mets
# This code is distributed under the MIT license
#
#  This script is just to test waverunner functionality as a module
#
import numpy
from matplotlib import pyplot
import sys
import os
sys.path.append(os.path.expanduser('.'))
from waverunner import Waverunner

""" Example program to plot the Y-T data from one scope channel
    derived from capture_channel_1.py but using new interface methods """

# Initialize our scope
scope = Waverunner("127.0.0.1")
scope.grabData()
data1 = scope.getScaledWaveform(1)
data2 = scope.getScaledWaveform(2)

# Now, generate a time axis.
time = scope.getTimeAxis()

# See if we should use a different time axis
if (time[599] < 1e-3):
    time = time * 1e6
    tUnit = "uS"
elif (time[599] < 1):
    time = time * 1e3
    tUnit = "mS"
else:
    tUnit = "S"

# close interface
scope.close()

# Plot the data
pyplot.plot(time,data1)
pyplot.plot(time,data2)
pyplot.title("Oscilloscope Data")
pyplot.ylabel("Voltage (V)")
pyplot.xlabel("Time (" + tUnit + ")")
pyplot.xlim(time[0], time[599])
pyplot.show()