from __future__ import annotations

from dataclasses import dataclass
from num2words import num2words
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset
import pickle
from torchtext.data import get_tokenizer
from tqdm import tqdm
import torch.nn.functional as F
from typing import Dict, List, Set, Tuple

import torch
import torch.nn as nn


@dataclass
class HarryPotterBooks:
    path: str
    raw: str
    tokens: List[str]
    vocab: Set[str]
    word2id: Dict[str, int]
    id2word: Dict[int, str]

    @classmethod
    def from_file(cls, path: str, fraction: float = 1) -> HarryPotterBooks:
        with open(path, "r") as f:
            raw = f.read()
            raw = " ".join(raw.split())

        numbers = lambda t: num2words(t) if type(t) in [int, float] else t
        tokenizer = get_tokenizer("basic_english")
        tokens = tokenizer(raw)
        tokens = list(map(numbers, tokens))
        tokens = tokens[:int(len(tokens) * fraction)]

        vocab = set(tokens)

        word2id = {t: i for i, t in enumerate(vocab)}
        id2word = {i: t for i, t in enumerate(vocab)}
        
        return cls(path, raw, tokens, vocab, word2id, id2word)


class HarryPotterDataset(Dataset):
    def __init__(self, hpb: HarryPotterBooks, window: int) -> None:
        super(HarryPotterDataset, self).__init__()
        self.hpb = hpb
        self.window = window

    def __len__(self) -> int:
        return len(self.hpb.tokens) - (self.window + 1)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        data = self.hpb.tokens[idx:idx + self.window + 1]
        sentence_tokens, word_token = data[:-1], data[-1]

        sentence = torch.tensor([
            self.hpb.word2id[t] for t in sentence_tokens
        ], dtype=torch.long)
        word = self.hpb.word2id[word_token]
        
        return sentence, word



class JKRowling(nn.Module):
    def __init__(
        self,
        vocab: int,
        embedding: int,
        hidden: int,
        layers: int,
        dropout: float,
    ) -> None:
        super(JKRowling, self).__init__()
        self.embedding = nn.Embedding(vocab, embedding)
        self.lstm = nn.LSTM(
            embedding,
            hidden,
            num_layers=layers,
            dropout=dropout,
            batch_first=True,
        )
        self.do  = nn.Dropout(p=0.5)
        self.act = nn.ReLU()
        self.fc = nn.Linear(hidden, vocab)

    def forward(self, sentence: torch.Tensor) -> torch.Tensor:
        embedding = self.embedding(sentence)
        _, (hn, cn) = self.lstm(embedding)
        word = self.fc(hn[-1])
        
        return word


ROOT = "./dataset/full.txt"

FRACTION = 1.0
MAX_SEQ = 50

EMBEDDING = 64
HIDDEN = 64
LAYERS = 2
DROPOUT = 0.2

BATCH_SIZE = 256
NUM_WORKERS = 4

EPOCHS = 10
LR = 1e-3
DECAY = 1e-4
VOCAB = "vocab.pkl"

with open(VOCAB, "rb") as f:
    hpb = HarryPotterBooks(*pickle.load(f))
hpd = HarryPotterDataset(hpb, window=MAX_SEQ)
loader = DataLoader(
    hpd,
    batch_size=BATCH_SIZE,
    shuffle=False,
    num_workers=NUM_WORKERS,
    pin_memory=True,
)


print("===== STATS ====")
print("RAW    :", len(hpb.raw))
print("TOKENS :", len(hpb.tokens))
print("VOCAB  :", len(hpb.vocab))
print("DATASET:", len(hpd))
print("================")
print()




model = JKRowling(len(hpb.vocab), EMBEDDING, HIDDEN, LAYERS, DROPOUT).cuda()
criterion = nn.CrossEntropyLoss().cuda()

optim = AdamW(model.parameters(), lr=LR, weight_decay=DECAY)

def perplexity(model: torch.nn.Module, dataset: HarryPotterDataset, start: int = 50) -> float:
    n = 10
    prod = 1
    for k in range(start, start + n):
        sentence_ids, wid = dataset[k]

        _word_id = F.softmax(model(sentence_ids.unsqueeze(0).cuda()).squeeze(0).cpu(), dim=0)
        wk = _word_id[wid].item()

        prod *= 1/wk
    return prod ** 1/float(n)


if __name__ == "__main__":

    for epoch in range(EPOCHS):
        with tqdm(loader, desc="Train") as pbar:
            total_loss = 0.0
            model = model.train()
            
            for sentence, word in pbar:
                sentence, word = sentence.cuda(), word.cuda()
                optim.zero_grad()

                out = model(sentence)
                loss = criterion(out, word)
                loss.backward()
                optim.step()

                total_loss += loss.item() / len(loader)
                pbar.set_postfix(loss=total_loss)

        with torch.no_grad():
            model = model.eval()

            print(f'PP(JKRowling) = {perplexity(model, hpd, 10)}')

            for idx, i in enumerate([0, 10, 100]):
                sentence_ids, word_id = hpd[i]

                _word_id = model(sentence_ids.unsqueeze(0).cuda())
                _word_id = torch.argmax(_word_id.squeeze(0)).cpu().item()
                
                sentence = [hpb.id2word[i.item()] for i in sentence_ids]
                word = hpb.id2word[word_id]
                _word = hpb.id2word[_word_id]

                print()
                print(f"===== TEST{idx + 1} ====")
                print("SENTENCE  :", " ".join(sentence))
                print("TARGET    :", word)
                print("PREDICTION:", _word)
                print("================")
                print()

    m = torch.jit.script(model)
    m.save(f"Challenge02_TEST_TEST.ts") 