import torch
import torch.nn as nn
import torch.optim as optim

import albumentations as A
from albumentations.pytorch import ToTensorV2

from tqdm import tqdm
from UNET import UNET
from utils import (
    load_checkpoint,
    save_checkpoint,
    get_loaders,
    check_accuracy,
    save_predictions_as_images
)

learning_rate = 1e-04
batch_size = 16
num_epochs = 3
num_workers = 2
image_height = 160
image_width = 240
pin_memory = True
load_model = False
train_img_dir = 'dataset/image/'
train_mask_dir = 'dataset/mask/'
test_img_dir = 'dataset/test_image/'
test_mask_dir = 'dataset/test_mask/'

def train_fn(loader, model, optimizer, loss_fn, scaler):
    loop = tqdm(loader)

    for batch_idx, (data, targets) in enumerate(loop):
        targets = targets.float().unsqueeze(1)

        # Forward
        with torch.cuda.amp.autocast():
            predictions = model(data)
            loss = loss_fn(predictions, targets)

        # Backward
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        #update tqdm_loop
        loop.set_postfix(loss=loss.item())

def main():
    train_transform = A.Compose(
        [
            A.Resize(height=image_height, width=image_width),
            A.Rotate(limit=35, p=1.0),
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.1),
            A.Normalize(
                mean=[0.0,0.0,0.0],
                std=[1.0,1.0,1.0],
                max_pixel_value=255.0
            ),
            ToTensorV2()
        ]
    )

    test_transforms = A.Compose(
        [
            A.Resize(height=image_height, width=image_width),
            A.Normalize(
                        mean=[0.0,0.0,0.0],
                        std=[1.0,1.0,1.0],
                        max_pixel_value=255.0
            ),
            ToTensorV2()
        ]
    )

    model = UNET(in_channels=3, out_channels=1)
    loss_fn = nn.BCEWithLogitsLoss() #Since we are not doing Sigmoid on the output of the model.
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    train_loader, test_loader = get_loaders(
        train_img_dir,
        train_mask_dir,
        test_img_dir,
        test_mask_dir,
        batch_size,
        train_transform,
        test_transforms,
        num_workers,
        pin_memory
    )

    if load_model:
        load_checkpoint(torch.load('my_checkpoint.pth.tar'), model)

    check_accuracy(test_loader, model)
    scaler = torch.cuda.amp.GradScaler()
    for epoch in range(num_epochs):
        train_fn(train_loader, model, optimizer, loss_fn, scaler)
        checkpoint = {
            'state_dict': model.state_dict(),
            'optimizer': optimizer.state_dict()
        }
        save_checkpoint(checkpoint)
        check_accuracy(test_loader, model)
        save_predictions_as_images(test_loader, model, folder='saved_images/')


if __name__ == "__main__":
    main()