import torch
from eztokenizer import EzTokenizer
import torch.nn as nn

tokenizer = EzTokenizer()

dataset = """


blub blub blubs blubby
ooh food
bloob blub blip blop
WAIT IS THAT A SHARK NO PLEZ HELP ME
Mama? i want food
beep blab bloop blob boloob food
no shark yipe!!!!!!!!!!!
PLEZ MR SHARK NO EAT ME TANKZ
ooh i like seaweed em nom nom
ooh coral reef wow
yummi coral reef nom nom chomp blub


"""

tokenizer.train(dataset)
enc = tokenizer.encode(dataset)

def shift(tokens, context_len: int):
    x_sequences = []
    y_sequences = []
    for position in range(len(tokens) - context_len):
        context = tokens[position:position + context_len]
        x_sequences.append(context)
        y_context = tokens[position+1: position+context_len+1]
        y_sequences.append(y_context)
    return x_sequences, y_sequences
x, y = shift(enc, 32)

x = torch.tensor(x)
y = torch.tensor(y)

embed_dim = 6
num_heads = 3

embedding = torch.nn.Embedding(len(tokenizer.vocab) + 1, 6)
attention = torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True)

output = torch.nn.Linear(6, len(tokenizer.vocab) + 1, bias=True)
loss = torch.nn.CrossEntropyLoss()

class Fih(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = embedding
        self.attention = attention
        self.output = output

    def forward_pass(self, x):

      input_tokens = x
      em_x = self.embedding(x)
      fihs_attention = self.attention(em_x, em_x, em_x)[0]
      fih_output = self.output(fihs_attention)
      return fih_output