# coding: utf-8
""" data module """
import torch
import os
from torch.utils.data import Dataset


def get_shakespeare():
    """loads the shakespeare text"""
    return open(
        os.path.join(os.path.dirname(__file__), "data/shakespeare.txt"),
        "rt",
        encoding="utf-8",
    ).read()


def preprocess(text: str) -> tuple:
    """preprocesses the shakespeare text"""
    alphabet = sorted(set(text))
    letter_to_int = {let: ind for ind, let in enumerate(alphabet)}
    int_to_letter = dict(enumerate(alphabet))
    letter_ints = [letter_to_int[letter] for letter in text]
    alphabet_size = len(alphabet)
    return int_to_letter, letter_to_int, alphabet_size, letter_ints


class ShakespeareDataset(Dataset):
    """Shakespeare Dataset"""

    def __init__(self, seq_len: int):
        _, _, _, self.text = preprocess(get_shakespeare())
        self.data = torch.LongTensor(self.text[:-1])  # get the data
        self.labels = torch.LongTensor(self.text[1:])
        self.seq_len = seq_len  # set the sequence length

    def __len__(self) -> int:
        return (
            len(self.text) - self.seq_len - 1
        )  # length of corpora minus sequence length minus shift

    def __getitem__(self, index: int) -> tuple:
        return (
            self.data[index : index + self.seq_len],
            self.data[index : index + self.seq_len],
        )  # return tuple of (sample, label)
