import os
import sys
import glob
import argparse
import shutil
import PIL
from PIL import Image
import random
from random import shuffle

def number_snapshots(input_dir, count):
    # import pdb; pdb.set_trace()
    count = count
    oldwd = os.getcwd()
    os.chdir(input_dir) # go to /images/
    print "Count is: {0}".format(count)
    print glob.glob("./*") # ./train, ./test, ./val
    for ttv in glob.glob("./*"):
        folders_list = glob.glob(ttv + "/*") # ./train/face ...
        for folder in folders_list:
            print folder
            fldername = folder.split("/")[-1] # face or no_face
            outfolder = "./" + ttv.split("/")[-1] + "/new_" + fldername + "/"
            files_list = glob.glob(folder + "/*")
            for f in files_list:
                try:
                    if not os.path.exists(outfolder):
                        os.makedirs(outfolder)
                    Image.open(f).save(outfolder + str(count) + ".jpg", "JPEG")
                    count += 1
                except IOError:
                    print "Cannot convert: {0}".format(f)
                    sys.exit(1)
            # remove the old face/ and not_face/ dirs
            shutil.rmtree(folder)
            os.rename(outfolder, folder)

    os.chdir(oldwd)
    print "Done."

def resize(input_dir, resize_dim):
    oldwd = os.getcwd()
    os.chdir(input_dir)
    print glob.glob("./*") # ./train, ./teset, ./val
    for ttv in glob.glob("./*"):
        folders_list = glob.glob(ttv + "/*")
        for folder in folders_list:
            print folder
            fldername = folder.split("/")[-1]
            outfolder = "./" + ttv.split("/")[-1] + "/new_" + fldername + "/"
            files_list = glob.glob(folder + "/*")
            for f in files_list:
                outfile = f.split('/')[-1]
                try:
                    if not os.path.exists(outfolder):
                        os.makedirs(outfolder)
                    ### DO THIS to Preserve aspect ratio ###
                    # basewidth = int(args.resize_width) # ex. 300 px
                    # img = Image.open(f)
                    # wpercent = (basewidth/float(img.size[0]))
                    # hsize = int((float(img.size[1])*float(wpercent)))
                    # img = img.resize((basewidth,hsize), PIL.Image.ANTIALIAS)
                    # img.save(outfolder + str(count) + ".jpg", "JPEG")
                    #### DO This to Create Square Images ###
                    img = Image.open(f)
                    size = tuple(map(lambda k: int(k), resize_dim.split('x')))
                    img = img.resize(size, Image.ANTIALIAS)
                    img.save(outfolder + outfile, "JPEG")
                except IOError:
                    print "Cannot convert: {0}".format(f)
                    sys.exit(1)
            # remove the old face no_face dirs and rename the cropped dirs
            shutil.rmtree(folder)
            os.rename(outfolder, folder)
    os.chdir(oldwd)
    print "Done."

def convert2txt(input_dir, txt_dir):
    """ Creates train.txt, test.txt, val.txt files from the input images folder """
    images_root = input_dir # ./data/images/
    txt_dir = txt_dir # ./data/txt/

    for ttv in glob.glob(images_root + "*"):
        # ttv will be test, train, val
        print ttv+"\n"
        data_src = ttv.split('/')[-1] # train, test, val
        writefolder = txt_dir
        writefilename = data_src + '.txt'
        print "Writing to: {0}\n".format(writefilename)
        if not os.path.exists(writefolder):
            os.makedirs(writefolder)
        old_pwd = os.getcwd()
        print "Changing dir"
        os.chdir(old_pwd + "/" + writefolder)

        with open(writefilename, 'w') as wf:
            print "Opened file: " + writefilename
            os.chdir(old_pwd)
            print glob.glob(ttv + "/*")
            for label in glob.glob(ttv + "/*"):
                # label might be .data/images/train/face or ..../no_face
                print label
                for imgsrc in glob.glob(label + "/*"):
                    line = imgsrc.replace(input_dir, "") # should be in form 'train/no_face/472.jpg'
                    if "no_face" in label.split('/'):
                        # label them 0
                        line = line + " 0\n" # 'train/no_face/472.jpg 0'
                    else:
                        line = line + " 1\n" # 'train/face'
                    wf.write(line)

    print "Done."

def shuffle_all_data(images_root, ratio):
    """ images_root is going to be ./data/images/
        ratio determines train/val/test data ratios
        by default it is set to 48:32:20 percentages
    """
    ratios = [int(e) for e in ratio.split(":")] # ratios [train : test : val]
    image_paths = []
    ttv_s = glob.glob(images_root + "*")
    # collect all image paths in one big list
    for ttv in ttv_s:
        for folder in glob.glob(ttv + "/*"):
            image_paths.extend(glob.glob(folder + "/*"))
    # import pdb; pdb.set_trace()
    random.shuffle(image_paths, random.random) # in-place shuffling
    size = len(image_paths)
    train_size = size*ratios[0]/100
    test_size = size*ratios[1]/100
    val_size = size - train_size - test_size
    if not os.path.exists(images_root + "train_shuffled/"):
        os.makedirs(images_root + "train_shuffled/")
        os.makedirs(images_root + "train_shuffled/face/")
        os.makedirs(images_root + "train_shuffled/no_face/")
    if not os.path.exists(images_root + "test_shuffled/"):
        os.makedirs(images_root + "test_shuffled/")
        os.makedirs(images_root + "test_shuffled/face/")
        os.makedirs(images_root + "test_shuffled/no_face/")
    if not os.path.exists(images_root + "val_shuffled/"):
        os.makedirs(images_root + "val_shuffled/")
        os.makedirs(images_root + "val_shuffled/face/")
        os.makedirs(images_root + "val_shuffled/no_face/")
    train_paths = image_paths[0 : train_size]
    test_paths = image_paths[train_size : train_size+test_size]
    val_paths = image_paths[-val_size:]
    for srcpath in train_paths:
        destpath = srcpath.replace("/train/","/train_shuffled/").replace("/test/",
            "/train_shuffled/").replace("/val/","/train_shuffled/")
        print "Source: {0}  | Dest {1}\n".format(srcpath, destpath)
        os.rename(srcpath, destpath)

    print "Training Shuffled...\n"

    for srcpath in test_paths:
        destpath = srcpath.replace("/train/", "/test_shuffled/").replace("/test/",
            "/test_shuffled/").replace("/val/", "/test_shuffled/")
        print "Source: {0}  | Dest {1}\n".format(srcpath, destpath)
        os.rename(srcpath, destpath)

    print "Testing Shuffled...\n"

    for srcpath in val_paths:
        destpath = srcpath.replace("/train/", "/val_shuffled/").replace("/test/",
            "/val_shuffled/").replace("/val/", "/val_shuffled/")
        print "Source: {0}  | Dest {1}\n".format(srcpath, destpath)
        os.rename(srcpath, destpath)

    print "Val Shuffled...\n"
    
    shutil.rmtree(images_root + "train/")
    os.rename(images_root + "train_shuffled/", images_root + "train/")
    shutil.rmtree(images_root + "test/")
    os.rename(images_root + "test_shuffled/", images_root + "test/")
    shutil.rmtree(images_root + "val/")
    os.rename(images_root + "val_shuffled/", images_root + "val/")

    print "Done."