import os
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import torch.optim as optim
import torch.nn as nn
import json
from torchvision import models
from utils import compute_batch
from elasticsearch_utils import ESAPI



datasetdir = '../catalog_images_17k-Processed'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_modified = models.vgg16(pretrained=True)
for param in model_modified.parameters():
    param.requires_grad = False
model_modified.classifier = torch.nn.Sequential(nn.Linear(25088, 4096, bias=True),
                                       nn.ReLU(inplace=True),
                                       nn.Dropout(),
                                       nn.Linear(4096, 100, bias=True))
model_modified.to(device)

es = ESAPI()

#Prepare inputs
all = np.arange(len(os.listdir(datasetdir)))
no_label_ids = es.get_all_taxonomy_ids('no-label')
all = np.setdiff1d(all, no_label_ids)
train = np.random.choice(all , size = int(len(all)*0.8), replace=False)
val = np.setdiff1d(all, train)


word_vectors = torch.load('../Word-Vectors/word_vectors.pt')
with open('../ES-Data/documents.json', 'r') as file:
  documents = json.load(file)


class FashionDataset(Dataset):

    def __init__(self, images):
        self.dir = datasetdir
        self.images = images

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

    def __getitem__(self, idx):
        image_id = self.images[idx]
        image_id = str(image_id)
        image_name = documents[image_id]['file']
        input_batch = compute_batch(os.path.join(self.dir, image_name))
        taxonomy = documents[image_id]['taxonomy']
        taxonomy = word_vectors[taxonomy]
        sample = {'input_batch': input_batch, 'taxonomy': taxonomy}
        return sample


train_ds = FashionDataset(train)
val_ds = FashionDataset(val)
train_dl = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4)
val_dl = DataLoader(val_ds, batch_size=32, shuffle=True, num_workers=4)


#Train model
'''
dataloaders -> [train,val]
'''
def train_model(model, dataloaders, criterion, optimizer, num_epochs=40):
    name = 'best-word-model.pt'
    best_loss = 100.0

    all_train_loss = []
    all_val_loss = []
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)
        # Each epoch has a training and validation phase
        for i, dataloader in enumerate(dataloaders):
            if i == 0:
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode
            running_loss = 0.0
            # Iterate over data.
            for j, samples in enumerate(dataloader):
                print(f"Epoch {epoch}: Batch {j} / {len(dataloader)}")
                inputs_batch = samples['input_batch']
                labels = samples['taxonomy']
                inputs_batch = inputs_batch.to(device)
                labels = labels.to(device)
                # zero the parameter gradients
                optimizer.zero_grad()
                # forward
                # track history if only in train
                with torch.set_grad_enabled(i == 0):
                    # Get model outputs and calculate loss
                    outputs = model(inputs_batch)
                    loss = criterion(outputs, labels)
                    # backward + optimize only if in training phase
                    if i == 0:
                        loss.backward()
                        optimizer.step()
                # statistics
                running_loss += loss.item() * inputs_batch.size(0)
            epoch_loss = running_loss / len(dataloader.dataset)
            if i == 0:
                phase = 'Train'
            else:
                phase = 'Val'
            print('{} Loss: {:.4f}'.format(phase, epoch_loss))
            if phase == 'Val' and epoch_loss < best_loss:
                print('------------- Model updated ---------------')
                best_loss = epoch_loss
                torch.save(model.state_dict(), name)
            if phase == 'Val':
                all_val_loss.append(epoch_loss)
            if phase == 'Train':
                all_train_loss.append(epoch_loss)
        print()

    print('Best val loss: {:4f}'.format(best_loss))
    print('Train:', all_train_loss)
    print('Val:', all_val_loss)
    model.load_state_dict(torch.load(name))
    return model


params_to_update = []
for name, param in model_modified.named_parameters():
    if param.requires_grad:
        params_to_update.append(param)
        print(name, 'will be updated.')

optimizer = optim.SGD(params_to_update, lr=0.001, momentum=0.9)
criterion = nn.MSELoss()
model_modified = train_model(model_modified, [train_dl, val_dl], criterion, optimizer)