from IPython.core.display import display, HTML
display(HTML("<style>.container { width:100% !important; }</style>"))

# Python ≥3.5 is required
import sys
assert sys.version_info >= (3, 5)

# Scikit-Learn ≥0.20 is required
import sklearn
assert sklearn.__version__ >= "0.20"

# Common imports
import numpy as np
import os
import os.path as osp
from PIL import Image
import skimage
from skimage import io
import os
import re
from PIL import Image
import random
import shutil
import time

# to make this notebook's output stable across runs
np.random.seed(42)

# To plot pretty figures
#%matplotlib inline
import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.rc('axes', labelsize=14)
mpl.rc('xtick', labelsize=12)
mpl.rc('ytick', labelsize=12)
plt.rc('font', size=12) 
plt.rc('figure', figsize = (12, 5))

import pandas as pd
pd.set_option('display.max_rows', 25)
pd.set_option('display.max_columns', 500)
pd.set_option('display.max_colwidth', 50)

# Ignore useless warnings (see SciPy issue #5998)
import warnings
warnings.filterwarnings(action="ignore", message="^internal gelsd")

def humans(y, human):
    res = [y == human]
    res = res[0]
    return res.any()

def plot_images(dataset, images, human, indices:list=None, nrows:int=1, ncols:int=1, size=(15, 15), root='./', imagesdir='', annotationdir='', train=False):
    if indices is None:
        indices = [x for x in  range(len(images))]
    
    sub_path = 'training' if train else 'validation'
    fig = plt.figure(figsize=size)
    for ax in range(1, (nrows * ncols),2):
        plt.subplot(nrows, ncols, ax)
        
        index = random.choice(indices)
        y = io.imread(f'{root}/{annotationdir}/{sub_path}/{images[index]}.png')
        
        while(not humans(y, human)):
            index = random.choice(indices)
            y = io.imread(f'{root}/{annotationdir}/{sub_path}/{images[index]}.png')
        
        img = io.imread(f"{root}/{imagesdir}/{sub_path}/{images[index]}.jpg")
        
        plt.imshow(img)
        plt.subplot(nrows, ncols, ax + 1)
        
        plt.imshow(y)
    
    plt.savefig(f'./DatasetGraphs/{dataset}_{sub_path}.png')
    plt.show()
    

def calculate_pixels_with_value(y, label):
    res = [y == label]
    res = np.asarray(res[0])
    indices = res.nonzero()
    #      label pixels / total pixels              label pixsels    total pixels
    return len(indices[0])/(y.shape[0]*y.shape[1]), len(indices[0]), (y.shape[0]*y.shape[1])

def calc_data_info(images_path, annotation_path):
    images_names = os.listdir(images_path)
    y_names = os.listdir(annotation_path)
    #print(len(y_names))
    
    total_pixels = 0
    human_pixels = 0
    sum_percent_per_image = 0
    total_size = 0
    
    for index, value in enumerate(images_names):
        split = re.split(r'\.(?!\d)', value)
        full_name_y = f'{annotation_path}{split[0]}.png'
        label = io.imread(full_name_y)
        total_size += 1
        if label.any():
            perc, human, total = calculate_pixels_with_value(label, 1)
            sum_percent_per_image += perc 
            human_pixels += human
            total_pixels += total
        else:
            total_pixels += label.shape[0]*label.shape[1]
    
    result1 = sum_percent_per_image/total_size
    resutl2 = human_pixels/total_pixels
    print("total size ---------------------------\t", total_size)
    print("pixels -------------------------------\t", total_pixels)
    print("mean human pixels percentage per image\t", result1)
    print("human pixels percentage per all pixels\t", resutl2)


def get_curv(
    input_images_path,
    input_annotations_path,
    train=False,
    sub_path='validation'):
    
    values = list()
    content = os.listdir(f'{input_annotations_path}/{sub_path}')
    #print(len(content))
    for index, value in enumerate(content):
        label = io.imread(f'{input_annotations_path}/{sub_path}/{value}')
        if label.any():
            perc, human, total = calculate_pixels_with_value(label, 1)
            values.append(perc)
    
    return values
    

def show_curv(data_set,input_images_path,input_annotations_path,train=False, sub_path='validation'):
    results = get_curv(input_images_path,input_annotations_path,train, sub_path)
    res = pd.Series(results)
    print("human size ---------------------------\t", len(results))
    res.hist(bins=len(res))
    plt.xlabel("Percentage of human pixels")
    plt.ylabel("nº samples")
    plt.savefig(f'./{data_set}_{sub_path}.png')
    

def mkdir_if_not_exist(dir_name, mode=0o777):
    if dir_name == '':
        return
    dir_name = osp.expanduser(dir_name)
    os.makedirs(dir_name, mode=mode, exist_ok=True)

def clean_data(
    input_images_path,
    input_annotations_path,
    output_images_path=None,
    output_annotations_path=None,
    train=False, 
    threshold_down=0.05,
    threshold_up=0.7):
    
    sub_path = 'training' if train else 'validation'
    dummy = 'removed'
    nohuman = 'nohuman'
    
    content = os.listdir(f'{input_annotations_path}/{sub_path}')
    print(len(content))
    
    if not output_images_path or not output_annotations_path:
        return
    
    mkdir_if_not_exist(output_images_path)
    mkdir_if_not_exist(f'{output_images_path}/{sub_path}')
    mkdir_if_not_exist(f'{output_images_path}/{dummy}')
    mkdir_if_not_exist(f'{output_images_path}/{nohuman}')
    
    mkdir_if_not_exist(output_annotations_path)
    mkdir_if_not_exist(f'{output_annotations_path}/{sub_path}')
    mkdir_if_not_exist(f'{output_annotations_path}/{dummy}')
    mkdir_if_not_exist(f'{output_annotations_path}/{nohuman}')
    
    human_counter = 0
    human_saved = 0
    human_not_saved = 0
    not_human_saved = 0
    
    for index, value in enumerate(content):
        split = re.split(r'\.(?!\d)', value)
        full_name_img = f'{input_images_path}/{sub_path}/{split[0]}.jpg'
        full_name_y = f'{input_annotations_path}/{sub_path}/{split[0]}.png'
        img = io.imread(full_name_img)
        label = io.imread(full_name_y)
        
        aux = dummy
        
        if label.any():
            human_counter += 1
            perc, human, total = calculate_pixels_with_value(label, 1)
            if perc < threshold_down or perc > threshold_up:
                aux = dummy
                human_not_saved += 1
            else:
                aux = sub_path
                human_saved += 1  
        else:
            aux = nohuman
            not_human_saved += 1
        
        io.imsave(
            '{0}/{1}/{2}.png'.format(
                output_annotations_path,
                aux, 
                split[0]
            ),
            label
        )

        io.imsave(
            '{0}/{1}/{2}.jpg'.format(
                output_images_path,
                aux,
                split[0]
            ),
            img
        )
        
    print(human_counter)
    print(human_saved)
    print(human_not_saved)
    print(not_human_saved)
    

def remove_not_needed_files(
    abs_path,
    images_path,
    annotations_path,
    sub_path=['validation', 'training'],
    file_name='files_to_remove.txt',
    folder_name='removed'):
    
    counter = 0
    counter_removed = 0
    
    for sub in sub_path:
        content = os.listdir(f'./{annotations_path}/{sub}')
        images = os.listdir(f'./{images_path}/{sub}')
        print(sub)
        assert (len(content) == len(images))
        print(len(images))
    
    mkdir_if_not_exist(f'./{images_path}/{folder_name}')
    mkdir_if_not_exist(f'./{annotations_path}/{folder_name}')
    
    # remove images detected by human detector
    files_to_remove = pd.read_csv(file_name, header=None, delimiter="\s+")
    files_to_remove = files_to_remove.to_numpy()
    
    
    index = 0
    while index < len(files_to_remove):   
        for sub in sub_path:
            value = files_to_remove[index][0]
            split = re.split(r'\.(?!\d)', value)

            src_img = '/{0}/{1}/{2}/{3}.jpg'.format(abs_path, images_path, sub, split[0])
            src_y = '/{0}/{1}/{2}/{3}.png'.format(abs_path, annotations_path, sub, split[0])

            dist_img_folder = '/{0}/{1}/{2}/'.format(abs_path, images_path, folder_name)
            dist_y_folder = '/{0}/{1}/{2}/'.format(abs_path, annotations_path, folder_name)

            try:
                shutil.move(src_y, dist_y_folder)
                shutil.move(src_img, dist_img_folder)
                time.sleep(0.01)
                counter_removed += 1
                sys.stdout.flush()
                sys.stdout.write('\r'+f'res: {counter_removed}/{len(files_to_remove)}')
            except Exception as e:
                time.sleep(0.01)
                sys.stdout.flush()
                sys.stdout.write('\r'+f'res: {counter_removed}/{len(files_to_remove)}')
        
        index += 1

    sys.stdout.flush()
    sys.stdout.write('\r'+f'res: {counter_removed}/{len(files_to_remove)}')

    sys.stdout.flush()
    sys.stdout.write('\r'+f'res: {counter_removed}/{len(files_to_remove)}')


def merge_and_split_images(
    abs_path,
    images_path,
    annotations_path,
    division=0.2):
    
    sub_path_training = 'training'
    sub_path_validation = 'validation'
    sub_path_testing = 'test'
    
    c_val = 0
    c_train = 0
    c_test = 0
    
    content1 = os.listdir(f'{annotations_path}/{sub_path_training}')
    images1 = os.listdir(f'{images_path}/{sub_path_training}')
    #print(len(content1))
    #print(len(images1))
    
    content2 = os.listdir(f'{annotations_path}/{sub_path_validation}')
    images2 = os.listdir(f'{images_path}/{sub_path_validation}')
    #print(len(content2))
    #print(len(images2))
    
    total = len(content1) + len(content2)
    print(total)

    split_validation = int(division*total)+1
    split_test = int(division*0.5*total)+1
    split_training = total-split_validation
    split_training -= split_test
    
    print('trainig:', split_training, split_training/total)
    print('validation:', split_validation, split_validation/total)
    print('test', split_test, split_test/total)

    sub_path_training_x = f'{sub_path_training}_{1-(division+division*0.5)}'
    sub_path_validation_x = f'{sub_path_validation}_{division}'
    sub_path_testing_x = f'{sub_path_testing}_{division*0.5}'
    
    mkdir_if_not_exist(f'{images_path}/{sub_path_training_x}')
    mkdir_if_not_exist(f'{images_path}/{sub_path_validation_x}')
    mkdir_if_not_exist(f'{images_path}/{sub_path_testing_x}')
    
    mkdir_if_not_exist(f'{annotations_path}/{sub_path_training_x}')
    mkdir_if_not_exist(f'{annotations_path}/{sub_path_validation_x}')
    mkdir_if_not_exist(f'{annotations_path}/{sub_path_testing_x}')
     
    for index, value in enumerate(content2):
        split = re.split(r'\.(?!\d)', value)

        src_img = '/{0}/{1}/{2}/{3}.jpg'.format(abs_path, images_path, sub_path_validation, split[0])
        src_y = '/{0}/{1}/{2}/{3}.png'.format(abs_path, annotations_path, sub_path_validation, split[0])

        if ((random.randrange(0,1) == 0 and c_val <= split_validation) or (c_train >= split_training and c_test >= split_test)): 
            # validation
            dist_img_folder = '/{0}/{1}/{2}/'.format(abs_path, images_path, sub_path_validation_x)
            dist_y_folder = '/{0}/{1}/{2}/'.format(abs_path, annotations_path, sub_path_validation_x)
            c_val = c_val + 1
        else: 
            if ((random.randrange(0,2) == 0 and c_test <= split_test) or c_train >= split_training): 
                # validation
                dist_img_folder = '/{0}/{1}/{2}/'.format(abs_path, images_path, sub_path_testing_x)
                dist_y_folder = '/{0}/{1}/{2}/'.format(abs_path, annotations_path, sub_path_testing_x)
                c_test = c_test + 1
            else:
                # training
                dist_img_folder = '/{0}/{1}/{2}/'.format(abs_path, images_path, sub_path_training_x)
                dist_y_folder = '/{0}/{1}/{2}/'.format(abs_path, annotations_path, sub_path_training_x)
                c_train = c_train + 1

        shutil.copy(src_y, dist_y_folder)
        shutil.copy(src_img, dist_img_folder)
        time.sleep(0.01)
        sys.stdout.flush()
        sys.stdout.write('\r'+f'train: {c_train}/{c_train/split_training} - val: {c_val}/{c_val/split_validation} - test: {c_test}/{c_test/split_test}')
    
        
    for index, value in enumerate(content1):
        split = re.split(r'\.(?!\d)', value)

        src_img = '/{0}/{1}/{2}/{3}.jpg'.format(abs_path, images_path, sub_path_training, split[0])
        src_y = '/{0}/{1}/{2}/{3}.png'.format(abs_path, annotations_path, sub_path_training, split[0])

        if ((random.randrange(0,2) == 0 and c_val <= split_validation) or (c_train >= split_training and c_test >= split_test)): 
            # validation
            dist_img_folder = '/{0}/{1}/{2}/'.format(abs_path, images_path, sub_path_validation_x)
            dist_y_folder = '/{0}/{1}/{2}/'.format(abs_path, annotations_path, sub_path_validation_x)
            c_val = c_val + 1
        else: 
            if ((random.randrange(0,2) == 0 and c_test <= split_test) or c_train >= split_training): 
                # validation
                dist_img_folder = '/{0}/{1}/{2}/'.format(abs_path, images_path, sub_path_testing_x)
                dist_y_folder = '/{0}/{1}/{2}/'.format(abs_path, annotations_path, sub_path_testing_x)
                c_test = c_test + 1
            else:
                # training
                dist_img_folder = '/{0}/{1}/{2}/'.format(abs_path, images_path, sub_path_training_x)
                dist_y_folder = '/{0}/{1}/{2}/'.format(abs_path, annotations_path, sub_path_training_x)
                c_train = c_train + 1

        shutil.copy(src_y, dist_y_folder)
        shutil.copy(src_img, dist_img_folder)
        time.sleep(0.01)
        sys.stdout.flush()
        sys.stdout.write('\r'+f'train: {c_train}/{c_train/split_training} - val: {c_val}/{c_val/split_validation} - test: {c_test}/{c_test/split_test}')
    
    sys.stdout.flush()
    sys.stdout.write('\r'+f'train: {c_train}/{c_train/split_training} - val: {c_val}/{c_val/split_validation} - test: {c_test}/{c_test/split_test}')
    
    
def change_annotation(y, human):
    res = [y == human]
    res = np.asarray(res[0])
    indices = res.nonzero()
    ziped = list(zip(list(indices[0]), list(indices[1])))
    y[:,:] = 0
    for i,j in ziped:
        y[i,j] = 1
    return y

def load_outside_data(outisdepath, sub_path, distpath_images, distpath_annot):
    counter = 0
    c = 0
    images_names = os.listdir(f'{outisdepath}')
    print(len(images_names))
    for index, value in enumerate(images_names):
        human = False
        split = re.split(r'\.(?!\d)', value)
        full_name_img = f'{outisdepath}{value}'
        img = io.imread(full_name_img)
        label = np.asarray(img[...,0]).copy()
        result = change_annotation(label, 13)
        counter += 1
        io.imsave('{0}/{1}/{2}.png'.format(distpath_annot, sub_path, split[0]), result)
        io.imsave('{0}/{1}/{2}.jpg'.format(distpath_images, sub_path, split[0]), img)

    print(counter)
def add_outside_data(
    abs_path,
    images_path,
    annotations_path,
    outisdepath):
    
    outside = 'outised'
    counter = 0
    mkdir_if_not_exist(f'{images_path}/{outside}')
    mkdir_if_not_exist(f'{annotations_path}/{outside}')
    load_outside_data(outisdepath, outside, images_path, annotations_path)
    
def add_nohuman_data(
    abs_path,
    images_path,
    annotations_path,
    sub_path='validation',
    addition=0.2):
    
    nohuman = 'nohuman'
    counter = 0
    content_nohumn = os.listdir(f'{annotations_path}/{nohuman}')
    images_nohuman = os.listdir(f'{images_path}/{nohuman}')
    print(len(content_nohumn))
    print(len(images_nohuman))
    
    content = os.listdir(f'{annotations_path}/{sub_path}')
    images = os.listdir(f'{images_path}/{sub_path}')
    print(len(content))
    print(len(images))
    
    amount = int(addition*len(content))+1
    print('adding', amount)
    
    sub_path_x_nohuman = f'{sub_path}_{addition}_nohuman'
    
    mkdir_if_not_exist(f'{images_path}/{sub_path_x_nohuman}')
    mkdir_if_not_exist(f'{annotations_path}/{sub_path_x_nohuman}')
    
    index = 0
    
    while counter < amount and index < len(content_nohumn):
        value = content_nohumn[index]
        split = re.split(r'\.(?!\d)', value)
        
        src_img = '/{0}/{1}/{2}/{3}.jpg'.format(abs_path, images_path, nohuman, split[0])
        src_y = '/{0}/{1}/{2}/{3}.png'.format(abs_path, annotations_path, nohuman, split[0])
        
        dist_img_folder = '/{0}/{1}/{2}/'.format(abs_path, images_path, sub_path_x_nohuman)
        dist_y_folder = '/{0}/{1}/{2}/'.format(abs_path, annotations_path, sub_path_x_nohuman)
        
        if (random.randrange(0,1) == 0):
            counter += 1
            shutil.move(src_y, dist_y_folder)
            shutil.move(src_img, dist_img_folder)
            time.sleep(0.001)
            sys.stdout.flush()
            sys.stdout.write('\r'+f"{counter}")
        
        index += 1
        
    sys.stdout.flush()
    sys.stdout.write('\r'+f"{counter}")
    
def move_images(
    abs_path,
    images_path,
    annotations_path,
    sub_path='training',
    addition=0.05):
    
    counter = 0
    content = os.listdir(f'{annotations_path}/{sub_path}')
    np.random.shuffle(content)
    
    print(len(content))
    
    amount = int(addition*len(content))
    print('adding', amount)
    
    sub_path_test = f'test'
    
    mkdir_if_not_exist(f'{images_path}/{sub_path_test}')
    mkdir_if_not_exist(f'{annotations_path}/{sub_path_test}')
    
    index = 0
    
    while counter < amount and index < len(content):
        value = content[index]
        split = re.split(r'\.(?!\d)', value)
        
        src_img = '/{0}/{1}/{2}/{3}.jpg'.format(abs_path, images_path, sub_path, split[0])
        src_y = '/{0}/{1}/{2}/{3}.png'.format(abs_path, annotations_path, sub_path, split[0])
        
        dist_img_folder = '/{0}/{1}/{2}/'.format(abs_path, images_path, sub_path_test)
        dist_y_folder = '/{0}/{1}/{2}/'.format(abs_path, annotations_path, sub_path_test)
        
        counter += 1
        shutil.move(src_y, dist_y_folder)
        shutil.move(src_img, dist_img_folder)
        time.sleep(0.001)
        sys.stdout.flush()
        sys.stdout.write('\r'+f"{counter} - {index}")
        
        index += 1
        
    sys.stdout.flush()
    sys.stdout.write('\r'+f"{counter} - {index}")
    
