import csv
import os
from PIL import Image
from torchvision import transforms
from torchvision.utils import save_image
import torch


def transform_image(img_in, target_dir, transformation, suffix):
    """
    Transforms an image according to provided transformation.

        Parameters:
            img_in (path): Image to transform
            target_dir (path): Destination path
            transformation (callable): Transformation to be applied
            suffix (str): Suffix of resulting image.

        Returns:
            binary_sum (str): Binary string of the sum of a and b
    """
    if suffix == 'rot':
        im = Image.open(img_in)
        im = im.rotate(270)
        tensor = transforms.ToTensor()(im)
        save_image(tensor, target_dir + os.sep + suffix + '.jpg')
    elif suffix == 'sobel':
        im = Image.open(img_in)
        tensor = transforms.ToTensor()(im)
        sobel_filter = torch.tensor([[1., 2., 1.], [0., 0., 0.], [-1., -2., -1.]])
        f = sobel_filter.expand(1, 3, 3, 3)
        tensor = torch.conv2d(tensor, f, stride=1, padding=1 )
        save_image(tensor, target_dir + os.sep + suffix + '.jpg')
    elif suffix == 'noise':
        im = Image.open(img_in)
        tensor = transforms.ToTensor()(im)
        tensor = tensor + (torch.randn(tensor.size()) * 0.2 + 0)
        save_image(tensor, target_dir + os.sep + suffix + '.jpg')
    elif suffix == 'cutout':
        print("asd")
    else:
        im = Image.open(img_in)
        im_trans = transformation(im)
        im_trans.save(target_dir + os.sep + suffix + '.jpg')


def sort_and_rename_images(excel_path: str):
    """Renames images and sorts them according to csv."""
    base_dir = excel_path.rsplit(os.sep, 1)[0]
    dir_all = base_dir + os.sep + 'all'
    if not os.path.isdir(dir_all):
        os.mkdir(dir_all)
    dir_hem = base_dir + os.sep + 'hem'
    if not os.path.isdir(dir_hem):
        os.mkdir(dir_hem)

    with open(excel_path, mode='r') as file:
        csv_file = csv.reader(file)
        for lines in csv_file:
            print(lines)
            if lines[2] == '1':
                os.rename(base_dir + os.sep + lines[1], dir_all + os.sep + lines[0])
            elif lines[2] == '0':
                os.rename(base_dir + os.sep + lines[1], dir_hem + os.sep + lines[0])


def drop_color_channels(source_dir, target_dir, rgb):
    """Rotates all images in in source dir."""
    if rgb == 0:
        suffix = "red_only"
        drop_1 = 1
        drop_2 = 2
    elif rgb == 1:
        suffix = "green_only"
        drop_1 = 0
        drop_2 = 2
    elif rgb == 2:
        suffix = "blue_only"
        drop_1 = 0
        drop_2 = 1
    elif rgb == 3:
        suffix = "no_red"
        drop_1 = 0
    elif rgb == 4:
        suffix = "no_green"
        drop_1 = 1
    elif rgb == 5:
        suffix = "no_blue"
        drop_1 = 2
    else:
        suffix = ""
        print("Invalid RGB-channel")
    if suffix != "":
        dirs = os.listdir(source_dir)
        for item in dirs:
            if os.path.isfile(source_dir + os.sep + item):
                im = Image.open(source_dir + os.sep + item)
                tensor = transforms.ToTensor()(im)
                tensor[drop_1, :, :] = 0
                if rgb < 3:
                    tensor[drop_2, :, :] = 0
                save_image(tensor, target_dir + os.sep + item, 'bmp')


def rotate_images(target_dir, source_dir, rotate, theta):
    """Rotates all images in in source dir."""
    dirs = os.listdir(source_dir)
    for item in dirs:
        if os.path.isfile(source_dir + os.sep + item):
            for i in range(0, rotate):
                im = Image.open(source_dir + os.sep + item)
                im = im.rotate(i*theta)
                tensor = transforms.ToTensor()(im)
                save_image(tensor, target_dir + os.sep + str(i) + '_' + item, 'bmp')


def grayscale_image(source_dir, target_dir):
    """Grayscale transforms all images in path."""
    t = transforms.Grayscale()
    dirs = os.listdir(source_dir)
    if not os.path.isdir(target_dir):
        os.mkdir(target_dir)
    for item in dirs:
        if os.path.isfile(source_dir + os.sep + item):
            im = Image.open(source_dir + os.sep + item).convert('RGB')
            im_resize = t(im)
            tensor = transforms.ToTensor()(im_resize)
            padding = torch.zeros(1, tensor.shape[1], tensor.shape[2])
            tensor = torch.cat((tensor, padding), 0)
            im_resize.save(target_dir + os.sep + item, 'bmp')


def resize(source_dir):
    """Rotates all images in in source dir."""
    t = transforms.Compose([transforms.Resize((128, 128))])
    dirs = os.listdir(source_dir)
    target_dir = source_dir + os.sep + 'resized'
    if not os.path.isdir(target_dir):
        os.mkdir(target_dir)
    for item in dirs:
        if os.path.isfile(source_dir + os.sep + item):
            im = Image.open(source_dir + os.sep + item)
            im_resize = t(im)
            im_resize.save(source_dir + os.sep + 'resized' + os.sep + item, 'bmp')


def crop_image(source_dir):
    """Center Crops all images in path."""
    t = transforms.CenterCrop((224, 224))
    dirs = os.listdir(source_dir)
    target_dir = source_dir + os.sep + 'cropped'
    if not os.path.isdir(target_dir):
        os.mkdir(target_dir)
    for item in dirs:
        if os.path.isfile(source_dir + os.sep + item):
            im = Image.open(source_dir + os.sep + item)
            im_resize = t(im, )
            im_resize.save(source_dir + os.sep + 'cropped' + os.sep + item, 'bmp')


def mk_dirs(target_dir):
    dir_0 = target_dir + r"\fold_0"
    dir_1 = target_dir + r"\fold_1"
    dir_2 = target_dir + r"\fold_2"
    dir_3 = target_dir + r"\phase2"
    dir_4 = target_dir + r"\phase3"
    dir_0_all = dir_0 + r"\all"
    dir_0_hem = dir_0 + r"\hem"
    dir_1_all = dir_1 + r"\all"
    dir_1_hem = dir_1 + r"\hem"
    dir_2_all = dir_2 + r"\all"
    dir_2_hem = dir_2 + r"\hem"
    if not os.path.isdir(dir_0):
        os.mkdir(dir_0)
    if not os.path.isdir(dir_1):
        os.mkdir(dir_1)
    if not os.path.isdir(dir_2):
        os.mkdir(dir_2)
    if not os.path.isdir(dir_3):
        os.mkdir(dir_3)
    if not os.path.isdir(dir_4):
        os.mkdir(dir_4)

    if not os.path.isdir(dir_0_all):
        os.mkdir(dir_0_all)
    if not os.path.isdir(dir_0_hem):
        os.mkdir(dir_0_hem)
    if not os.path.isdir(dir_1_all):
        os.mkdir(dir_1_all)
    if not os.path.isdir(dir_1_hem):
        os.mkdir(dir_1_hem)
    if not os.path.isdir(dir_2_all):
        os.mkdir(dir_2_all)
    if not os.path.isdir(dir_2_hem):
        os.mkdir(dir_2_hem)
    return dir_0_all, dir_0_hem, dir_1_all, dir_1_hem, dir_2_all, dir_2_hem, dir_3, dir_4