import argparse
import csv
import logging
import os
import shutil
import time
import sys
import zipfile
from torch.utils.data.sampler import SubsetRandomSampler
import numpy as np
import PIL
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.nn.functional as F

import cv2
import matplotlib.pyplot as plt
import pandas as pd
from get_mean_std import get_params

sys.path.append(os.path.join(os.path.dirname(__name__)))

from model import mobilenetv3, EfficientNet
from torchvision.utils import save_image
from focal_loss import FocalLoss
from visualize.grad_cam import make_grad_cam
from utils import accuracy, AverageMeter, get_args_from_yaml, MyImageFolder, printlog, FastDataLoader
from PIL import Image
import torchvision.transforms.functional as TF
#from utils import restapi, preprocessing

global error_case_idx, correct_case_idx

logger = logging.getLogger(os.path.dirname(__name__))
logger.setLevel(logging.INFO)

streamHandler = logging.StreamHandler()
logger.addHandler(streamHandler)

def make_type_dir():
    if not os.path.exists('test_result'):
        os.mkdir('test_result')
    if not os.path.exists('test_result/Type'):
        os.mkdir('test_result/Type')  
    if not os.path.exists('test_result/Type/Double'):
        os.mkdir('test_result/Type/Double')
    if not os.path.exists('test_result/Type/Flip'):
        os.mkdir('test_result/Type/Flip')
    if not os.path.exists('test_result/Type/Scratch'):
        os.mkdir('test_result/Type/Scratch')
    if not os.path.exists('test_result/Type/Leave'):
        os.mkdir('test_result/Type/Leave')
    if not os.path.exists('test_result/Type/Empty'):
        os.mkdir('test_result/Type/Empty')
    if not os.path.exists('test_result/Type/Crack'):
        os.mkdir('test_result/Type/Crack')

def make_all_dir():
    if not os.path.exists('test_result'):
        os.mkdir('test_result')
    if not os.path.exists('test_result/All'):
        os.mkdir('test_result/All')  
    if not os.path.exists('test_result/All/Double'):
        os.mkdir('test_result/All/Double')
    if not os.path.exists('test_result/All/Flip'):
        os.mkdir('test_result/All/Flip')
    if not os.path.exists('test_result/All/Scratch'):
        os.mkdir('test_result/All/Scratch')
    if not os.path.exists('test_result/All/Leave'):
        os.mkdir('test_result/All/Leave')
    if not os.path.exists('test_result/All/Normal'):
        os.mkdir('test_result/All/Normal')
    if not os.path.exists('test_result/All/Empty'):
        os.mkdir('test_result/All/Empty')
    if not os.path.exists('test_result/All/Crack'):
        os.mkdir('test_result/All/Crack')
    if not os.path.exists('test_result/All/Normal'):
        os.mkdir('test_result/All/Normal')

def make_error_dir():
    if not os.path.exists('test_result'):
        os.mkdir('test_result')
    if not os.path.exists('test_result/Error'):
        os.mkdir('test_result/Error')
    if not os.path.exists('test_result/Error/Normal'):
        os.mkdir('test_result/Error/Normal')
    if not os.path.exists('test_result/Error/Error'):
        os.mkdir('test_result/Error/Error')

def get_savepath_classes_args(mode):
    if mode == "Error":
        save_path = './test_result/Error'
        classes = ['Error', 'Normal']
        args = get_args_from_yaml("configs/Error_config.yml")

    elif mode == "Type":
        save_path = './test_result/Type'
        classes = ['Crack', 'Double', 'Empty', 'Flip', 'Leave','Pollute', 'Scratch']
        args = get_args_from_yaml('configs/ErrorType_config.yml')

    else:
        save_path = './test_result/All'
        classes = ['Crack','Double', 'Empty', 'Flip', 'Leave', 'Normal','Pollute', 'Scratch']
        args = get_args_from_yaml('configs/All_config.yml')

    return save_path, classes, args
    

# 여러개의 인풋을 Test 수행할 때 사용되는 함수.
def test(testloader, model, mode):
     with torch.no_grad():
        save_path, classes, _ = get_savepath_classes_args(mode)
        model.eval()
        for _, data in enumerate(testloader):
            (input, _), (path, _) = data
            if torch.cuda.is_available():
                input = input.cuda()

            output = model(input)
            prob = F.softmax(output, dim=1)

            for idx, p in enumerate(prob):
                values = torch.topk(p,2).values.tolist()
                indices = torch.topk(p,2).indices.tolist()
                img = cv2.imread(path[idx])
                cv2.imwrite(f"{save_path}/{classes[indices[0]]}/{classes[indices[0]]}={values[0]}__{classes[indices[1]]}={values[1]}.bmp", img)

# Test input이 하나의 파일일 때 사용되는 함수.
# Path = 데이터 원본의 경로, mode = 수행하는 Task.
def single_file_test(input, model, path, mode, q):
    with torch.no_grad():
        save_path, classes, _ = get_savepath_classes_args(mode)
        model.eval()
        if torch.cuda.is_available():
            input = input.cuda()

        start = time.time()    
        output = model(input)
        prob = F.softmax(output, dim=1)
        q.put(f"Inference time 1 image : {str(round(time.time() - start , 5))}")

        for idx, p in enumerate(prob):
            values = torch.topk(p,2).values.tolist()        # 확률
            indices = torch.topk(p,2).indices.tolist()      # 인덱스
            img = cv2.imread(path)
            cv2.imwrite(f"{save_path}/{classes[indices[0]]}/{classes[indices[0]]}={values[0]}__{classes[indices[1]]}={values[1]}.bmp", img)
                
# 유저가 지정해준 checkpoint가 없으면 config 에 있는 checkpoint를 사용.
# data는 config에 지정된 data를 활용.
def UI_validate(mode, q, **kwargs):
    try:
        _, _, args = get_savepath_classes_args(mode)
        args['model']['blocks'] = kwargs['blocknum']
        args['data']['val'] = kwargs["data_path"]

        q.put(f"using user's checkpoint {kwargs['ck_path']}")
        args['checkpoint'] = kwargs['ck_path']

        timestring = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
        args['id'] = "validate_" + timestring
        main(args, q)
    except Exception as ex:
        q.put(f"실행 중 에러가 발생하였습니다. 자세한 사항은 보시려면 로그를 확인해 주세요.")
        logger.info(ex)

# test는 항상 유저가 지정해주는 data를 활용.
# 만약 없다면 demo 즉 UI단에서 지정한 default data를 활용.
# mode: Error, ErrorType, All
# path: data path
# test_mode: File or dir
def UI_test(mode, path, test_mode, q, **kwargs):
    try:
        _, _, args = get_savepath_classes_args(mode)
        make_error_dir()
        make_type_dir()
        make_all_dir()

        args['model']['blocks'] = kwargs['blocknum']
        args['train']['size'] = kwargs['size']

        q.put(f"using user's checkpoint {kwargs['ck_path']}")
        args['checkpoint'] = kwargs['ck_path']

        timestring = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
        args['id'] = "test_" + timestring
        gpus = args['gpu']
        resize_size = args['train']['size']
        model = mobilenetv3(n_class=args['model']['class'], blocknum=args['model']['blocks'])
        if torch.cuda.is_available():
            torch.cuda.set_device(gpus[0])
            with torch.cuda.device(gpus[0]):
                model = model.cuda()
            model = torch.nn.DataParallel(model, device_ids=gpus, output_device=gpus[0])
        else:
            model = torch.nn.DataParallel(model)
            device = torch.device("cpu")
            model.to(device)

        q.put("loading checkpoint...")
        if torch.cuda.is_available():
            checkpoint = torch.load(args['checkpoint'])
        else:
            checkpoint = torch.load(args['checkpoint'],map_location=torch.device('cpu'))

        model.load_state_dict(checkpoint['state_dict'])
        q.put("checkpoint already loaded!")
        q.put("start test")

        # 테스트 데이터가 디렉토일 경우.
        # 다중 데이터를 받아야 하므로 Pytorch에서 구성된 Dataloader 이용.
        # 해당 코드에서는 Data의 Path까지 출력해주는 Loader을 추가로 구성함. (저장하기 위하여)
        if test_mode == 'dir':
            normalize = transforms.Normalize(mean=[0.4015], std=[0.2165])
            transform_test = transforms.Compose([
                transforms.Resize((resize_size,resize_size)),
                transforms.Grayscale(),
                transforms.ToTensor(),
                normalize
            ])
            q.put(f"data path directory is {path}")
            testset = MyImageFolder(path, transform=transform_test)
            test_loader = FastDataLoader(testset, batch_size=args['predict']['batch-size'], shuffle=False, num_workers=8)
            start = time.time()
            test(test_loader, model, mode)
            q.put(f"Inference time {len(testset)} images : {str(round(time.time() - start , 5))}")
        
        # 테스트 데이터가 하나의 파일일 경우.
        # 하나의 데이터이므로 바로 이미지를 텐서로 바꿈.
        # Dataloader에서는 transforms.compose로 데이터 Preprocessing을 묶었지만
        # 여기서는 Dataloader를 사용하지 않기 때문에 transforms.functional을 이용하여 직접 변경.
        # trasnforms 함수와 같은 형태를 가지고 있기 때문에 쉽게 이해 가능.
        else:
            image = Image.open(path)
            x = TF.resize(image, (resize_size,resize_size))     # 리사이즈
            x = TF.to_grayscale(x)                              # 그레이스케일 적용
            x = TF.to_tensor(x)                                 # 텐서 변환.
            x.unsqueeze_(0)                                     # 0-dim에 차원 추가.
            start = time.time()
            single_file_test(x, model, path, mode, q)
            q.put(f"Inference time 1 image : {str(round(time.time() - start , 5))}")
        q.put('finish test')

    except Exception as ex:
        q.put("실행 중 에러가 발생하였습니다. 자세한 사항은 보시려면 로그를 확인해 주세요.")
        logger.info(ex)

def UI_temp(path,q,model):
    try:
        resize_size = 64
        image = Image.open(path)
        x = TF.resize(image, (resize_size, resize_size))
        x = TF.to_grayscale(x)
        x = TF.to_tensor(x)
        x.unsqueeze_(0)
        
        single_file_test(x, model, path, "Error", q)
        
        q.put('temp test finish')
    except Exception as ex:
        q.put("실행 중 에러가 발생하였습니다. 자세한 사항은 보시려면 로그를 확인해 주세요.")
        logger.info(ex)

def UI_temp2():
    batch_size = 256

    train_transforms = transforms.Compose([
        transforms.Resize((256,256)),
        transforms.ToTensor()
    ])

    train_dataset = datasets.ImageFolder("../data/Fifth_data/All", train_transforms)

    train_loader = FastDataLoader(dataset=train_dataset,
    batch_size=batch_size, shuffle=True)
    i=0
    start = time.time()
    for x,y in train_loader:
        i = i+1
        pass
    end = time.time()
    print((end - start)/i)

def main(args, q=None):
    try:
        logdir = f'logs/runs/{args["task"]}/'
        
        if not os.path.exists(logdir):
            os.makedirs(logdir)

        fileHandler = logging.FileHandler(logdir + f'{args["id"]}.log')
        logger.addHandler(fileHandler)

        # 2. eval
        run_model(args, q)

        # 3. Done
        printlog(f"[{args['id']}] done", logger, q)

    except Exception as ex:
        printlog("실행 중 에러가 발생하였습니다. 자세한 사항은 보시려면 로그를 확인해 주세요.", logger, q)
        logger.info(ex)

def run_model(args, q=None):
    resize_size = args['train']['size']

    gpus = args['gpu']

    mean, std = get_params(args['data']['val'], resize_size)

    normalize = transforms.Normalize(mean=[mean[0].item()],
                         std=[std[0].item()])
    normalize_factor = [mean, std]

    # data loader
    transform_test = transforms.Compose([
        transforms.Resize((resize_size,resize_size)),
        transforms.Grayscale(),
        transforms.ToTensor(),
        normalize
    ])
    kwargs = {'num_workers': args['predict']['worker'], 'pin_memory': True}
    test_data = MyImageFolder(args['data']['val'], transform_test)

    random_seed = 10
    validation_ratio = 0.1
    num_test = len(test_data)
    indices = list(range(num_test))
    split = int(np.floor(validation_ratio * num_test))

    np.random.seed(random_seed)
    np.random.shuffle(indices)

    valid_idx = indices[:split]
    valid_sampler = SubsetRandomSampler(valid_idx)


    val_loader = FastDataLoader(
        test_data, batch_size=args['predict']['batch-size'], sampler=valid_sampler,
        **kwargs)

    criterion = nn.CrossEntropyLoss()

    # load model
    model = mobilenetv3(n_class=args['model']['class'], blocknum=args['model']['blocks'])

    # get the number of model parameters
    logger.info('Number of model parameters: {}'.format(
        sum([p.data.nelement() for p in model.parameters()])))

    if torch.cuda.is_available():
        torch.cuda.set_device(gpus[0])
        with torch.cuda.device(gpus[0]):
            model = model.cuda()
            criterion = criterion.cuda()
        model = torch.nn.DataParallel(model, device_ids=gpus, output_device=gpus[0])
    else:
        model = torch.nn.DataParallel(model)
        device = torch.device("cpu")
        model.to(device)
        criterion.to(device)


    logger.info("=> loading checkpoint '{}'".format(args['checkpoint']))

    if torch.cuda.is_available():
        checkpoint = torch.load(args['checkpoint'])
    else:
        checkpoint = torch.load(args['checkpoint'], map_location=torch.device('cpu'))
    args['start_epoch'] = checkpoint['epoch']
    best_prec1 = checkpoint['best_prec1']
    model.load_state_dict(checkpoint['state_dict'])
    logger.info("=> loaded checkpoint '{}' (epoch {})"
          .format(args['checkpoint'], checkpoint['epoch']))

    cudnn.benchmark = True

    # define loss function (option 2)
    #criterion = FocalLoss(
    #   gamma=args['loss']['gamma'], alpha=args['loss']['alpha']).cuda()

    # evaluate on validation set
    prec1 = validate(val_loader, model, criterion, normalize_factor, args ,q)

    # remember best prec@1 and save checkpoint
    best_prec1 = max(prec1, best_prec1)
    logger.info(f'Best accuracy: {best_prec1}')


def validate(val_loader, model, criterion, normalize_factor, args, q):
    """Perform validation on the validation set"""
    with torch.no_grad():
        batch_time = AverageMeter()
        losses = AverageMeter()
        top1 = AverageMeter()

        # switch to evaluate mode
        model.eval()

        end = time.time()
        for i, data in enumerate(val_loader):
            (input, target), (path, _) = data
            if torch.cuda.is_available():
                target = target.cuda()
                input = input.cuda()

            # compute output
            output = model(input)

            loss = criterion(output, target)
            
            # measure accuracy and record loss
            prec1 = accuracy(output.data, target, topk=(1,))[0]

            #save error case
            #save correct = 맞은거 까지 저장하는지 마는지.
            if args['predict']['save']:
                save_error_case(output.data, target, path, args, topk=(1,), input=input, save_correct=False)

            losses.update(loss.item(), input.size(0))
            top1.update(prec1.item(), input.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % 1 == 0:
                logger.info('Test: [{0}/{1}]\t'
                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                          i, len(val_loader), batch_time=batch_time, loss=losses,
                          top1=top1))
            
        printlog(' * Prec@1 {top1.avg:.3f}'.format(top1=top1), logger, q)
        
    if args["predict"]["cam"]:
        logger.info("Creating CAM")

        #print grad cam
        if args['predict']['normalize']:
            make_grad_cam(f"eval_results/{args['task']}/error_case",
                    f"eval_results/{args['task']}/error_case/cam" , model, normalize_factor, cam_class=args['predict']['cam-class'], args=args)
        else:
            make_grad_cam(f"eval_results/{args['task']}/error_case",
                    f"eval_results/{args['task']}/error_case/cam" , model, normalize_factor=None, cam_class=args['predict']['cam-class'], args = args)

    return top1.avg
    
def save_error_case(output, target, path, args , topk=(1,), input=None, save_correct=False):
    global error_case_idx, correct_case_idx

    error_case_idx = 0
    correct_case_idx = 0

    _, class_arr, _ = get_savepath_classes_args(args['task'])

    p = F.softmax(output, dim=1)

    values = torch.topk(p,2).values.tolist()
    indices = torch.topk(p,2).indices.tolist()

    maxk = max(topk)
    batch_size = target.size(0)
    _, pred = output.topk(maxk, 1, True, True)

    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    pred = pred.view(batch_size)
    correct = correct.view(batch_size)

    if not os.path.exists(f'eval_results'):
        os.mkdir(f'eval_results')

    if not os.path.exists(f"eval_results/{args['task']}"):
        os.mkdir(f"eval_results/{args['task']}")

    if not os.path.exists(f"eval_results/{args['task']}/error_case"):
        os.mkdir(f"eval_results/{args['task']}/error_case")
    
    if not os.path.exists(f"eval_results/{args['task']}/correct_case") and save_correct:
        os.mkdir(f"eval_results/{args['task']}/correct_case")

    for idx, correct_element in enumerate(correct):
        # 틀린 경우
        if correct_element.item() == 0:
            #save_image(input[idx], f"eval_results/{args['task']}/error_case/idx_{error_case_idx}_label_{target[idx]}_pred_{pred[idx]}.bmp")
            img = cv2.imread(path[idx])
            cv2.imwrite(f"eval_results/{args['task']}/error_case/idx_{error_case_idx}_label_{class_arr[target[idx]]}_pred_{class_arr[indices[idx][0]]}={round(values[idx][0]*100,1)}_{class_arr[indices[idx][1]]}={round(values[idx][1]*100,1)}_real.bmp" ,img)
            error_case_idx = error_case_idx + 1

        # 맞는 경우에도 저장.
        if save_correct and correct_element.item() == 1:
            #save_image(input[idx], f"eval_results/{args['task']}/correct_case/idx_{correct_case_idx}_label_{target[idx]}.bmp")
            img = cv2.imread(path[idx])
            cv2.imwrite(f"eval_results/{args['task']}/correct_case/idx_{correct_case_idx}_label_{class_arr[target[idx]]}_real.bmp" ,img)
            correct_case_idx = correct_case_idx + 1


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="configs/Error_config.yml", help="train config file")       #config 파일을 디폴트로 받음.
    args = parser.parse_args()

    if args.config == 'Error':
        args = get_args_from_yaml("configs/Error_config.yml")
    elif args.config == 'Type':
        args = get_args_from_yaml('configs/ErrorType_config.yml')
    else:
        args = get_args_from_yaml('configs/All_config.yml')
    args['id'] = 'eval'

    main(args)