import os

from tensorflow.contrib.learn.python.learn.datasets import base
import numpy as np
import IPython

from subprocess import call

from keras.preprocessing import image

from influence.dataset import DataSet
from influence.inception_v3 import preprocess_input

BASE_DIR = 'data' # TODO: change

def fill(X, Y, idx, label, img_path, img_side):
    img = image.load_img(img_path, target_size=(img_side, img_side))
    x = image.img_to_array(img)
    X[idx, ...] = x
    Y[idx] = label

     
def extract_and_rename_animals():
    class_maps = [
        ('dog', 'n02084071'),
        ('cat', 'n02121808'),
        ('bird', 'n01503061'),
        ('fish', 'n02512053'),
        ('horse', 'n02374451'),
        ('monkey', 'n02484322'),
        ('zebra', 'n02391049'),
        ('panda', 'n02510455'),
        ('lemur', 'n02496913'),
        ('wombat', 'n01883070'),
        ]


    for class_string, class_id in class_maps:
        
        class_dir = os.path.join(BASE_DIR, class_string)
        print(class_dir)
        call('mkdir %s' % class_dir, shell=True)
        call('tar -xf %s.tar -C %s' % (os.path.join(BASE_DIR, class_id), class_dir), shell=True)
        
        for filename in os.listdir(class_dir):

            file_idx = filename.split('_')[1].split('.')[0]
            src_filename = os.path.join(class_dir, filename)
            dst_filename = os.path.join(class_dir, '%s_%s.JPEG' % (class_string, file_idx))
            os.rename(src_filename, dst_filename)



def load_animals(num_train_ex_per_class=300, 
                 num_test_ex_per_class=100,
                 num_valid_ex_per_class=0,
                 classes=None,
                 ):    

    num_channels = 3
    img_side = 299

    if num_valid_ex_per_class == 0:
        valid_str = ''
    else:
        valid_str = '_valid-%s' % num_valid_examples

    if classes is None:
        classes = ['dog', 'cat', 'bird', 'fish', 'horse', 'monkey', 'zebra', 'panda', 'lemur', 'wombat']
        data_filename = os.path.join(BASE_DIR, 'dataset_train-%s_test-%s%s.npz' % (num_train_ex_per_class, num_test_ex_per_class, valid_str))
    else:
        data_filename = os.path.join(BASE_DIR, 'dataset_%s_train-%s_test-%s%s.npz' % ('-'.join(classes), num_train_ex_per_class, num_test_ex_per_class, valid_str))

    num_classes = len(classes)
    num_train_examples = num_train_ex_per_class * num_classes
    num_test_examples = num_test_ex_per_class * num_classes
    num_valid_examples = num_valid_ex_per_class * num_classes

    if os.path.exists(data_filename):
        print('Loading animals from disk...')
        f = np.load(data_filename)
        X_train = f['X_train']
        X_test = f['X_test']
        Y_train = f['Y_train']
        Y_test = f['Y_test']

        if 'X_valid' in f:
            X_valid = f['X_valid']
        else:
            X_valid = None

        if 'Y_valid' in f:
            Y_valid = f['Y_valid']
        else:
            Y_valid = None

    else:
        print('Reading animals from raw images...')
        X_train = np.zeros([num_train_examples, img_side, img_side, num_channels])
        X_test = np.zeros([num_test_examples, img_side, img_side, num_channels])
        # X_valid = np.zeros([num_valid_examples, img_side, img_side, num_channels])
        X_valid = None

        Y_train = np.zeros([num_train_examples])
        Y_test = np.zeros([num_test_examples])
        # Y_valid = np.zeros([num_valid_examples])
        Y_valid = None

        for class_idx, class_string in enumerate(classes):
            print('class: %s' % class_string)            
            # For some reason, a lot of numbers are skipped.
            i = 0
            num_filled = 0
            while num_filled < num_train_ex_per_class:        
                img_path = os.path.join(BASE_DIR, '%s/%s_%s.JPEG' % (class_string, class_string, i))
                print(img_path)
                if os.path.exists(img_path):
                    fill(X_train, Y_train, num_filled + (num_train_ex_per_class * class_idx), class_idx, img_path, img_side)
                    num_filled += 1
                    print(num_filled)
                i += 1

            num_filled = 0
            while num_filled < num_test_ex_per_class:        
                img_path = os.path.join(BASE_DIR, '%s/%s_%s.JPEG' % (class_string, class_string, i))
                if os.path.exists(img_path):
                    fill(X_test, Y_test, num_filled + (num_test_ex_per_class * class_idx), class_idx, img_path, img_side)
                    num_filled += 1
                    print(num_filled)
                i += 1

            num_filled = 0
            while num_filled < num_valid_ex_per_class:        
                img_path = os.path.join(BASE_DIR, '%s/%s_%s.JPEG' % (class_string, class_string, i))
                if os.path.exists(img_path):
                    fill(X_valid, Y_valid, num_filled + (num_valid_ex_per_class * class_idx), class_idx, img_path, img_side)
                    num_filled += 1
                    print(num_filled)
                i += 1

        X_train = preprocess_input(X_train)
        X_test = preprocess_input(X_test)
        X_valid = preprocess_input(X_valid)

        np.random.seed(0)
        permutation_idx = np.arange(num_train_examples)
        np.random.shuffle(permutation_idx)
        X_train = X_train[permutation_idx, :]
        Y_train = Y_train[permutation_idx]
        permutation_idx = np.arange(num_test_examples)
        np.random.shuffle(permutation_idx)
        X_test = X_test[permutation_idx, :]
        Y_test = Y_test[permutation_idx]
        permutation_idx = np.arange(num_valid_examples)
        np.random.shuffle(permutation_idx)
        X_valid = X_valid[permutation_idx, :]
        Y_valid = Y_valid[permutation_idx]

        np.savez_compressed(data_filename, X_train=X_train, Y_train=Y_train, X_test=X_test, Y_test=Y_test, X_valid=X_valid, Y_valid=Y_valid)


    train = DataSet(X_train, Y_train)
    if (X_valid is not None) and (Y_valid is not None):
        # validation = DataSet(X_valid, Y_valid)
        validation = None
    else:
        validation = None

    test = DataSet(X_test, Y_test)

    return base.Datasets(train=train, validation=validation, test=test)


def load_koda():
    num_channels = 3
    img_side = 299

    data_filename = os.path.join(BASE_DIR, 'dataset_koda.npz')

    if os.path.exists(data_filename):
        print('Loading Koda from disk...')
        f = np.load(data_filename)
        X = f['X']
        Y = f['Y']
    else:
        # Returns all class 0
        print('Reading Koda from raw images...')

        image_files = [image_file for image_file in os.listdir(os.path.join(BASE_DIR, 'koda')) if (image_file.endswith('.jpg'))]
        # Hack to get the image files in the right order
        # image_files = [image_file for image_file in os.listdir(os.path.join(BASE_DIR, 'koda')) if (image_file.endswith('.jpg') and not image_file.startswith('124'))]
        # image_files += [image_file for image_file in os.listdir(os.path.join(BASE_DIR, 'koda')) if (image_file.endswith('.jpg') and image_file.startswith('124'))]


        num_examples = len(image_files)
        X = np.zeros([num_examples, img_side, img_side, num_channels])
        Y = np.zeros([num_examples])

        class_idx = 0
        for counter, image_file in enumerate(image_files):
            img_path = os.path.join(BASE_DIR, 'koda', image_file)
            fill(X, Y, counter, class_idx, img_path, img_side)

        X = preprocess_input(X)

        np.savez(data_filename, X=X, Y=Y)

    return X, Y
    

def load_dogfish_with_koda():        
    classes = ['dog', 'fish']
    X_test, Y_test = load_koda()

    data_sets = load_animals(num_train_ex_per_class=900, 
                 num_test_ex_per_class=300,
                 num_valid_ex_per_class=0,
                 classes=classes)
    train = data_sets.train
    validation = data_sets.validation
    test = DataSet(X_test, Y_test)

    return base.Datasets(train=train, validation=validation, test=test)


def load_dogfish_with_orig_and_koda():
    classes = ['dog', 'fish']
    X_test, Y_test = load_koda()
    X_test = np.reshape(X_test, (X_test.shape[0], -1))

    data_sets = load_animals(num_train_ex_per_class=900, 
                 num_test_ex_per_class=300,
                 num_valid_ex_per_class=0,
                 classes=classes)
    train = data_sets.train
    validation = data_sets.validation

    test = DataSet(
        np.concatenate((data_sets.test.x, X_test), axis=0), 
        np.concatenate((data_sets.test.labels, Y_test), axis=0))

    return base.Datasets(train=train, validation=validation, test=test)


[((361, 419), 'keras.preprocessing.image.load_img', 'image.load_img', (['img_path'], {'target_size': '(img_side, img_side)'}), '(img_path, target_size=(img_side, img_side))\n', (375, 419), False, 'from keras.preprocessing import image\n'), ((428, 451), 'keras.preprocessing.image.img_to_array', 'image.img_to_array', (['img'], {}), '(img)\n', (446, 451), False, 'from keras.preprocessing import image\n'), ((2507, 2536), 'os.path.exists', 'os.path.exists', (['data_filename'], {}), '(data_filename)\n', (2521, 2536), False, 'import os\n'), ((5927, 5952), 'influence.dataset.DataSet', 'DataSet', (['X_train', 'Y_train'], {}), '(X_train, Y_train)\n', (5934, 5952), False, 'from influence.dataset import DataSet\n'), ((6132, 6155), 'influence.dataset.DataSet', 'DataSet', (['X_test', 'Y_test'], {}), '(X_test, Y_test)\n', (6139, 6155), False, 'from influence.dataset import DataSet\n'), ((6168, 6228), 'tensorflow.contrib.learn.python.learn.datasets.base.Datasets', 'base.Datasets', ([], {'train': 'train', 'validation': 'validation', 'test': 'test'}), '(train=train, validation=validation, test=test)\n', (6181, 6228), False, 'from tensorflow.contrib.learn.python.learn.datasets import base\n'), ((6309, 6351), 'os.path.join', 'os.path.join', (['BASE_DIR', '"""dataset_koda.npz"""'], {}), "(BASE_DIR, 'dataset_koda.npz')\n", (6321, 6351), False, 'import os\n'), ((6360, 6389), 'os.path.exists', 'os.path.exists', (['data_filename'], {}), '(data_filename)\n', (6374, 6389), False, 'import os\n'), ((7934, 7957), 'influence.dataset.DataSet', 'DataSet', (['X_test', 'Y_test'], {}), '(X_test, Y_test)\n', (7941, 7957), False, 'from influence.dataset import DataSet\n'), ((7970, 8030), 'tensorflow.contrib.learn.python.learn.datasets.base.Datasets', 'base.Datasets', ([], {'train': 'train', 'validation': 'validation', 'test': 'test'}), '(train=train, validation=validation, test=test)\n', (7983, 8030), False, 'from tensorflow.contrib.learn.python.learn.datasets import base\n'), ((8148, 8189), 'numpy.reshape', 'np.reshape', (['X_test', '(X_test.shape[0], -1)'], {}), '(X_test, (X_test.shape[0], -1))\n', (8158, 8189), True, 'import numpy as np\n'), ((8595, 8655), 'tensorflow.contrib.learn.python.learn.datasets.base.Datasets', 'base.Datasets', ([], {'train': 'train', 'validation': 'validation', 'test': 'test'}), '(train=train, validation=validation, test=test)\n', (8608, 8655), False, 'from tensorflow.contrib.learn.python.learn.datasets import base\n'), ((954, 990), 'os.path.join', 'os.path.join', (['BASE_DIR', 'class_string'], {}), '(BASE_DIR, class_string)\n', (966, 990), False, 'import os\n'), ((1024, 1064), 'subprocess.call', 'call', (["('mkdir %s' % class_dir)"], {'shell': '(True)'}), "('mkdir %s' % class_dir, shell=True)\n", (1028, 1064), False, 'from subprocess import call\n'), ((1195, 1216), 'os.listdir', 'os.listdir', (['class_dir'], {}), '(class_dir)\n', (1205, 1216), False, 'import os\n'), ((1991, 2113), 'os.path.join', 'os.path.join', (['BASE_DIR', "('dataset_train-%s_test-%s%s.npz' % (num_train_ex_per_class,\n    num_test_ex_per_class, valid_str))"], {}), "(BASE_DIR, 'dataset_train-%s_test-%s%s.npz' % (\n    num_train_ex_per_class, num_test_ex_per_class, valid_str))\n", (2003, 2113), False, 'import os\n'), ((2596, 2618), 'numpy.load', 'np.load', (['data_filename'], {}), '(data_filename)\n', (2603, 2618), True, 'import numpy as np\n'), ((3028, 3092), 'numpy.zeros', 'np.zeros', (['[num_train_examples, img_side, img_side, num_channels]'], {}), '([num_train_examples, img_side, img_side, num_channels])\n', (3036, 3092), True, 'import numpy as np\n'), ((3110, 3173), 'numpy.zeros', 'np.zeros', (['[num_test_examples, img_side, img_side, num_channels]'], {}), '([num_test_examples, img_side, img_side, num_channels])\n', (3118, 3173), True, 'import numpy as np\n'), ((3301, 3331), 'numpy.zeros', 'np.zeros', (['[num_train_examples]'], {}), '([num_train_examples])\n', (3309, 3331), True, 'import numpy as np\n'), ((3349, 3378), 'numpy.zeros', 'np.zeros', (['[num_test_examples]'], {}), '([num_test_examples])\n', (3357, 3378), True, 'import numpy as np\n'), ((5073, 5098), 'influence.inception_v3.preprocess_input', 'preprocess_input', (['X_train'], {}), '(X_train)\n', (5089, 5098), False, 'from influence.inception_v3 import preprocess_input\n'), ((5116, 5140), 'influence.inception_v3.preprocess_input', 'preprocess_input', (['X_test'], {}), '(X_test)\n', (5132, 5140), False, 'from influence.inception_v3 import preprocess_input\n'), ((5159, 5184), 'influence.inception_v3.preprocess_input', 'preprocess_input', (['X_valid'], {}), '(X_valid)\n', (5175, 5184), False, 'from influence.inception_v3 import preprocess_input\n'), ((5194, 5211), 'numpy.random.seed', 'np.random.seed', (['(0)'], {}), '(0)\n', (5208, 5211), True, 'import numpy as np\n'), ((5238, 5267), 'numpy.arange', 'np.arange', (['num_train_examples'], {}), '(num_train_examples)\n', (5247, 5267), True, 'import numpy as np\n'), ((5276, 5310), 'numpy.random.shuffle', 'np.random.shuffle', (['permutation_idx'], {}), '(permutation_idx)\n', (5293, 5310), True, 'import numpy as np\n'), ((5426, 5454), 'numpy.arange', 'np.arange', (['num_test_examples'], {}), '(num_test_examples)\n', (5435, 5454), True, 'import numpy as np\n'), ((5463, 5497), 'numpy.random.shuffle', 'np.random.shuffle', (['permutation_idx'], {}), '(permutation_idx)\n', (5480, 5497), True, 'import numpy as np\n'), ((5609, 5638), 'numpy.arange', 'np.arange', (['num_valid_examples'], {}), '(num_valid_examples)\n', (5618, 5638), True, 'import numpy as np\n'), ((5647, 5681), 'numpy.random.shuffle', 'np.random.shuffle', (['permutation_idx'], {}), '(permutation_idx)\n', (5664, 5681), True, 'import numpy as np\n'), ((5780, 5917), 'numpy.savez_compressed', 'np.savez_compressed', (['data_filename'], {'X_train': 'X_train', 'Y_train': 'Y_train', 'X_test': 'X_test', 'Y_test': 'Y_test', 'X_valid': 'X_valid', 'Y_valid': 'Y_valid'}), '(data_filename, X_train=X_train, Y_train=Y_train, X_test\n    =X_test, Y_test=Y_test, X_valid=X_valid, Y_valid=Y_valid)\n', (5799, 5917), True, 'import numpy as np\n'), ((6446, 6468), 'numpy.load', 'np.load', (['data_filename'], {}), '(data_filename)\n', (6453, 6468), True, 'import numpy as np\n'), ((7170, 7228), 'numpy.zeros', 'np.zeros', (['[num_examples, img_side, img_side, num_channels]'], {}), '([num_examples, img_side, img_side, num_channels])\n', (7178, 7228), True, 'import numpy as np\n'), ((7241, 7265), 'numpy.zeros', 'np.zeros', (['[num_examples]'], {}), '([num_examples])\n', (7249, 7265), True, 'import numpy as np\n'), ((7490, 7509), 'influence.inception_v3.preprocess_input', 'preprocess_input', (['X'], {}), '(X)\n', (7506, 7509), False, 'from influence.inception_v3 import preprocess_input\n'), ((7519, 7552), 'numpy.savez', 'np.savez', (['data_filename'], {'X': 'X', 'Y': 'Y'}), '(data_filename, X=X, Y=Y)\n', (7527, 7552), True, 'import numpy as np\n'), ((8465, 8515), 'numpy.concatenate', 'np.concatenate', (['(data_sets.test.x, X_test)'], {'axis': '(0)'}), '((data_sets.test.x, X_test), axis=0)\n', (8479, 8515), True, 'import numpy as np\n'), ((8526, 8581), 'numpy.concatenate', 'np.concatenate', (['(data_sets.test.labels, Y_test)'], {'axis': '(0)'}), '((data_sets.test.labels, Y_test), axis=0)\n', (8540, 8581), True, 'import numpy as np\n'), ((1306, 1339), 'os.path.join', 'os.path.join', (['class_dir', 'filename'], {}), '(class_dir, filename)\n', (1318, 1339), False, 'import os\n'), ((1367, 1431), 'os.path.join', 'os.path.join', (['class_dir', "('%s_%s.JPEG' % (class_string, file_idx))"], {}), "(class_dir, '%s_%s.JPEG' % (class_string, file_idx))\n", (1379, 1431), False, 'import os\n'), ((1444, 1481), 'os.rename', 'os.rename', (['src_filename', 'dst_filename'], {}), '(src_filename, dst_filename)\n', (1453, 1481), False, 'import os\n'), ((7371, 7413), 'os.path.join', 'os.path.join', (['BASE_DIR', '"""koda"""', 'image_file'], {}), "(BASE_DIR, 'koda', image_file)\n", (7383, 7413), False, 'import os\n'), ((3767, 3840), 'os.path.join', 'os.path.join', (['BASE_DIR', "('%s/%s_%s.JPEG' % (class_string, class_string, i))"], {}), "(BASE_DIR, '%s/%s_%s.JPEG' % (class_string, class_string, i))\n", (3779, 3840), False, 'import os\n'), ((3892, 3916), 'os.path.exists', 'os.path.exists', (['img_path'], {}), '(img_path)\n', (3906, 3916), False, 'import os\n'), ((4257, 4330), 'os.path.join', 'os.path.join', (['BASE_DIR', "('%s/%s_%s.JPEG' % (class_string, class_string, i))"], {}), "(BASE_DIR, '%s/%s_%s.JPEG' % (class_string, class_string, i))\n", (4269, 4330), False, 'import os\n'), ((4350, 4374), 'os.path.exists', 'os.path.exists', (['img_path'], {}), '(img_path)\n', (4364, 4374), False, 'import os\n'), ((4713, 4786), 'os.path.join', 'os.path.join', (['BASE_DIR', "('%s/%s_%s.JPEG' % (class_string, class_string, i))"], {}), "(BASE_DIR, '%s/%s_%s.JPEG' % (class_string, class_string, i))\n", (4725, 4786), False, 'import os\n'), ((4806, 4830), 'os.path.exists', 'os.path.exists', (['img_path'], {}), '(img_path)\n', (4820, 4830), False, 'import os\n'), ((1104, 1136), 'os.path.join', 'os.path.join', (['BASE_DIR', 'class_id'], {}), '(BASE_DIR, class_id)\n', (1116, 1136), False, 'import os\n'), ((6660, 6690), 'os.path.join', 'os.path.join', (['BASE_DIR', '"""koda"""'], {}), "(BASE_DIR, 'koda')\n", (6672, 6690), False, 'import os\n')]