import argparse
import logging
import os
import random
import sys
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
from pathlib import Path
from torch import optim
from torch.utils.data import DataLoader, random_split
from tqdm import tqdm

from evaluate import evaluate
from unet.unet_model import UNet
from utils.data_loading import BasicDataset, CarvanaDataset
from utils.dice_score import dice_loss

device = torch.device('cuda:4' if torch.cuda.is_available() else 'cpu')

PRED_MODEL = './epoch_26_acc_0.90_best_val_acc.pth'

dir_img = Path('./data/test/imgs/')
dir_mask = Path('./data/test/masks/')
#dir_checkpoint = Path('./out_checkpoints/')

def test_model(
        model, device, 
        epochs: int = 1,
        batch_size: int = 1,
        learning_rate: float=0.001,
        img_scale: float = 0.5,
        amp: bool = False,
        weight_decay: float = 1e-8,
    ):
    
    data_transform = transforms.Compose([
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    try:
        dataset = CarvanaDataset(dir_img, dir_mask, img_scale)
    except (AssertionError, RuntimeError, IndexError):
        dataset = BasicDataset(dir_img, dir_mask, img_scale, data_transform)

    loader_args = dict(batch_size=batch_size, num_workers=os.cpu_count(), pin_memory=True)
    test_loader = DataLoader(dataset, shuffle=True, **loader_args)
    optimizer = optim.Adam(model.parameters(),lr=learning_rate, weight_decay=weight_decay)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'max', patience=5)

    test_score = evaluate(model, test_loader, device, amp)
    scheduler.step(test_score)

    logging.info('Test Dice score: {}'.format(test_score))


def get_args():
    parser = argparse.ArgumentParser(description='Train the UNet on images and target masks')
    parser.add_argument('--model', '-m', default= PRED_MODEL, metavar='FILE',help="Specify the file in which the model is stored")
    parser.add_argument('--input', '-i', metavar='INPUT', nargs='+', help='Filenames of input images', default = dir_img)
    parser.add_argument('--output', '-o', metavar='OUTPUT', nargs='+', help='Filenames of output images')
    parser.add_argument('--viz', '-v', action='store_true',
                        help='Visualize the images as they are processed')
    parser.add_argument('--no-save', '-n', action='store_true', help='Do not save the output masks')
    parser.add_argument('--mask-threshold', '-t', type=float, default=0.5,
                        help='Minimum probability value to consider a mask pixel white')
    parser.add_argument('--scale', '-s', type=float, default=0.5,
                        help='Scale factor for the input images')
    parser.add_argument('--bilinear', action='store_true', default=False, help='Use bilinear upsampling')
    parser.add_argument('--classes', '-c', type=int, default=2, help='Number of classes')
    parser.add_argument('--amp', action='store_true', default=False, help='Use mixed precision')

    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()

    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
    #device = torch.device('cuda:7' if torch.cuda.is_available() else 'cpu')
    logging.info(f'Using device {device}')

    """
    Change here to adapt to your data
    n_channels=3 for RGB images
    n_classes is the number of probabilities you want to get per pixel
    """
    model = UNet(n_channels=1, n_classes=5, bilinear=True)

    #Load pre-trained model
    model.load_state_dict(torch.load(PRED_MODEL, map_location=device))

    model = model.to(memory_format=torch.channels_last)

    logging.info(f'Network:\n'
                 f'\t{model.n_channels} input channels\n'
                 f'\t{model.n_classes} output channels (classes)\n'
                 f'\t{"Bilinear" if model.bilinear else "Transposed conv"} upscaling')

    # if args.load:
    #     state_dict = torch.load(args.load, map_location=device)
    #     del state_dict['mask_values']
    #     model.load_state_dict(state_dict)
    #     logging.info(f'Model loaded from {args.load}')

    model.to(device=device)

    test_model(
        model=model,
        device=device,
        img_scale=args.scale,
        amp=args.amp
    )