import json
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torch.nn import TransformerDecoder, TransformerDecoderLayer
from torch.nn.functional import cross_entropy
from keras.preprocessing.sequence import pad_sequences
from keras.preprocessing.text import Tokenizer





NUM_TOKENS = 6000
EPOCHS = 100
BATCH_SIZE = 96
MAX_INSTRUCTION_LENGTH = 50
MAX_OUTPUT_LENGTH = 50  



# Check if CUDA (GPU support) is available and set device accordingly
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_data(json_file_path):
    with open(json_file_path, "r") as f:
        return json.load(f)

json_file_path = "C:/Users/Dell-PC/Desktop/Projets/Datasets/code_alpaca_6k.json"
data = load_data(json_file_path)
instructions = [item["instruction"] for item in data]
outputs = ["<start> " + item["output"] + " <end>" for item in data]

tokenizer = Tokenizer(num_words=NUM_TOKENS, filters="", lower=False)
tokenizer.fit_on_texts(instructions + outputs)

instr_sequences = tokenizer.texts_to_sequences(instructions)
out_sequences = tokenizer.texts_to_sequences(outputs)

instr_padded = pad_sequences(instr_sequences, maxlen=MAX_INSTRUCTION_LENGTH, padding="post")
output_padded = pad_sequences(out_sequences, maxlen=MAX_OUTPUT_LENGTH, padding="post")


class CodeAlpacaDataset(Dataset):
    def init(self, instruction_padded, decoder_input_data, decoder_target_data):
        self.instruction_padded = instruction_padded
        self.decoder_input_data = decoder_input_data
        self.decoder_target_data = decoder_target_data

    def len(self):
        return len(self.instruction_padded)

    def getitem(self, idx):
        return (torch.tensor(self.instruction_padded[idx]),
                torch.tensor(self.decoder_input_data[idx]),
                torch.tensor(self.decoder_target_data[idx].flatten()))




# Assuming instr_padded, decoder_input_data, and decoder_target_data are numpy arrays from the Keras code above
train_dataset = CodeAlpacaDataset(instr_padded, decoder_input_data, decoder_target_data)

# Create a DataLoader instance
data_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)


class CustomTransformerDecoder(nn.Module):
    def init(self, vocab_size, d_model, nhead, num_decoder_layers, dim_feedforward, max_seq_length):
        super(CustomTransformerDecoder, self).init()
        self.embeddings = nn.Embedding(vocab_size, d_model)
        self.positional_encodings = nn.Parameter(torch.zeros(max_seq_length, d_model))
        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward)
        self.transformer_decoder = TransformerDecoder(decoder_layer, num_decoder_layers)
        self.output_layer = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt, tgt_mask):
        src_embeddings = self.embeddings(src) + self.positional_encodings[:src.size(1)]
        tgt_embeddings = self.embeddings(tgt) + self.positional_encodings[:tgt.size(1)]
        memory = self.transformer_decoder(src_embeddings, tgt_embeddings, tgt_mask)
        out = self.output_layer(memory)
        return out

def create_target_mask(size):
    # Mask out subsequent positions (to prevent lookahead)
    mask = torch.triu(torch.ones(size, size) * float("-inf"), diagonal=1)
    return mask

# Model parameters - these would be more carefully chosen in a real application
vocab_size = 20000
d_model = 512
nhead = 4
num_decoder_layers = 6
dim_feedforward = 1024
max_seq_length = 20

# Create the model
model = CustomTransformerDecoder(vocab_size, d_model, nhead, num_decoder_layers, dim_feedforward, max_seq_length)

# Prepare data
# Define your custom dataset, data loaders, etc.

# Define loss function and optimizer
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0)
scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=2.5e-4, total_steps=2000)

# Training loop
for epoch in range(EPOCHS):
    model.train()
    total_loss = 0
    
    for instruction_seqs, decoder_inputs, targets in data_loader:
        optimizer.zero_grad()
        instruction_seqs = instruction_seqs.to(device)  # Move to the right device (e.g., GPU or CPU)
        decoder_inputs = decoder_inputs.to(device)
        targets = targets.to(device)
        
        # Create the mask for the transformer
        tgt_mask = create_target_mask(decoder_inputs.size(1)).to(device)
        
        # Perform the forward pass
        outputs = model(instruction_seqs, decoder_inputs, tgt_mask)
        outputs = outputs.view(-1, vocab_size)  # Flatten the output for cross-entropy
        
        # Calculate loss and backpropagate errors
        loss = loss_function(outputs, targets)
        loss.backward()
        optimizer.step()
        scheduler.step()
        
        total_loss += loss.item()
    
    print(f"Epoch: {epoch}, Loss: {total_loss / len(data_loader):.4f}")