DATADIR = './dogImages'
OUTDIR = './model-k'
import os
import re
from glob import glob

def make_sample_dict(sample_dir):
    """Return the dictionaries between dog_breeds and number"""
    breed_list = []
    samples = [f for f in os.listdir(sample_dir) if os.path.isdir(os.path.join(sample_dir, f))]
    samples.sort()
    for sample in samples:
        breed_list.append(sample.split(".")[1])
    id2name = {i: name for i, name in enumerate(breed_list)}
    name2id = {name: i for i, name in id2name.items()}
    print('Number of dictionary is {}'.format(len(breed_list)))
    print('Example: {}'.format(breed_list[0]))
    return id2name, name2id

id2name, name2id = make_sample_dict(os.path.join(DATADIR, 'test'))

def load_data(data_dir):
    pattern = re.compile("(.+)\/(.+)\/([0-9][0-9][0-9]).(.+)\/(.+)jpg")
    print(pattern)
    all_files = glob(os.path.join(data_dir, "*/*/*jpg"))
    
    train, val, test = [], [], []
    for entry in all_files:
        r = re.match(pattern, entry)
        if r:
            regime, kind= r.group(2),r.group(4)
            kind_number = name2id[kind]
            sample = (entry, kind_number)
            
            if regime == 'train':
                train.append(sample)
            elif regime == 'valid':
                val.append(sample)
            elif regime == 'test':
                test.append(sample)
            else:
                raise ValueError('The regime does not much. Regime: {}'.format(regime))
                
    print('There are {} train, {} valid and {} test samples.'.format(len(train), len(val), len(test)))
    return train, val, test
                
trainset, valset, testset = load_data(DATADIR)                          


print(trainset[1000])
print(trainset[6000])

num_classes = len(id2name)

def extract_data(dataset):
    """The feature possibly filepath now. In that case, it's necessary to convert it to data."""
    feature_data, label_data = [], []
    for feature, label in dataset:
        feature_data.append(feature)
        label_data.append(label)
    return feature_data, label_data
"""
train_data_path, train_label = extract_data(trainset)
valid_data_path, valid_label = extract_data(valset)
test_data_path, test_label = extract_data(testset)
"""

"""
Explanation: One of a Basic TensorFlow Tutorial for CNN

In this notebook, I'll show how to use TensorFlow for CNN.
https://www.tensorflow.org/tutorials/layers

1. Data Preparation

The './dogImages' direcory has three directories; 'train', 'test' and 'valid'. 
All of these directories have directiries containing images. These directories' names like '001.Affenpinsher'.
So, it's necessary to make sets of labels and filepaths. This is easily done by using Keras. But I'll try other approach.
End of explanation
"""


def file2resized_img(filepath):
    """Return image converted and resized to ndarray from filepath.
    Args:
        :filepath: filepath to image file.
    Returns:
        resized ndarray from the filepath.
    """
    # img: PIL Image instance.
    img = tf.keras.preprocessing.image.load_img(filepath, target_size=(DATA_HEIGHT, DATA_WIDTH))
    #plt.imshow(img)
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    return img_array

def normarizer(feature_data):
    """
    Args:
        :feature_data: 3-D ndarray.
    Returns:
        3-D ndarray divided 255.
    """
    img_max = 255
    resized_data = []
    for data in feature_data:
        resized_data.append(data/img_max)
    resized_data = np.array(resized_data)
    return resized_data

def dataset_to_feature_and_label(dataset):
    feature_data, label_data = [], []
    for filepath, label in dataset:
        try:
            feature_data.append(file2resized_img(filepath))
            label_data.append(label)
        except:
            continue
    return feature_data, label_data

train_data, train_label = dataset_to_feature_and_label(trainset)
valid_data, valid_label = dataset_to_feature_and_label(valset)
test_data, test_label = dataset_to_feature_and_label(testset)

train_data = normarizer(train_data)
valid_data = normarizer(valid_data)

test_data = normarizer(test_data)

train_label = np.asarray(train_label, dtype=np.int32)

valid_label = np.asarray(valid_label, dtype=np.int32)
test_label = np.asarray(test_label, dtype=np.int32)


type(train_label[0])

train_data[0]

"""
Explanation: http://www.mathgram.xyz/entry/keras/preprocess/img
https://github.com/tensorflow/tensorflow/blob/r1.4/tensorflow/python/keras/_impl/keras/preprocessing/image.py
https://www.tensorflow.org/versions/r0.12/api_docs/python/image/cropping
https://www.tensorflow.org/api_docs/python/tf/keras/preprocessing/image/img_to_array
End of explanation
"""


check_str = ['hello', 'hello/world', 'dogImages/test/001.Affenpinsher/Alaska_001.jpg']
pattern = re.compile("(.+)\/(.+)\/([0-9][0-9][0-9]).(.+)\/(.+)jpg")
print(pattern)
for entry in check_str:
    r = re.match(pattern, entry)
    if r:
        print(entry)
        print(r.group(0), r.group(1), r.group(4))

"""
Explanation: COLUMN: An example how to use RE

https://docs.python.jp/3/library/re.html#re.regex.match
End of explanation
"""


import numpy as np
def one_hot(dataset, num_classes):
    """Change the label of the data into one-hot
    Args:
        dataset: A list of tuples. Each label is like (path_to_the_file, int_label)
    Returns:
        A list of tuples. Each label is like (path_to_the_file, int_label, ndarray-one-hot_label)"""
    ret_list = []
    for item in dataset:
        path_to_file, int_label = item
        one_hot_label = np.zeros(num_classes)
        one_hot_label[int_label] = 1
        ret_list.append((path_to_file, int_label, one_hot_label))
    return ret_list

"""
trainset = one_hot(trainset, num_classes)
valset = one_hot(valset, num_classes)
testset = one_hot(testset, num_classes)
"""

"""
Explanation: One-hot encodeing

I made this function, but it's may be not necessary.
In my trial, it's necessary to change labels into one-hot encoding.
End of explanation
"""


DATA_WIDTH = 64
DATA_HEIGHT = 64

import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
%matplotlib inline
tf.logging.set_verbosity(tf.logging.INFO)

def flatten(x_tensor):
    """
    Flatten x_tensor to (Batch Size, Flattened Image Size)
    Args:
        x_tensor: A Tensor of size (Batch Size, ...), where ... are the image dimensions.
    Returns:
        A tensor of size (Batch Size, Flattened Image Size)
    """
    tensor_size = x_tensor.get_shape().as_list()
    flatten_image_size = tensor_size[1] * tensor_size[2] * tensor_size[3]
    flatten_tensor = tf.reshape(x_tensor, [-1, flatten_image_size])
    return flatten_tensor

def cnn_model_fn(features, labels, mode):
    input_layer = tf.reshape(features["x"], [-1, DATA_HEIGHT, DATA_WIDTH, 3])
    
    # Conv Layer 1
    conv1 = tf.layers.conv2d(
        inputs=input_layer,
        filters=16,
        kernel_size=[3, 3],
        padding="same",
        activation=tf.nn.elu)
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2,2], strides=2)
    
    # Conv Layer 2
    conv2 = tf.layers.conv2d(
        inputs=pool1,
        filters=32,
        kernel_size=[3, 3],
        padding="same",
        activation=tf.nn.elu)
    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
    
    # Dense Layer
    pool2_flat = flatten(pool2)
    dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.elu)
    dropout = tf.layers.dropout(inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
    
    # Logits Layer
    logits = tf.layers.dense(inputs=dropout, units=num_classes)
    
    predictions = {
            "classes":  tf.argmax(input=logits, axis=1),
            "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
    }
    
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
    
    onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth = num_classes)
    loss = tf.losses.softmax_cross_entropy(
        onehot_labels=onehot_labels, logits=logits)
    
    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
        train_op = optimizer.minimize(
            loss=loss,
            global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
    
    eval_metric_ops = {
        "accuracy": tf.metrics.accuracy(
            labels=labels, predictions=predictions["classes"])
    }
    return tf.estimator.EstimatorSpec(
        mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

def main(train_data, train_labels, valid_data, valid_labels):
    
    image_classifier = tf.estimator.Estimator(
        model_fn=cnn_model_fn, model_dir=model_dir)
    
    tensors_to_log = {"probabilities": "softmax_tensor"}
    logging_hook = tf.train.LoggingTensorHook(
        tensors=tensors_to_log, every_n_iter=50)
    
    train_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x": train_data},
        y=train_labels,
        batch_size=100,
        num_epochs=None,
        shuffle=True)
    
    image_classifier.train(
        input_fn = train_input_fn,
        steps=20000,
        hooks=[logging_hook])
    
    eval_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x": valid_data},
        y=valid_labels,
        num_epochs=1,
        shuffle=False)
    eval_results = image_classifier.evaluate(input_fn=eval_input_fn)
    print(eval_results)

os.makedirs(os.path.join(OUTDIR, 'eval'), exist_ok=True)
model_dir=OUTDIR

tf.app.run(main=main(train_data, train_label, valid_data, valid_label))

"""
Explanation: 2. Make a Model
End of explanation
"""


from tqdm import tqdm
from sklearn.metrics import accuracy_score
def test(test_data, test_labels):
    
    image_classifier = tf.estimator.Estimator(
        model_fn=cnn_model_fn, model_dir=model_dir)
    
    test_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x": test_data},
        y=test_labels,
        batch_size=100,
        num_epochs=1,
        shuffle=False)
    it = image_classifier.predict(input_fn=test_input_fn)
    #saver = tf.train.Saver()
    #sess = tf.Session()
    #saver.restore(sess, model_dir)
    predict_list = []
    for t in tqdm(it):
        label =  [t["classes"]]
        predict_list.append(label)
    accuracy = accuracy_score(test_labels, predict_list)
    print('accuracy is {}'.format(accuracy))
    

test(test_data, test_label)

"""
Explanation: 4. Test
End of explanation
"""


