from typing import Sequence
import torch
import torchvision
from torchvision.transforms import functional as F
from torchvision import models, transforms
from PIL import Image


def pad_to_square(img):
    w, h = img.size
    max_wh = max(w, h)
    pad = ((max_wh - w) // 2, (max_wh - h) // 2)
    padding = (pad[0], pad[1], max_wh - w - pad[0], max_wh - h - pad[1])
    return F.pad(img, padding, fill=0, padding_mode='constant')


def build_batch(paths: Sequence[str], size=(224, 224), additional_transforms=[]) -> torch.Tensor:
    """Exercise 1.1

    @param paths: A sequence (e.g. list) of strings, each specifying the location of an image file.
    @param transform: One or multiple image transformations for augmenting the batch images.
    @return: Returns one single tensor that contains every image.
    """
    preprocess = transforms.Compose([
        transforms.Lambda(pad_to_square),
        transforms.Resize(size)]
        + additional_transforms
        +
        [transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406],
                              std=[0.229, 0.224, 0.225]),
         ])
    imgs = []
    
    for path in paths:
        img = Image.open(path).convert('RGB')
        img = preprocess(img)
        imgs.append(img)
    batch = torch.stack(imgs)
    return batch


def get_model() -> torch.nn.Module:
    """Exercise 1.2

    @return: Returns a neural network, initialised with pretrained weights.
    """
    model = models.resnet18(
        weights=models.ResNet18_Weights.DEFAULT)
    return model


def forward_pass(paths, batch, model):
    with torch.no_grad():
        outputs = model(batch)

    max_scores, preds = outputs.max(dim=1)

    class_names = torchvision.models.ResNet18_Weights.DEFAULT.meta["categories"]

    for i, (p, s) in enumerate(zip(preds, max_scores)):
        print(f"Image: {paths[i]}")
        print(f"Model output score: {s.item():.4f}")
        print(f"Predicted class: {class_names[p.item()]}")
        print()


def main():
    """Exercise 1.3

    Put all your code for exercise 1.3 here.
    """

    paths = [
        "./images/golden retriever.jpg",
        "./images/koala.jpg",
        "./images/pacifier.jpg",
        "./images/rubber duck sculpture.jpg",
        "./images/rubber ducks.jpg",
        "./images/shoehorn.jpg",
        "./images/zoo.jpg",
    ]
    batch_a = build_batch(paths)
    model = get_model()
    print("Batch A:")
    forward_pass(paths, batch_a, model)

    print("Batch B:")
    batch_b = build_batch(paths, (400, 400))
    forward_pass(paths, batch_b, model)

    print("Batch C:")
    batch_c = build_batch(paths, additional_transforms=[
                          transforms.RandomVerticalFlip(1)])
    forward_pass(paths, batch_c, model)

if __name__ == "__main__":
    main()