import matplotlib.pyplot as plt
import random
import pickle
from skimage.transform import rotate
from scipy import ndimage
from skimage.util import img_as_ubyte
from joblib import Parallel, delayed
from sklearn.ensemble.forest import _generate_unsampled_indices
from sklearn.ensemble.forest import _generate_sample_indices
import numpy as np
from sklearn.ensemble import BaggingClassifier
from sklearn.tree import DecisionTreeClassifier
from itertools import product
import keras

from keras import layers

from joblib import Parallel, delayed
from multiprocessing import Pool

import tensorflow as tf

from numba import cuda

import sys
sys.path.append("../../proglearn/")
from progressive_learner import ProgressiveLearner
from deciders import SimpleArgmaxAverage
from transformers import TreeClassificationTransformer, NeuralClassificationTransformer
from voters import TreeClassificationVoter, KNNClassificationVoter

def cross_val_data(data_x, data_y, total_cls=10):
    x = data_x.copy()
    y = data_y.copy()
    idx = [np.where(data_y == u)[0] for u in np.unique(data_y)]


    for i in range(total_cls):
        indx = idx[i]#np.roll(idx[i],(cv-1)*100)
        random.shuffle(indx)

        if i==0:
            train_x1 = x[indx[0:250],:]
            train_x2 = x[indx[250:500],:]
            train_y1 = y[indx[0:250]]
            train_y2 = y[indx[250:500]]

            test_x = x[indx[500:600],:]
            test_y = y[indx[500:600]]
        else:
            train_x1 = np.concatenate((train_x1, x[indx[0:250],:]), axis=0)
            train_x2 = np.concatenate((train_x2, x[indx[250:500],:]), axis=0)
            train_y1 = np.concatenate((train_y1, y[indx[0:250]]), axis=0)
            train_y2 = np.concatenate((train_y2, y[indx[250:500]]), axis=0)

            test_x = np.concatenate((test_x, x[indx[500:600],:]), axis=0)
            test_y = np.concatenate((test_y, y[indx[500:600]]), axis=0)


    return train_x1, train_y1, train_x2, train_y2, test_x, test_y

def LF_experiment(data_x, data_y, angle, model, granularity, reps=1, ntrees=29, acorn=None):
    if acorn is not None:
        np.random.seed(acorn)

    errors = np.zeros(2)

    for rep in range(reps):
        print("Starting Rep {} of Angle {}".format(rep, angle))
        train_x1, train_y1, train_x2, train_y2, test_x, test_y = cross_val_data(data_x, data_y, total_cls=10)


        #change data angle for second task
        tmp_data = train_x2.copy()
        _tmp_ = np.zeros((32,32,3), dtype=int)
        total_data = tmp_data.shape[0]

        for i in range(total_data):
            tmp_ = image_aug(tmp_data[i],angle)
            tmp_data[i] = tmp_

        if model == "uf":
            train_x1 = train_x1.reshape((train_x1.shape[0], train_x1.shape[1] * train_x1.shape[2] * train_x1.shape[3]))
            tmp_data = tmp_data.reshape((tmp_data.shape[0], tmp_data.shape[1] * tmp_data.shape[2] * tmp_data.shape[3]))
            test_x = test_x.reshape((test_x.shape[0], test_x.shape[1] * test_x.shape[2] * test_x.shape[3]))

        with tf.device('/gpu:'+str(int(angle //  granularity) % 4)):
            default_transformer_class = NeuralClassificationTransformer

            network = keras.Sequential()
            network.add(layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu', input_shape=np.shape(train_x1)[1:]))
            network.add(layers.BatchNormalization())
            network.add(layers.Conv2D(filters=32, kernel_size=(3, 3), strides = 2, padding = "same", activation='relu'))
            network.add(layers.BatchNormalization())
            network.add(layers.Conv2D(filters=64, kernel_size=(3, 3), strides = 2, padding = "same", activation='relu'))
            network.add(layers.BatchNormalization())
            network.add(layers.Conv2D(filters=128, kernel_size=(3, 3), strides = 2, padding = "same", activation='relu'))
            network.add(layers.BatchNormalization())
            network.add(layers.Conv2D(filters=254, kernel_size=(3, 3), strides = 2, padding = "same", activation='relu'))

            network.add(layers.Flatten())
            network.add(layers.BatchNormalization())
            network.add(layers.Dense(2000, activation='relu'))
            network.add(layers.BatchNormalization())
            network.add(layers.Dense(2000, activation='relu'))
            network.add(layers.BatchNormalization())
            network.add(layers.Dense(units=10, activation = 'softmax'))

            default_transformer_kwargs = {"network" : network,
                                          "euclidean_layer_idx" : -2,
                                          "num_classes" : 10,
                                          "optimizer" : keras.optimizers.Adam(3e-4)
                                         }

            default_voter_class = KNNClassificationVoter
            default_voter_kwargs = {"k" : int(np.log2(len(train_x1)))}

            default_decider_class = SimpleArgmaxAverage

            progressive_learner = ProgressiveLearner(default_transformer_class = default_transformer_class,
                                         default_transformer_kwargs = default_transformer_kwargs,
                                         default_voter_class = default_voter_class,
                                         default_voter_kwargs = default_voter_kwargs,
                                         default_decider_class = default_decider_class)

            progressive_learner.add_task(
                X = train_x1,
                y = train_y1,
                transformer_voter_decider_split = [0.67, 0.33, 0],
                decider_kwargs = {"classes" : np.unique(train_y1)}
            )

            progressive_learner.add_transformer(
                X = tmp_data,
                y = train_y2,
                transformer_data_proportion = 1,
                backward_task_ids = [0]
            )


            llf_task1=progressive_learner.predict(test_x, task_id=0)
            llf_single_task=progressive_learner.predict(test_x, task_id=0, transformer_ids=[0])

            errors[1] = errors[1]+(1 - np.mean(llf_task1 == test_y))
            errors[0] = errors[0]+(1 - np.mean(llf_single_task == test_y))

    errors = errors/reps
    print("Errors For Angle {}: {}".format(angle, errors))
    with open('rotation_results/angle_'+str(angle)+'_'+model+'.pickle', 'wb') as f:
        pickle.dump(errors, f, protocol = 2)

def image_aug(pic, angle, centroid_x=23, centroid_y=23, win=16, scale=1.45):
    im_sz = int(np.floor(pic.shape[0]*scale))
    pic_ = np.uint8(np.zeros((im_sz,im_sz,3),dtype=int))

    pic_[:,:,0] = ndimage.zoom(pic[:,:,0],scale)

    pic_[:,:,1] = ndimage.zoom(pic[:,:,1],scale)
    pic_[:,:,2] = ndimage.zoom(pic[:,:,2],scale)

    image_aug = rotate(pic_, angle, resize=False)
    #print(image_aug.shape)
    image_aug_ = image_aug[centroid_x-win:centroid_x+win,centroid_y-win:centroid_y+win,:]

    return img_as_ubyte(image_aug_)

### MAIN HYPERPARAMS ###
model = "dnn"
granularity = 2
reps = 4
########################


(X_train, y_train), (X_test, y_test) = keras.datasets.cifar100.load_data()
data_x = np.concatenate([X_train, X_test])
data_y = np.concatenate([y_train, y_test])
data_y = data_y[:, 0]

def perform_angle(angle):
    LF_experiment(data_x, data_y, angle, model, granularity, reps=reps, ntrees=16, acorn=1)

if model == "dnn":
    for angle_adder in range(30, 180, granularity * 4):
        angles = angle_adder + np.arange(0, granularity * 4, granularity)
        with Pool(4) as p:
            p.map(perform_angle, angles)
elif model == "uf":
    angles = np.arange(30,180,2)
    Parallel(n_jobs=-1)(delayed(LF_experiment)(data_x, data_y, angle, model, granularity, reps=20, ntrees=16, acorn=1) for angle in angles)

[((26, 0, 26, 35), 'sys.path.append', 'sys.path.append', ({(26, 16, 26, 34): '"""../../proglearn/"""'}, {}), "('../../proglearn/')", False, 'import sys\n'), ((174, 39, 174, 74), 'keras.datasets.cifar100.load_data', 'keras.datasets.cifar100.load_data', ({}, {}), '()', False, 'import keras\n'), ((175, 9, 175, 42), 'numpy.concatenate', 'np.concatenate', ({(175, 24, 175, 41): '[X_train, X_test]'}, {}), '([X_train, X_test])', True, 'import numpy as np\n'), ((176, 9, 176, 42), 'numpy.concatenate', 'np.concatenate', ({(176, 24, 176, 41): '[y_train, y_test]'}, {}), '([y_train, y_test])', True, 'import numpy as np\n'), ((66, 13, 66, 24), 'numpy.zeros', 'np.zeros', ({(66, 22, 66, 23): '2'}, {}), '(2)', True, 'import numpy as np\n'), ((156, 18, 156, 48), 'scipy.ndimage.zoom', 'ndimage.zoom', ({(156, 31, 156, 41): 'pic[:, :, (0)]', (156, 42, 156, 47): 'scale'}, {}), '(pic[:, :, (0)], scale)', False, 'from scipy import ndimage\n'), ((158, 18, 158, 48), 'scipy.ndimage.zoom', 'ndimage.zoom', ({(158, 31, 158, 41): 'pic[:, :, (1)]', (158, 42, 158, 47): 'scale'}, {}), '(pic[:, :, (1)], scale)', False, 'from scipy import ndimage\n'), ((159, 18, 159, 48), 'scipy.ndimage.zoom', 'ndimage.zoom', ({(159, 31, 159, 41): 'pic[:, :, (2)]', (159, 42, 159, 47): 'scale'}, {}), '(pic[:, :, (2)], scale)', False, 'from scipy import ndimage\n'), ((161, 16, 161, 49), 'skimage.transform.rotate', 'rotate', (), '', False, 'from skimage.transform import rotate\n'), ((165, 11, 165, 35), 'skimage.util.img_as_ubyte', 'img_as_ubyte', ({(165, 24, 165, 34): 'image_aug_'}, {}), '(image_aug_)', False, 'from skimage.util import img_as_ubyte\n'), ((40, 8, 40, 28), 'random.shuffle', 'random.shuffle', ({(40, 23, 40, 27): 'indx'}, {}), '(indx)', False, 'import random\n'), ((64, 8, 64, 29), 'numpy.random.seed', 'np.random.seed', ({(64, 23, 64, 28): 'acorn'}, {}), '(acorn)', True, 'import numpy as np\n'), ((75, 16, 75, 46), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((150, 8, 150, 44), 'pickle.dump', 'pickle.dump', (), '', False, 'import pickle\n'), ((153, 16, 153, 44), 'numpy.floor', 'np.floor', ({(153, 25, 153, 43): 'pic.shape[0] * scale'}, {}), '(pic.shape[0] * scale)', True, 'import numpy as np\n'), ((154, 20, 154, 55), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((188, 13, 188, 32), 'numpy.arange', 'np.arange', ({(188, 23, 188, 25): '30', (188, 26, 188, 29): '180', (188, 30, 188, 31): '2'}, {}), '(30, 180, 2)', True, 'import numpy as np\n'), ((35, 11, 35, 32), 'numpy.where', 'np.where', ({(35, 20, 35, 31): '(data_y == u)'}, {}), '(data_y == u)', True, 'import numpy as np\n'), ((35, 45, 35, 62), 'numpy.unique', 'np.unique', ({(35, 55, 35, 61): 'data_y'}, {}), '(data_y)', True, 'import numpy as np\n'), ((51, 23, 51, 75), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((52, 23, 52, 77), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((53, 23, 53, 73), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((54, 23, 54, 75), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((56, 21, 56, 73), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((57, 21, 57, 71), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((90, 22, 90, 40), 'keras.Sequential', 'keras.Sequential', ({}, {}), '()', False, 'import keras\n'), ((120, 34, 124, 87), 'progressive_learner.ProgressiveLearner', 'ProgressiveLearner', (), '', False, 'from progressive_learner import ProgressiveLearner\n'), ((184, 31, 184, 73), 'numpy.arange', 'np.arange', ({(184, 41, 184, 42): '(0)', (184, 44, 184, 59): '(granularity * 4)', (184, 61, 184, 72): 'granularity'}, {}), '(0, granularity * 4, granularity)', True, 'import numpy as np\n'), ((185, 13, 185, 20), 'multiprocessing.Pool', 'Pool', ({(185, 18, 185, 19): '(4)'}, {}), '(4)', False, 'from multiprocessing import Pool\n'), ((189, 4, 189, 23), 'joblib.Parallel', 'Parallel', (), '', False, 'from joblib import Parallel, delayed\n'), ((92, 24, 92, 51), 'keras.layers.BatchNormalization', 'layers.BatchNormalization', ({}, {}), '()', False, 'from keras import layers\n'), ((93, 24, 93, 119), 'keras.layers.Conv2D', 'layers.Conv2D', (), '', False, 'from keras import layers\n'), ((94, 24, 94, 51), 'keras.layers.BatchNormalization', 'layers.BatchNormalization', ({}, {}), '()', False, 'from keras import layers\n'), ((95, 24, 95, 119), 'keras.layers.Conv2D', 'layers.Conv2D', (), '', False, 'from keras import layers\n'), ((96, 24, 96, 51), 'keras.layers.BatchNormalization', 'layers.BatchNormalization', ({}, {}), '()', False, 'from keras import layers\n'), ((97, 24, 97, 120), 'keras.layers.Conv2D', 'layers.Conv2D', (), '', False, 'from keras import layers\n'), ((98, 24, 98, 51), 'keras.layers.BatchNormalization', 'layers.BatchNormalization', ({}, {}), '()', False, 'from keras import layers\n'), ((99, 24, 99, 120), 'keras.layers.Conv2D', 'layers.Conv2D', (), '', False, 'from keras import layers\n'), ((101, 24, 101, 40), 'keras.layers.Flatten', 'layers.Flatten', ({}, {}), '()', False, 'from keras import layers\n'), ((102, 24, 102, 51), 'keras.layers.BatchNormalization', 'layers.BatchNormalization', ({}, {}), '()', False, 'from keras import layers\n'), ((103, 24, 103, 61), 'keras.layers.Dense', 'layers.Dense', (), '', False, 'from keras import layers\n'), ((104, 24, 104, 51), 'keras.layers.BatchNormalization', 'layers.BatchNormalization', ({}, {}), '()', False, 'from keras import layers\n'), ((105, 24, 105, 61), 'keras.layers.Dense', 'layers.Dense', (), '', False, 'from keras import layers\n'), ((106, 24, 106, 51), 'keras.layers.BatchNormalization', 'layers.BatchNormalization', ({}, {}), '()', False, 'from keras import layers\n'), ((107, 24, 107, 70), 'keras.layers.Dense', 'layers.Dense', (), '', False, 'from keras import layers\n'), ((112, 56, 112, 83), 'keras.optimizers.Adam', 'keras.optimizers.Adam', ({(112, 78, 112, 82): '(0.0003)'}, {}), '(0.0003)', False, 'import keras\n'), ((144, 39, 144, 67), 'numpy.mean', 'np.mean', ({(144, 47, 144, 66): '(llf_task1 == test_y)'}, {}), '(llf_task1 == test_y)', True, 'import numpy as np\n'), ((145, 39, 145, 73), 'numpy.mean', 'np.mean', ({(145, 47, 145, 72): '(llf_single_task == test_y)'}, {}), '(llf_single_task == test_y)', True, 'import numpy as np\n'), ((189, 24, 189, 46), 'joblib.delayed', 'delayed', ({(189, 32, 189, 45): 'LF_experiment'}, {}), '(LF_experiment)', False, 'from joblib import Parallel, delayed\n'), ((130, 46, 130, 65), 'numpy.unique', 'np.unique', ({(130, 56, 130, 64): 'train_y1'}, {}), '(train_y1)', True, 'import numpy as np\n'), ((91, 101, 91, 119), 'numpy.shape', 'np.shape', ({(91, 110, 91, 118): 'train_x1'}, {}), '(train_x1)', True, 'import numpy as np\n')]