import os
import torch
import matplotlib.pyplot as plt

from torch import nn
from torch import optim
from os.path import join
from torchvision import models
from tqdm.notebook import tqdm
from torchvision import transforms
from torch.nn import functional as F
from torchvision.datasets import ImageFolder
from torch.utils.data import Subset, DataLoader, random_split

plt.rcParams["figure.figsize"] = (12, 12)

# Clone dataset from GitHub
!git clone https://github.com/Horea94/Fruit-Images-Dataset

# Helper functions

def load_dataset(path):
    data = load_files(path)
    files = np.array(data['filenames'])
    targets = np.array(data['target'])
    target_labels = np.array(data['target_names'])
    return files,targets,target_labels
    
def get_default_device():
    """Pick GPU if available, else CPU"""
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')
    
def to_device(data, device):
    """Move tensor(s) to chosen device"""
    if isinstance(data, (list,tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)

def predict_and_plot(images, labels, plot_size=[1, 6]):
    prediction_limit = plot_size[0] * plot_size[1]
    output = []
    
    for idx, (image, label) in enumerate(zip(images, labels)):
        if idx >= prediction_limit:
            break
            
        xb = image.unsqueeze(0)
        xb = to_device(xb, device)
        yb = model(xb)
        _, preds  = torch.max(yb, dim=1)
        prediction = preds[0].item()
        output.append(prediction)
    
    result = torch.IntTensor(output)
    plot_img(images, labels, result)
    
def plot_img(images, labels, predictions=None, plot_size=[1, 6]):
    fig, axes = plt.subplots(*plot_size)
    fig.tight_layout()
    
    for idx, ax in enumerate(axes):
        if predictions != None:
            predictions_text = f'\nPrediction:\n{predictions[idx].item()}: {train_dataset.classes[predictions[idx]]}'
            ax.set_xlabel(predictions_text)
            ax.xaxis.label.set_color('green' if predictions[idx].item() == labels[idx].item() else 'red')
        ax.title.set_text(f'{labels[idx].item()}: {train_dataset.classes[labels[idx]]}')
        ax.imshow(images[idx].cpu().permute(1, 2, 0))
        
@torch.no_grad()
def predict_dl(dl, model):
    torch.cuda.empty_cache()
    batch_probs = []
    for xb, targets in tqdm(dl):
        probs = model(xb)
        _, preds = torch.max(probs, dim=1)
        correct = torch.sum(preds == targets).item()
        total = targets.size(0)
        batch_probs.append((correct, total))
    correct = sum(list(zip(*batch_probs))[0])
    total = sum(list(zip(*batch_probs))[1])
    return correct / total

def plot_chart(title, data_labels, data, x_label, y_label):
    fig = plt.figure(figsize=(20,10))
    plt.title(title)
    
    for data_label, data_list in zip(data_labels, data):
        plt.plot(data_list, label=data_label)

    plt.xlabel(x_label, fontsize=12)
    plt.ylabel(y_label, fontsize=12)
    plt.legend(loc='best')

class DeviceDataLoader():
    """Wrap a dataloader to move data to a device"""
    def __init__(self, dl, device):
        self.dl = dl
        self.device = device
        
    def __iter__(self):
        """Yield a batch of data after moving it to device"""
        for b in self.dl: 
            yield to_device(b, self.device)

    def __len__(self):
        """Number of batches"""
        return len(self.dl)

# Model class

class FruitModel(nn.Module):
    def __init__(self):
        super().__init__()
        # Use a pretrained model
        self.network = models.resnet18(pretrained=True)
        # Replace last layer
        num_ftrs = self.network.fc.in_features
        self.network.fc = nn.Linear(num_ftrs, len(train_dataset.classes))
    
    def forward(self, xb):
        return torch.sigmoid(self.network(xb))
    
    def training_step(self, batch):
        images, targets = batch 
        out = self(images)
        _, preds = torch.max(out, dim=1)
        correct = torch.sum(preds == targets).item()
        total = targets.size(0)
        loss = F.cross_entropy(out, targets)      
        return loss, correct, total
    
    def validation_step(self, batch):
        images, targets = batch 
        out = self(images)  # Generate predictions
        _, preds = torch.max(out, dim=1)
        correct = torch.sum(preds == targets).item()
        total = targets.size(0)
        loss = F.cross_entropy(out, targets)  # Calculate loss
        return {'val_loss': loss.detach(), 'val_correct': correct, 'val_total': total}
        
    def validation_epoch_end(self, outputs):
        batch_losses = [x['val_loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()   # Combine losses
        correct = sum([x['val_correct'] for x in outputs])
        total = sum([x['val_total'] for x in outputs])
        return {'val_loss': epoch_loss.item(), 'val_correct': correct, 'val_total': total}
    
    def epoch_end(self, epoch, result):
        print("Epoch [{}], train_loss: {:.4f}, val_loss: {:.4f}, train_acc: {:.2%}, val_acc: {:.2%}".format(
            epoch, result['train_loss'], result['val_loss'], result['train_correct'] / result['train_total'], result['val_correct'] / result['val_total']))