import os
import sys
import pickle
import tarfile
import urllib.request
from pathlib import Path
import numpy as np
import tensorflow as tf


def untar(fname):
    if fname.endswith("tar.gz"):
        tar = tarfile.open(fname)
        tar.extractall()
        tar.close()
        print ("Extracted in the same directory")
    else:
        print ("Not a tar.gz file: '%s '" % sys.argv[0])


def unpickle(file):
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


def download_data(url, file_name):
    u = urllib.request.urlopen(url)
    f = open(file_name, 'wb')
    meta = u.info()
    file_size = int(meta["Content-Length"])
    print ("Downloading: %s, File size in bytes: %s" % (file_name, file_size))
  
    file_size_dl = 0
    block_sz = 8192
    while True:
        buffer = u.read(block_sz)
        if not buffer:
            break

        file_size_dl += len(buffer)
        f.write(buffer)
        status = r"%10d [%3.2f%%]" % (file_size_dl, file_size_dl * 100. / file_size)
        status = status + chr(8)*(len(status)+1)
        print (status),

    f.close()
    return True


def main():
    data_file_name = "cifar-10-python.tar.gz"
    url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"

    cwd = os.getcwd()
    data_file = Path(os.path.join(cwd, data_file_name))
    data_dir = Path(os.path.join(cwd, "cifar-10-batches-py"))

    # Check if 'cifar-10-python.tar.gz' is located in the same directory
    if not data_file.exists():
        download_data(url, data_file_name)
        print("Data is downloaded.")
    else:
        print("Data File Exists.")

    # Load the dataset
    if not data_dir.is_dir():
        print("Directory not detected, unpacking.")
        untar(data_file_name)
    else:
        print("Directory detected, not unpacking.")

    # Note: The dict has 2 entries 'data' and 'labels'. data is a 10000 x 3072
    #  numpy array of uint8.
    # each row is a 32 x 32 colour image. first 1024 is the r channel,
    #  next 1024 g channel and last 1024 b channel
    # is a 10000 number of 0 - 9

    data_dict = unpickle(os.path.join(str(data_dir), "data_batch_1"))
    filenames = data_dict[b'filenames']
    labels = data_dict[b'labels']
    data = data_dict[b'data']
    batch_label = data_dict[b'batch_label']

    print(len(data))
    print(len(labels))
    print(data[0])

    # create an image placeholder
    img = data[1].reshape([1, 3, 32, 32])
    img = np.moveaxis(img, 1, -1)
    print(img.shape)

    # shape of [batch, height, width, channels]
    img_tensor = tf.placeholder(tf.uint8, shape=[1, 32, 32, 3])
    img_summary = tf.summary.image("test_image", img_tensor)

    with tf.Session() as sess:

        file_writer = tf.summary.FileWriter('./logs', sess.graph)
        summary = sess.run(img_summary, feed_dict={img_tensor: img})
        file_writer.add_summary(summary)
        # train_writer.add_graph(sess.graph)


if __name__ == '__main__':
    # if len(sys.argv) < 2:
    #  print ("Usage: '%s filename'" % sys.argv[0])
    #  sys.exit(0)
    # untar(sys.argv[1])
    main()
