import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import random

num_images = 1

class CustomTensorDataset(Dataset):
    def __init__(self, path):
        self.data = torch.load(path)

    def __len__(self):
        return len(self.data['Entries']) * 1000

    def __getitem__(self, idx):
        idx = idx % len(self.data['Entries'])
        label = torch.zeros(self.data['NumLabels'], dtype=torch.float)
        entry = self.data['Entries'][idx]
        embedding = get_random_parts(entry['f'])
        label[entry['l']] = 1
        mask = torch.ones(num_images, dtype=torch.bool)
        for idx in range(0, embedding.size(0)):
            mask[idx] = False

        current_length = embedding.size(0)
    
        if current_length < num_images:
            padding_size = num_images - current_length
            embedding = F.pad(embedding, (0, 0, 0, padding_size))
        return {
            'label': label,
            'embeddings': embedding,
            'mask': mask
		}
    
def get_random_parts(tensor):
    length = random.randint(1, num_images)  # determine the random length
    start = random.randint(0, max(0, tensor.size(0) - length))  # determine the starting position
    return tensor[start:start + length]

class MultiTaskModel(nn.Module):
    def __init__(self, input_dim, num_labels, style_dim):
        super(MultiTaskModel, self).__init__()
        #self.style_transformer = StyleTransformer(input_dim, num_heads=8, num_layers=4, style_dim=style_dim)
        #style_dim = style_dim*2
        self.smth = nn.Sequential(
            nn.Linear(input_dim, style_dim),
            nn.ReLU(),
            nn.Linear(style_dim, style_dim),
            nn.ReLU(),
            nn.Linear(style_dim, style_dim),
            nn.ReLU(),
            nn.Linear(style_dim, style_dim),
            nn.ReLU(),
            nn.Linear(style_dim, style_dim),
            nn.ReLU(),
            nn.Linear(style_dim, style_dim),
            )
        self.fc_label = nn.Linear(style_dim, num_labels)

    def forward(self, x, mask):
        #style_embedding = F.normalize(self.style_transformer(x, mask))
        style_embedding = (self.smth(x.squeeze(1)))
        logits = self.fc_label(style_embedding)
        return logits 

model = MultiTaskModel(768, dataset.data['NumLabels'], 768).to('cuda')

optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss().to('cuda')
#dataloader = DataLoader(dataset, batch_size=dataset.data['NumLabels']*2, shuffle=True)
dataloader = DataLoader(dataset, dataset.data['NumLabels']*16, shuffle=True, drop_last=True)
#model.style_transformer.load_state_dict(torch.load("G:/style ai/model.pth"))


for epoch in range(1000):
    model.train()
    
    for batch in dataloader:
        optimizer.zero_grad()
        labels = batch['label'].to('cuda')
        embeddings = batch['embeddings'].to('cuda')
        masks = batch['mask'].to('cuda')
        logits = model(embeddings, masks)
        
        loss = criterion(logits, labels)
        
        loss.backward()
        optimizer.step()
        print(f'Epoch: {epoch+1}, Loss: {loss.item()}')




Model not learning much, loss decreasing extremly slow, label count = 1700