import argparse
import datetime
import logging
import os
import shutil
import time

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 adabound
import torch.utils.data
from torch.utils.data.sampler import SubsetRandomSampler
import torchvision.models as models
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torchvision.utils import save_image
from PIL.ImageOps import grayscale
from PIL import Image
from get_mean_std import get_params

import random
from tensorboard_logger import configure, log_value
from model import mobilenetv3, EfficientNet
from focal_loss import FocalLoss
from utils import get_args_from_yaml, accuracy, precision, recall, AverageMeter, printlog, FastDataLoader
import threading

# make Logger
logger = logging.getLogger('Techwing_log_file')
logger.setLevel(logging.INFO)

# make Logger stream (콘솔창에 띄우고 싶으면 해당 주석처리 제거)
#streamHandler = logging.StreamHandler()
#logger.addHandler(streamHandler)

# used for logging to TensorBoard
best_prec1 = 0
print_dataset_statestics = False

# Set Ratio test & train set
validation_ratio = 0.1
random_seed = 10

#confidence
seed = random.randint(1,1000)

# UI 단에서 실행하는 Train 방식.
def UI_train(mode, q, **kwargs):
    try:
        if mode == 'Error':
            args = get_args_from_yaml("configs/Error_config.yml")
        elif mode == 'Type':
            args = get_args_from_yaml('configs/ErrorType_config.yml')
        else:
            args = get_args_from_yaml('configs/All_config.yml')

        args['model']['blocks'] = kwargs['blocknum']
        args['data']['train'] = kwargs["data_path"]
        args['train']['epochs'] = kwargs["epoch"]
        args['optimizer']['type'] = kwargs["optim"]
        args['optimizer']['lr'] = kwargs["lr"]
        args['train']["batch-size"] = kwargs["batch_size"]
        args['predict']['batch-size'] = kwargs["batch_size"]
        args['train']['size'] = kwargs["size"]

        if kwargs['resume']:
            q.put(f"resume training with checkpoint : {kwargs['ck_path']}")
            args['train']['resume'] = kwargs['ck_path']

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

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

        if q == None:
            streamHandler = logging.StreamHandler()
            logger.addHandler(streamHandler)

        # 따로 적는 Logfile 설정 및 stream 설정.
        fileHandler = logging.FileHandler(logdir + f"{args['id']}_{args['modelname']}_block_{args['model']['blocks']}.log")
        logger.addHandler(fileHandler)

        timestring = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
        logger.info(timestring)
        
        # Train & Validate
        run_model(args, q)

        timestring = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
        logger.info(timestring)
        logging.info(f"[{args['id']}] done")

    except Exception:
        logger.error("Fatal error in main loop", exc_info=True)
        logger.warn(f"[{args['id']}] failed")

# args : Yaml에서 가져온 설정 파일, q : Dialog 에 입력하기 위하여 쓰인 Queue
def run_model(args, q=None):
    global best_prec1
    
    #실험 정보 기입. (task= All, id= 랜덤 자연수, model= mobilenetv3, epoch= 1~3000, block= 1~11, class= 2~8)
    args['task'] = "%s/%s_model=%s-ep=%s-block=%s-class=%s" % (
        args['task'],
        args['id'],
        args['modelname'],
        args['train']['epochs'],
        args['model']['blocks'],
        args['model']['class'])

    logger.info(f"use seed {seed}")
    logger.info(f"use dataset : {args['data']['train']}")

    # get GPU information from configs file
    logger.info(args)
    gpus = args['gpu']
    resize_size = args['train']['size']

    ############################# 데이터 불러오는 과정 #############################
    # Data loading code
    #feature 노말라이즈 적용.
    mean, std = get_params(args['data']['train'], resize_size)
    normalize = transforms.Normalize(mean=[mean[0].item()],
                         std=[std[0].item()])

    #Train data loader에 적용하는 함수 (순서대로 적용됨.)
    if args['train']['augment']:
        transform_train = transforms.Compose([
            transforms.Resize((resize_size, resize_size)),      # 가로세로 크기 조정
            transforms.ColorJitter(0.2,0.2,0.2),                # 밝기, 대비, 채도 조정
            transforms.RandomRotation(2),                       # -2~ 2도 만큼 회전
            transforms.RandomAffine(5),                         # affine 변환 (평행사변형이 된다든지, 사다리꼴이 된다든지)
            transforms.RandomCrop(resize_size, padding=2),      # 원본에서 padding을 상하좌우 2로 둔 뒤, 64만큼 자름
            transforms.RandomHorizontalFlip(),                  # Data 변환 좌우 반전
            transforms.Grayscale(),
            transforms.ToTensor(),
            normalize
        ])
    # option not augment
    else:
        transform_train = transforms.Compose([
            transforms.Resize((resize_size, resize_size)),
            transforms.ToTensor(),
            normalize
        ])

    # Test loader에 적용하는 함수.
    transform_test = transforms.Compose([
        transforms.Resize((resize_size, resize_size)),
        transforms.Grayscale(),
        transforms.ToTensor(),
        normalize
    ])

    # num workers : GPU의 스레드개수, Pin memory는 GPU에 데이터를 올리는 설정.
    kwargs = {'num_workers': args['train']['worker'], 'pin_memory': True}

    # image Folder config에서 설정한 path로 해주면 해당 디렉토리 안의 디렉토리의 이름이 Class가 된다.
    train_data = datasets.ImageFolder(args['data']['train'], transform_train)
    val_data = datasets.ImageFolder(args['data']['train'],transform_test)

    # Train data를 전체 데이터로 설정하였기 때문에 validation을 진행하기 위해서
    # validation ratio = 0.1 즉 10%를 validation set으로 설정.
    random_seed = 10
    validation_ratio = 0.1
    num_train = len(train_data)
    indices = list(range(num_train))
    split = int(np.floor(validation_ratio * num_train))

    # 랜덤 시드 설정. (Train이나 ,Test 일때 모두 10 이므로 같은 데이터셋이라 할 수 있다)
    np.random.seed(random_seed)
    np.random.shuffle(indices)

    # Train set, Validation set 나누기.
    train_idx, valid_idx = indices[split:], indices[:split]
    train_sampler = SubsetRandomSampler(train_idx)
    valid_sampler = SubsetRandomSampler(valid_idx)
    train_loader = FastDataLoader(
        train_data, batch_size=args['train']['batch-size'], sampler=train_sampler,      #shuffle = True
        **kwargs)
    val_loader = FastDataLoader(
        val_data, batch_size=args['train']['batch-size'], sampler=valid_sampler,        #shuffle = False
        **kwargs)

    ############################## 모델 설정 과정 #############################
    # Convolution 초기화 할 때 Random이 사용되는데 그때 사용되는 Seed 설정.
    torch.manual_seed(seed)

    # 각 클래스마다 weight를 주어서 차등적으로 학습
    class_weights = torch.FloatTensor(args['train']['weight'])
    
    # Loss 함수 설정. Cross entropy 사용.
    criterion = nn.CrossEntropyLoss(weight=class_weights)
    
    # 모델로는 Mobilenet V3 사용
    model = mobilenetv3(n_class=args['model']['class'], blocknum=args['model']['blocks'], dropout=0.5)
  
    # SGD를 사용. (Adam이 성능이 우월하지만 Tuning을 잘한 SGD가 Adam보다 성능이 좋을 때가 많기 때문에 SGD 사용.)
    if args['optimizer']['type'] == "SGD":
        optimizer = torch.optim.SGD(model.parameters(), args['optimizer']['lr'],
                                momentum=args['optimizer']['momentum'],
                                nesterov=True,
                                weight_decay=args['optimizer']['weight_decay'])
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, len(train_loader))
        
    else:
        optimizer = torch.optim.Adam(model.parameters(), args['optimizer']['lr'],
                                  weight_decay=args['optimizer']['weight_decay'])
        scheduler = None            #Adam을 사용하면 Optimizer에서 LR을 줄여주므로 스케쥴러 사용하지 않음.

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

    # for training on multiple GPUs.
    # Use CUDA_VISIBLE_DEVICES=0,1 to specify which GPUs to use
    # 멀티 GPU 설정.
    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])        # 모델을 다른 GPU에 뿌려준 다음 Gradient를 한 군데에서 계산하기 때문에 보통 0번 GPU에 많은 메로리가 할당됨.
                                                                                            # 하나의 GPU에 많은 메모리가 할당되면 batchsize를 늘릴 수 없기 때문에 이를 해결하기 위하여 output_device를 할당.
                                                                                            # 해당 코드는 데이터의 크기가 작기 때문에 0번에다가 모두 처리하는 것으로 설정.
    else:
        model = torch.nn.DataParallel(model)
        device = torch.device("cpu")
        model.to(device)
        criterion.to(device)

    # 유저가 입력한 Checkpoint에서 다시 retraining 한다고 설정을 하였을 때.
    if args['train']['resume']:
        # 해당 경로에 실제로 파일이 있으면.
        if os.path.isfile(args['train']['resume']):
            logger.info(f"=> loading checkpoint '{args['train']['resume']}'")

            if torch.cuda.is_available():
                checkpoint = torch.load(args['train']['resume'])
            else:
                checkpoint = torch.load(args['train']['resume'], map_location=torch.device('cpu'))
            
            # 시작 Epoch또한 checkpoint의 epoch로 설정.
            args['train']['start-epoch'] = checkpoint['epoch']
            
            # checkpoint에서 나온 Best accuracy를 가져옴.
            best_prec1 = checkpoint['best_prec1']

            # 앞서 선언한 모델에 weight들을 설정.
            model.load_state_dict(checkpoint['state_dict'])
            logger.info(f"=> loaded checkpoint '{args['train']['resume']}' (epoch {checkpoint['epoch']})")
        # 파일이 없다면
        else:
            logger.info(f"=> no checkpoint found at '{args.resume}'")

    # True를 설정하면 Cudnn 라이브러리에서 hardware에 따라 사용하는 내부 알고리즘을 바꾸어 준다.
    # Tensor의 크기나 Gpu Memory에 따라 효율적인 convolution 알고리즘이 다르기 때문.
    cudnn.benchmark = True
    
    for epoch in range(args['train']['start-epoch'], args['train']['epochs']):
        # train for one epoch
        train(train_loader, model, criterion, optimizer, scheduler, epoch, args, q)

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

        # remember best prec@1 and save checkpoint
        is_best = prec1 >= best_prec1

        best_prec1 = max(prec1, best_prec1)
        checkpoint = save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'best_prec1': best_prec1,
        }, is_best, args, args['model']['blocks'], q)

    printlog(f'Best accuracy: {best_prec1}', logger, q)

    if args['model']['class'] !=2:
        for i in range(len(prec)):
            logger.info(' * Precision {prec.avg:.3f}'.format(prec=prec[i]))
            logger.info(' * recall {rec.avg:.3f}'.format(rec=rec[i]))
    else:
        logger.info(' * Precision {prec.avg:.3f}'.format(prec=prec))
        logger.info(' * recall {rec.avg:.3f}'.format(rec=rec))

    #count = count + 1 
    return checkpoint

# train_loader : 데이터 로더 
# Model : MobilenetV3(default).
# criterion : Crossentropy (default).
# scheduler : CosineAnnealing.
# epoch : 3000.
# args : config parameters. yaml 파일에서 확인하실 수 있습니다.
# q : UI창에서 글을 쓰는 역활을 하는 Queue. (쓰레드에서 Queue에서 지속적으로 가지고와서 입력.)
def train(train_loader, model, criterion, optimizer, scheduler, epoch, args, q=None):
    """Train for one epoch on the training set"""
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    prec = AverageMeter()

    # switch to train mode
    model.train()
    end = time.time()
    for i, (input, target) in enumerate(train_loader):
        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, target, topk=(1,))[0]

        losses.update(loss.item(), input.size(0))
        top1.update(prec1.item(), input.size(0))
        
        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if scheduler != None:
            scheduler.step()

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

        if i % args['etc']['print_freq'] == 0:
            # Error / Normal Case를 분류하는 Task의 결과 string.
            if args['model']['class'] == 2:
                logger.info('Epoch: [{0}][{1}/{2}]\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})\t'
                            'Precision {prec.val:.3f} ({prec.avg:.3f})'
                            .format(
                                epoch, i, len(train_loader), batch_time=batch_time,
                                loss=losses, top1=top1, prec=prec))
            # All task나 ErrorType Task를 실행하였을 때의 결과 string. (6개 이상의 Class의 Prcision을 표현하기에는 문제가 있을것 같아서 
            # precision 부분은 맨 마지막에 호출하는 것으로 구성.)
            else:
                logger.info('Epoch: [{0}][{1}/{2}]\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})\t'
                            .format(
                                epoch, i, len(train_loader), batch_time=batch_time,
                                loss=losses, top1=top1))
                                
    # log to TensorBoard
    if args['etc']['tensorboard']:
        log_value('train_loss', losses.avg, epoch)
        log_value('train_acc', top1.avg, epoch)

def validate(val_loader, model, criterion, epoch, args, q=None):
    """Perform validaadd_model_to_queuetion on the validation set"""
    with torch.no_grad():
        batch_time = AverageMeter()
        losses = AverageMeter()
        top1 = AverageMeter()
        prec = []
        rec = []
        if args['model']['class'] == 2:
            prec = AverageMeter()
            rec = AverageMeter()
        else:
            for i in range(args['model']['class']):
                prec.append(AverageMeter())
                rec.append(AverageMeter())
        # switch to evaluate mode
        model.eval()
        end = time.time()

        for i, (input, target) in enumerate(val_loader):
            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]

            losses.update(loss.item(), input.size(0))
            top1.update(prec1.item(), input.size(0))
            if args['model']['class'] == 2:
                prec.update(precision(output.data, target, target_class=0), input.size(0))
                rec.update(recall(output.data, target, target_class=0), input.size(0))
            else:
                for k in range(args['model']['class']):
                    prec[k].update(precision(output.data, target, target_class=k), input.size(0))
                    rec[k].update(recall(output.data, target, target_class=k), input.size(0))

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

            if i % args['etc']['print_freq'] == 0:
                if args['model']['class'] == 2:
                    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})\t'
                                'Precision {prec.val:.3f} ({prec.avg:.3f})'
                                .format(
                                    i, len(val_loader), batch_time=batch_time, loss=losses,
                                    top1=top1, prec=prec))
                else:
                    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(' * epoch: {epoch} Prec@1 {top1.avg:.3f}'.format(epoch=epoch,top1=top1), logger, q)

        if args['model']['class'] == 2:
            logger.info(' * Precision {prec.avg:.3f}'.format(prec=prec))
            logger.info(' * recall {rec.avg:.3f}'.format(rec=rec))
        
        # log to TensorBoard
        if args['etc']['tensorboard']:
            log_value('val_loss', losses.avg, epoch)
            log_value('val_acc', top1.avg, epoch)
    return top1.avg, prec, rec


def save_checkpoint(state, is_best, args, block, q, filename='checkpoint.pth.tar'):
    """Saves checkpoint to disk"""
    directory = "%s/%s/" % (args['output'], args['task'])
    if not os.path.exists(directory):
        os.makedirs(directory)
    filename = directory + filename
    torch.save(state, filename)
    logger.info(f"Checkpoint Saved: {filename}")
    best_filename = f"{args['output']}/{args['task']}/model_best.pth.tar"
    if is_best:
        shutil.copyfile(filename, best_filename)
        logger.info(f"New Best Checkpoint saved: {best_filename}")

    return best_filename



def save_error_case(output, target,epoch, topk=(1,), input=None):
    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)

    for idx, correct_element in enumerate(correct):
        image = input[idx]
        save_image(image, f"error_case/epoch_{epoch}_idx_{idx}_case_{pred[idx]}_{target[idx]}.bmp")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True, 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 == 'ErrorType':
        args = get_args_from_yaml('configs/ErrorType_config.yml')
    else:
        args = get_args_from_yaml('configs/All_config.yml')
        
    
    #job id
    args['id'] = str(random.randint(0,99999))
    main(args)