import torch

import torchvision

import os

def save_checkpoint(state, filename='tmp/checkpoint.pth.tar'):

    print('[INFO] Saving checkpoint')

    torch.save(state, filename)

def load_checkpoint(checkpoint, model):

    print('[INFO] Loading checkpoint')

    model.load_state_dict(checkpoint['state_dict'])

def check_accuracy(loader, model, device):

    num_correct = 0

    num_pixels = 0

    dice_score = 0

    with torch.no_grad():

        for x, y in loader:

            x = x.to(device)

            y = y.to(device)

            preds = model(x)

            preds = (preds > 0.5).float()

            num_correct += (preds == y).sum()

            num_pixels += torch.numel(preds)

            dice_score += (2 * (preds * y).sum()) / ((preds + y).sum() + 1e-9)

            print('Got {}/{} with acc {:2f}'.format(num_correct, num_pixels, num_correct / num_pixels * 100))

            print('Dice score {}'.format(dice_score / len(loader)))

            # wandb.log({"dice": dice_score})

            # wandb.log({"acc": (num_correct, num_pixels, num_correct / num_pixels * 100)})

            model.train()

def save_predictions_as_imgs(loader, model, device, folder='tmp/'):

    model.eval()

    for idx, (x, y) in enumerate(loader):

        x = x.to(device)

        with torch.no_grad():

            preds = model(x)

            preds = (preds > 0.5).float()

        torchvision.utils.save_image(preds, os.path.join(folder, 'pred_{}.png'.format(idx)))

        # torchvision.utils.save_image(y, os.path.join(folder, '{}.png'.format(idx)))

    model.train()