import torch
from tqdm import tqdm
import torch.nn  as nn
import torch.optim as optim
from Unet.model import UNET
from imutils import paths
from Dataset.datasetloader import MRIDataset
from Utils import utils
from torch.utils.data import DataLoader
from Preprocessing.preprocessing import Preprocessor
import argparse
# import wandb


# Hyperparam etc.
LEARNING_RATE = 1e-3
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
BATCH_SIZE = 8
NUM_EPOCHS = 100
LOAD_MODEL = False
NUM_WORKERS = 0


# wandb.init(project='Unet', entity='grubi')
# config = wandb.config
# config.learning_rate = LEARNING_RATE


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


    for batch_idx, (data, targets) in enumerate(loop):
        data = data.to(DEVICE)
        targets = targets.to(DEVICE)

        # 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())
        wandb.log({"loss": loss})




ap = argparse.ArgumentParser()
ap.add_argument("-i", "--images", required=True, help="path to images directory")
ap.add_argument("-l", "--labels", required=True, help="path to labels directory")
ap.add_argument("-vi", "--val_images", required=True, help="path to images directory")
ap.add_argument("-vl", "--val_labels", required=True, help="path to labels directory")
args = vars(ap.parse_args())

print("[INFO] loading images and labels...")
imagePath = list(paths.list_images(args["images"]))
labelPath = list(paths.list_images(args["labels"]))
val_imagePath = list(paths.list_images(args["val_images"]))
val_labelPath = list(paths.list_images(args["val_labels"]))

prep = Preprocessor(128, 256)
dl = DataLoader(MRIDataset(imgpath=imagePath, labelpath=labelPath, preprocessors=[prep], verbose=200),
                batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS)
dl_val = DataLoader(MRIDataset(imgpath=val_imagePath, labelpath=val_labelPath, preprocessors=[prep], verbose=200),
                    batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS)

model = UNET(in_channels=1, out_channels=1).to(DEVICE).float()
loss_fnc = nn.BCEWithLogitsLoss()
optimizer = optim.RMSprop(model.parameters(), lr=LEARNING_RATE)

if LOAD_MODEL:
    utils.load_checkpoint(torch.load('tmp/checkpoint.pth.tar'), model)

scaler = torch.cuda.amp.GradScaler()

# wandb.watch(model)
for epoch in range(NUM_EPOCHS):
    print(epoch+1)

    train_fn(dl, model, optimizer, loss_fnc, scaler)

    checkpoint = {
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict()
    }

    utils.save_checkpoint(checkpoint,filename='tmp/checkpoint.pth.tar')
    utils.check_accuracy(dl_val, model, DEVICE)
    utils.save_predictions_as_imgs(dl_val, model, DEVICE)