import torch
import torch.nn as nn
import torch.optim as optim
import json
import base64
from torch.utils.data import DataLoader, Dataset
import binascii
import chardet
import random
import math


class PositionalEncoding(nn.Module):
    def __init__(self, embedding_dim, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.embedding_dim = embedding_dim
        
        # Creating positional encoding matrix
        pe = torch.zeros(max_len, embedding_dim)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embedding_dim, 2).float() * (-math.log(10000.0) / embedding_dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        
        # Registering pe as a buffer since it’s not a model parameter
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return x

class DecoderLayer(nn.Module):
    def __init__(self, embedding_dim, heads, ff_dim):
        super(DecoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(embedding_dim, heads)
        
        # Feedforward network
        self.ffn = nn.Sequential(
            nn.Linear(embedding_dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, embedding_dim),
        )
        
        self.layer_norm1 = nn.LayerNorm(embedding_dim)
        self.layer_norm2 = nn.LayerNorm(embedding_dim)
    
    def forward(self, src):
        src2 = self.layer_norm1(src)
        attn_output, _ = self.self_attn(src2, src2, src2)
        src = src + attn_output
        src2 = self.layer_norm2(src)
        src = src + self.ffn(src2)
        return src

class Decoder(nn.Module):
    def __init__(self, vocab_size, embedding_dim, num_layers, heads, ff_dim):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.pos_encoding = PositionalEncoding(embedding_dim)
        self.decoder_layers = nn.ModuleList([DecoderLayer(embedding_dim, heads, ff_dim) for _ in range(num_layers)])
        self.final_layer = nn.Linear(embedding_dim, vocab_size)
    
    def forward(self, x):
        x = self.embedding(x)
        x = self.pos_encoding(x)
        
        for layer in self.decoder_layers:
            x = layer(x)
        
        output = self.final_layer(x)
        return output
  
class QAJsonlDataset(Dataset):
    def __init__(self, path, seq_len=512):
        super().__init__()
        self.pairs = self.load_data(path)
        self.seq_len = seq_len
        self.vocab = self.build_vocab([word for pair in self.pairs for item in pair for word in item])

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

    def __getitem__(self, idx):
        question, answer = self.pairs[idx]
        question_tokenized = self.tokenize([self.vocab[word] for word in question])
        answer_tokenized = self.tokenize([self.vocab[word] for word in answer])
        return torch.tensor(question_tokenized), torch.tensor(answer_tokenized)

    def load_data(self, path):
        pairs = []
        with open(path, "rb") as f:  # Open in binary mode
            for line in f:
                data = json.loads(line.decode("utf-8"))  # Decode binary data to JSON
                question, answer = self.decode_binary_data(data["user"], data["content"])  # Decode binary fields
                pairs.append((question.split(), answer.split()))
        return pairs

    def decode_binary_data(self, question_data, answer_data):
        try:
            question_encoding = chardet.detect(question_data.encode("utf-8"))['encoding']
            answer_encoding = chardet.detect(answer_data.encode("utf-8"))['encoding']
        # Decode data using the detected encoding
            question = base64.b64decode(question_data).decode(question_encoding)
            answer = base64.b64decode(answer_data).decode(answer_encoding)
        except (UnicodeDecodeError, binascii.Error):
        # Handle errors gracefully, e.g., log errors, skip entry
            print(f"Error decoding data: question - {question_data}, answer - {answer_data}")
            question = ""
            answer = ""
        return question, answer

    def tokenize(self, integer_sequence):
        tokens = integer_sequence[:self.seq_len]
        if len(tokens) < self.seq_len:
        # Pad with appropriate binary token
            padding_token = self.vocab["<pad>"]  # Consider a binary padding token if needed
            tokens.extend([padding_token] * (self.seq_len - len(tokens)))
        # Add <eos> token only if appropriate for binary representation
        return tokens

    def build_vocab(self, binary_tokens):
        unique_tokens = list(set(binary_tokens))
        self.vocab = {token: idx for idx, token in enumerate(unique_tokens)}
        self.vocab['<pad>'] = len(self.vocab)  # Add padding token
        self.pad_idx = self.vocab['<pad>']  # Store padding token index
        return self.vocab

class CustomDataLoader:
    def __init__(self, dataset, batch_size=32):
        self.dataset = dataset
        self.batch_size = batch_size

    def len(self):
        return len(self.dataset) // self.batch_size

    def __getitem__(self, idx):
        batch = self.dataset[idx * self.batch_size:(idx + 1) * self.batch_size]
        inputs, targets = zip(*batch)
        inputs_padded = torch.nn.utils.rnn.pad_sequence(inputs, batch_first=True, padding_value=0)
        targets_padded = torch.nn.utils.rnn.pad_sequence(targets, batch_first=True, padding_value=0)
        return inputs_padded, targets_padded

# Define model parameters
vocab_size = 10000  # Adjust based on your vocabulary size
embedding_dim = 128
num_layers = 2  # Number of decoder layers
heads = 2  # Number of attention heads
ff_dim = 512  # Feed-forward dimension

# Create decoder-only transformer model
model = Decoder(vocab_size, embedding_dim, num_layers, heads, ff_dim)

# Data loading and splitting
path_to_data = "data/Real_talk.jsonl"
dataset = QAJsonlDataset(path_to_data, 512)

# Shuffle the dataset
random.shuffle(dataset.pairs)

train_size = int(0.8 * len(dataset))  # 80% of data for training, adjust as necessary
val_size = len(dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=32)

# Define optimizer
optimizer = optim.Adam(model.parameters())
loss_fn = nn.CrossEntropyLoss()  # PyTorch version of SparseCategoricalCrossentropy with from_logits=True

def print_model_param_count(model):
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Total Trainable Parameters: {total_params}")

# Assuming your model’s variable name is model
print_model_param_count(model)

# Training loop (example)
for epoch in range(1, 6):
    for i, (inputs, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(inputs)
        loss = loss_fn(output.view(-1, vocab_size), targets.view(-1))  # Reshape for CrossEntropyLoss
        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch}, Loss: {loss.item()}')

def evaluate(model, val_loader, vocab):
    with torch.no_grad():
        for inputs, targets in val_loader:
            outputs = model(inputs)
            # Adjust dimension in argmax to correctly get predictions
            predictions = torch.argmax(outputs, dim=2)
            # Flatten tensors for element-wise comparison
            predictions_flat = predictions.view(-1)
            targets_flat = targets.view(-1)
            padding_value = dataset.pad_idx # Now accessible
            mask = targets_flat != padding_value

            # Use vocab to convert token IDs to words before calculating accuracy
            correct_words = [(vocab[pred.item()] if pred.item() != 0 else '') for pred in predictions_flat[mask]]
            accuracy = sum(word == target.item() for word, target in zip(correct_words, targets_flat[mask])) / mask.sum().item()

            return accuracy

# After training loop:
val_accuracy = evaluate(model, val_loader, dataset.vocab)
print(f'Validation Accuracy: {val_accuracy}')

def generate_text(model, initial_text, vocab, inv_vocab, seq_len=512, temperature=1.0):
    model.eval()  # Put the model in eval mode
    tokens = [vocab.get(word, vocab['<unk>']) for word in initial_text.split()]  # Convert initial text to tokens. Assumes &lt;unk&gt; for unknown words.
    with torch.no_grad():  # No need to track gradients
        for _ in range(seq_len - len(tokens)):
            input_tensor = torch.tensor([tokens], dtype=torch.long)  # Convert to tensor
            output = model(input_tensor)
            # Get the logits of the last token and apply temperature
            logits = output[0, -1, :] / temperature
            # Convert logits to probabilities (softmax) and sample the next token
            probabilities = torch.softmax(logits, dim=-1)
            next_token = torch.argmax(probabilities).item()
            tokens.append(next_token)  # Add predicted token to the sequence

            if next_token == vocab['<eos>']:  # Assuming you have an end-of-sequence token
                break

    generated_text = ' '.join([inv_vocab[token] for token in tokens])  # Convert tokens back to text
    return generated_text

# Access the vocabulary from the dataset after it's loaded
vocab = dataset.vocab

# Creating an inverse vocabulary for converting token IDs back to words
inv_vocab = {idx: token for token, idx in vocab.items()}  # vocab is your vocabulary dictionary from token to index

# Example usage
initial_text = "The meaning of life"  # Seed the generation with some text
generated_text = generate_text(model, initial_text, vocab, inv_vocab, seq_len=50)  # seq_len is the total desired length of generated sequence including the initial text

print(generated_text)