import os
import torch
import torch.nn
import numpy as np
from PIL import Image
from skimage.io import imread
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split


def read_IAM_label_txt_file(file_txt_labels):
    """
    ---------
    Arguments
    ---------
    file_txt_labels : str
        full path to the text file containing labels

    -------
    Returns
    -------
    a tuple of
    all_image_files : list
        a list of all image file names
    all_labels : list
        a list of all labels
    """
    label_file_handler = open(file_txt_labels, mode="r")
    all_lines = label_file_handler.readlines()
    num_lines = len(all_lines)

    all_image_files = []
    all_labels = []

    for cur_line_num in range(num_lines):
        if cur_line_num % 3 == 0:
            all_image_files.append(all_lines[cur_line_num].strip())
        elif cur_line_num % 3 == 1:
            all_labels.append(all_lines[cur_line_num].strip())
        else:
            continue

    return all_image_files, all_labels


class HWRecogIAMDataset(Dataset):
    """
    Main dataset class to be used only for training, validation and internal testing
    """

    CHAR_SET = " !\"#&'()*+,-./0123456789:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
    CHAR_2_LABEL = {char: i + 1 for i, char in enumerate(CHAR_SET)}
    LABEL_2_CHAR = {label: char for char, label in CHAR_2_LABEL.items()}

    def __init__(
        self,
        list_image_files,
        list_labels,
        dir_images,
        image_height=32,
        image_width=768,
        which_set="train",
    ):
        """
        ---------
        Arguments
        ---------
        list_image_files : list
            list of image files
        list_labels : list
            list of labels
        dir_images : str
            full path to directory containing images
        image_height : int
            image height (default: 32)
        image_width : int
            image width (default: 768)
        which_set : str
            a string indicating which set is being used (default: train)
        """
        self.list_labels = list_labels
        self.dir_images = dir_images
        self.list_image_files = list_image_files
        self.image_width = image_width
        self.image_height = image_height
        self.which_set = which_set

        if self.which_set == "train":
            # apply data augmentation only for train set
            self.transform = transforms.Compose(
                [
                    transforms.ToPILImage(),
                    transforms.Resize(
                        (self.image_height, self.image_width), Image.BILINEAR
                    ),
                    transforms.RandomAffine(
                        degrees=[-0.75, 0.75],
                        translate=[0, 0.05],
                        scale=[0.75, 1],
                        shear=[-35, 35],
                        interpolation=transforms.InterpolationMode.BILINEAR,
                        fill=255,
                    ),
                    transforms.ToTensor(),
                    transforms.Normalize(
                        mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225],
                    ),
                ]
            )
        else:
            self.transform = transforms.Compose(
                [
                    transforms.ToPILImage(),
                    transforms.Resize(
                        (self.image_height, self.image_width), Image.BILINEAR
                    ),
                    transforms.ToTensor(),
                    transforms.Normalize(
                        mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225],
                    ),
                ]
            )

    def __len__(self):
        return len(self.list_image_files)

    def __getitem__(self, idx):
        image_file_name = self.list_image_files[idx]
        image_gray = imread(os.path.join(self.dir_images, image_file_name))
        image_3_channel = np.repeat(np.expand_dims(image_gray, -1), 3, -1)
        image_3_channel = self.transform(image_3_channel)

        label_string = self.list_labels[idx]
        label_encoded = [self.CHAR_2_LABEL[c] for c in label_string]
        label_length = [len(label_encoded)]

        label_encoded = torch.LongTensor(label_encoded)
        label_length = torch.LongTensor(label_length)

        return image_3_channel, label_encoded, label_length


def IAM_collate_fn(batch):
    """
    collate function

    ---------
    Arguments
    ---------
    batch : tuple
        a batch of input data as a tuple

    -------
    Returns
    -------
    a collated tuple of
    images : tensor
        tensor of batch images
    labels : tensor
        tensor of batch labels
    label_lengths : tensor
        tensor of batch label lengths
    """
    images, labels, label_lengths = zip(*batch)
    images = torch.stack(images, 0)
    labels = torch.cat(labels, 0)
    label_lengths = torch.cat(label_lengths, 0)
    return images, labels, label_lengths


def split_dataset(file_txt_labels, for_train=True):
    """
    ---------
    Arguments
    ---------
    file_txt_labels : str
        full path to the text file containing labels
    for_train : bool
        indicating whether split is for training or internal testing

    -------
    Returns
    -------
    a tuple of files depending for train or internal testing
    """
    all_image_files, all_labels = read_IAM_label_txt_file(file_txt_labels)
    train_image_files, test_image_files, train_labels, test_labels = train_test_split(
        all_image_files, all_labels, test_size=0.1, random_state=4
    )
    train_image_files, valid_image_files, train_labels, valid_labels = train_test_split(
        train_image_files, train_labels, test_size=0.1, random_state=4
    )
    if for_train:
        return train_image_files, valid_image_files, train_labels, valid_labels
    else:
        return test_image_files, test_labels


def get_dataloaders_for_training(
    train_x,
    train_y,
    valid_x,
    valid_y,
    dir_images,
    image_height=32,
    image_width=768,
    batch_size=8,
):
    """
    ---------
    Arguments
    ---------
    train_x : list
        list of train file names
    train_y : list
        list of train labels
    valid_x : list
        list of validation file names
    valid_y : list
        list of validation labels
    dir_images : str
        full directory path containing the images
    image_height : int
        image height (default: 32)
    image_width : int
        image width (default: 768)
    batch_size : int
        batch size (default: 8)

    -------
    Returns
    -------
    a tuple of dataloaders objects
    train_loader : object
        object of train set dataloader
    valid_loader : object
        object of validation set dataloader
    """
    train_dataset = HWRecogIAMDataset(
        train_x,
        train_y,
        dir_images,
        image_height=image_height,
        image_width=image_width,
        which_set="train",
    )
    valid_dataset = HWRecogIAMDataset(
        valid_x,
        valid_y,
        dir_images,
        image_height=image_height,
        image_width=image_width,
        which_set="valid",
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=4,
        collate_fn=IAM_collate_fn,
    )
    valid_loader = DataLoader(
        valid_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        collate_fn=IAM_collate_fn,
    )
    return train_loader, valid_loader


def get_dataloader_for_testing(
    test_x, test_y, dir_images, image_height=32, image_width=768, batch_size=1
):
    """
    ---------
    Arguments
    ---------
    test_x : list
        list of test file names
    test_y : list
        list of test labels
    dir_images : str
        full directory path containing the images
    image_height : int
        image height (default: 32)
    image_width : int
        image width (default: 768)
    batch_size : int
        batch size (default: 1)

    -------
    Returns
    -------
    test_loader : object
        object of test set dataloader
    """
    test_dataset = HWRecogIAMDataset(
        test_x,
        test_y,
        dir_images=dir_images,
        image_height=image_height,
        image_width=image_width,
        which_set="test",
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        collate_fn=IAM_collate_fn,
    )
    return test_loader