# SmartinsideAI - 2021AIChamp
import argparse
import numpy as np
import os
import pandas as pd
from tqdm import tqdm
import yaml
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SequentialSampler, RandomSampler
from torch.optim import Adam, SGD

import time
from sklearn.metrics import accuracy_score
from contextlib import contextmanager
from sklearn.model_selection import StratifiedKFold
import matplotlib.pyplot as plt
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau

# Advanced augmentations
from albumentations import (
    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,
    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop,
    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout, ShiftScaleRotate, CenterCrop, Resize
)
from albumentations.pytorch import ToTensorV2
from albumentations import ImageOnlyTransform

# Utils functions
from utils import seed_everything
import dataset
import models

import warnings
warnings.filterwarnings("ignore")

#### Parser
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--cfg", default='configs/resnext50_32x4d.yaml', type=str)
parser.add_argument("--cuda", default='cuda:0', type=str)

args = parser.parse_args()
print(args.cuda)

SEED = 123
seed_everything(SEED)

if __name__ == "__main__":
    print(os.getcwd())
    with open(args.cfg) as f:
        cfg = yaml.load(f, Loader=yaml.FullLoader)
    print(cfg)

    print(cfg['model_name'])
    OUTPUT_DIR =  'weights/' + cfg['model_name'] + '_{}_{}_{}/'.format(cfg['size'],
                                                        cfg['scheduler'],
                                                        cfg['criterion'])
    
    os.makedirs(OUTPUT_DIR, exist_ok = True)

    # ====================================================
    # Preprocessing dataframe
    # ====================================================
    train_org = pd.read_csv("/media/daitran/Data_SSD/data/SmartinsideAI/2021AIChamp/preprocessed_data/preprocessed_train_df.csv")
    train = train_org.drop(["type", "fault"], axis = 1)
    train.columns = ["num","image_id", "label"]

    # IMAGE LOCATIONS
    TRAIN_PATH = '/media/daitran/Data_SSD/data/SmartinsideAI/2021AIChamp/11.한국전력공사_데이터/train'
    TEST_PATH = '/media/daitran/Data_SSD/data/SmartinsideAI/2021AIChamp/11.한국전력공사_데이터/test'

    # Set cuda 1 for Lab server
    device = torch.device("cuda:0")

    # Utils function
    def init_logger(log_file=OUTPUT_DIR+'train.log'):
        from logging import getLogger, INFO, FileHandler,  Formatter,  StreamHandler
        logger = getLogger(__name__)
        logger.setLevel(INFO)
        handler1 = StreamHandler()
        handler1.setFormatter(Formatter("%(message)s"))
        handler2 = FileHandler(filename=log_file)
        handler2.setFormatter(Formatter("%(message)s"))
        logger.addHandler(handler1)
        logger.addHandler(handler2)
        return logger
    
    def get_score(y_true, y_pred):
        return accuracy_score(y_true, y_pred)

    @contextmanager
    def timer(name):
        t0 = time.time()
        LOGGER.info(f'[{name}] start')
        yield
        LOGGER.info(f'[{name}] done in {time.time() - t0:.0f} s.')
        
    # DEFINE LOGGER
    LOGGER = init_logger()
    
    # KFOLD SPLIT
    folds = train.copy()
    Fold = StratifiedKFold(n_splits=cfg['n_fold'], shuffle=True, random_state=cfg['seed'])
    for n, (train_index, val_index) in enumerate(Fold.split(folds, folds[cfg['target_col']])):
        folds.loc[val_index, 'fold'] = int(n)
    folds['fold'] = folds['fold'].astype(int)
    
    # CHECK DATASET
#     train_dataset = dataset.TrainDataset(train, transform=None)

#     for i in range(1):
#         image, label = train_dataset[i]
#         plt.imshow(image)
#         plt.title(f'label: {label}')
#         plt.show() 

    # Augment

    def get_transforms(*, data):
        
        if data == 'train':
            return Compose([
                Resize(cfg['size'], cfg['size']),
                CLAHE(),
                # RandomResizedCrop(cfg['size'], cfg['size']),
                # Transpose(p=0.5),
                HorizontalFlip(p=0.5),
                VerticalFlip(p=0.5),
                # ShiftScaleRotate(p=0.5),
                Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),

                ToTensorV2(),
            ])

        elif data == 'valid':
            return Compose([
                Resize(cfg['size'], cfg['size']),
                Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                ),
                ToTensorV2(),
            ])

    # ADD LOSS FUNCTIONS
    # ====================================================
    # Label Smoothing
    # ====================================================
    class LabelSmoothingLoss(nn.Module): 
        def __init__(self, classes=cfg['target_size'], smoothing=0.0, dim=-1): 
            super(LabelSmoothingLoss, self).__init__() 
            self.confidence = 1.0 - smoothing 
            self.smoothing = smoothing 
            self.cls = classes 
            self.dim = dim 
        def forward(self, pred, target): 
            pred = pred.log_softmax(dim=self.dim) 
            with torch.no_grad():
                true_dist = torch.zeros_like(pred) 
                true_dist.fill_(self.smoothing / (self.cls - 1)) 
                true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence) 
            return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))
    # ====================================================
    # Focal loss
    # ====================================================
    class FocalLoss(nn.Module):
        def __init__(self, alpha=1, gamma=2, reduce=True):
            super(FocalLoss, self).__init__()
            self.alpha = alpha
            self.gamma = gamma
            self.reduce = reduce

        def forward(self, inputs, targets):
            BCE_loss = nn.CrossEntropyLoss()(inputs, targets)

            pt = torch.exp(-BCE_loss)
            F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss

            if self.reduce:
                return torch.mean(F_loss)
            else:
                return F_loss
    
    # ====================================================
    # FocalCosine loss
    # ====================================================
    class FocalCosineLoss(nn.Module):
        def __init__(self, alpha=1, gamma=2, xent=.1):
            super(FocalCosineLoss, self).__init__()
            self.alpha = alpha
            self.gamma = gamma

            self.xent = xent

            self.y = torch.Tensor([1]).cuda(device)

        def forward(self, input, target, reduction="mean"):
            cosine_loss = F.cosine_embedding_loss(input, F.one_hot(target, num_classes=input.size(-1)), self.y, reduction=reduction)

            cent_loss = F.cross_entropy(F.normalize(input), target, reduce=False)
            pt = torch.exp(-cent_loss)
            focal_loss = self.alpha * (1-pt)**self.gamma * cent_loss

            if reduction == "mean":
                focal_loss = torch.mean(focal_loss)

            return cosine_loss + self.xent * focal_loss
    # ====================================================
    # SymmetricCrossEntropy
    # ====================================================
    class SymmetricCrossEntropy(nn.Module):

        def __init__(self, alpha=0.1, beta=1.0, num_classes=cfg['target_size']):
            super(SymmetricCrossEntropy, self).__init__()
            self.alpha = alpha
            self.beta = beta
            self.num_classes = num_classes

        def forward(self, logits, targets, reduction='mean'):
            onehot_targets = torch.eye(self.num_classes)[targets].cuda(device)
            ce_loss = F.cross_entropy(logits, targets, reduction=reduction)
            rce_loss = (-onehot_targets*logits.softmax(1).clamp(1e-7, 1.0).log()).sum(1)
            if reduction == 'mean':
                rce_loss = rce_loss.mean()
            elif reduction == 'sum':
                rce_loss = rce_loss.sum()
            return self.alpha * ce_loss + self.beta * rce_loss
    # ====================================================
    # Bi-Tempered-Loss
    # ====================================================
    
    def log_t(u, t):
        """Compute log_t for `u'."""
        if t==1.0:
            return u.log()
        else:
            return (u.pow(1.0 - t) - 1.0) / (1.0 - t)

    def exp_t(u, t):
        """Compute exp_t for `u'."""
        if t==1:
            return u.exp()
        else:
            return (1.0 + (1.0-t)*u).relu().pow(1.0 / (1.0 - t))

    def compute_normalization_fixed_point(activations, t, num_iters):

        """Returns the normalization value for each example (t > 1.0).
        Args:
          activations: A multi-dimensional tensor with last dimension `num_classes`.
          t: Temperature 2 (> 1.0 for tail heaviness).
          num_iters: Number of iterations to run the method.
        Return: A tensor of same shape as activation with the last dimension being 1.
        """
        mu, _ = torch.max(activations, -1, keepdim=True)
        normalized_activations_step_0 = activations - mu

        normalized_activations = normalized_activations_step_0

        for _ in range(num_iters):
            logt_partition = torch.sum(
                    exp_t(normalized_activations, t), -1, keepdim=True)
            normalized_activations = normalized_activations_step_0 * \
                    logt_partition.pow(1.0-t)

        logt_partition = torch.sum(
                exp_t(normalized_activations, t), -1, keepdim=True)
        normalization_constants = - log_t(1.0 / logt_partition, t) + mu

        return normalization_constants

    def compute_normalization_binary_search(activations, t, num_iters):

        """Returns the normalization value for each example (t < 1.0).
        Args:
          activations: A multi-dimensional tensor with last dimension `num_classes`.
          t: Temperature 2 (< 1.0 for finite support).
          num_iters: Number of iterations to run the method.
        Return: A tensor of same rank as activation with the last dimension being 1.
        """

        mu, _ = torch.max(activations, -1, keepdim=True)
        normalized_activations = activations - mu

        effective_dim = \
            torch.sum(
                    (normalized_activations > -1.0 / (1.0-t)).to(torch.int32),
                dim=-1, keepdim=True).to(activations.dtype)

        shape_partition = activations.shape[:-1] + (1,)
        lower = torch.zeros(shape_partition, dtype=activations.dtype, device=activations.device)
        upper = -log_t(1.0/effective_dim, t) * torch.ones_like(lower)

        for _ in range(num_iters):
            logt_partition = (upper + lower)/2.0
            sum_probs = torch.sum(
                    exp_t(normalized_activations - logt_partition, t),
                    dim=-1, keepdim=True)
            update = (sum_probs < 1.0).to(activations.dtype)
            lower = torch.reshape(
                    lower * update + (1.0-update) * logt_partition,
                    shape_partition)
            upper = torch.reshape(
                    upper * (1.0 - update) + update * logt_partition,
                    shape_partition)

        logt_partition = (upper + lower)/2.0
        return logt_partition + mu

    class ComputeNormalization(torch.autograd.Function):
        """
        Class implementing custom backward pass for compute_normalization. See compute_normalization.
        """
        @staticmethod
        def forward(ctx, activations, t, num_iters):
            if t < 1.0:
                normalization_constants = compute_normalization_binary_search(activations, t, num_iters)
            else:
                normalization_constants = compute_normalization_fixed_point(activations, t, num_iters)

            ctx.save_for_backward(activations, normalization_constants)
            ctx.t=t
            return normalization_constants

        @staticmethod
        def backward(ctx, grad_output):
            activations, normalization_constants = ctx.saved_tensors
            t = ctx.t
            normalized_activations = activations - normalization_constants 
            probabilities = exp_t(normalized_activations, t)
            escorts = probabilities.pow(t)
            escorts = escorts / escorts.sum(dim=-1, keepdim=True)
            grad_input = escorts * grad_output

            return grad_input, None, None

    def compute_normalization(activations, t, num_iters=5):
        """Returns the normalization value for each example. 
        Backward pass is implemented.
        Args:
          activations: A multi-dimensional tensor with last dimension `num_classes`.
          t: Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).
          num_iters: Number of iterations to run the method.
        Return: A tensor of same rank as activation with the last dimension being 1.
        """
        return ComputeNormalization.apply(activations, t, num_iters)

    def tempered_sigmoid(activations, t, num_iters = 5):
        """Tempered sigmoid function.
        Args:
          activations: Activations for the positive class for binary classification.
          t: Temperature tensor > 0.0.
          num_iters: Number of iterations to run the method.
        Returns:
          A probabilities tensor.
        """
        internal_activations = torch.stack([activations,
            torch.zeros_like(activations)],
            dim=-1)
        internal_probabilities = tempered_softmax(internal_activations, t, num_iters)
        return internal_probabilities[..., 0]


    def tempered_softmax(activations, t, num_iters=5):
        """Tempered softmax function.
        Args:
          activations: A multi-dimensional tensor with last dimension `num_classes`.
          t: Temperature > 1.0.
          num_iters: Number of iterations to run the method.
        Returns:
          A probabilities tensor.
        """
        if t == 1.0:
            return activations.softmax(dim=-1)

        normalization_constants = compute_normalization(activations, t, num_iters)
        return exp_t(activations - normalization_constants, t)

    def bi_tempered_binary_logistic_loss(activations,
            labels,
            t1,
            t2,
            label_smoothing = 0.0,
            num_iters=5,
            reduction='mean'):

        """Bi-Tempered binary logistic loss.
        Args:
          activations: A tensor containing activations for class 1.
          labels: A tensor with shape as activations, containing probabilities for class 1
          t1: Temperature 1 (< 1.0 for boundedness).
          t2: Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).
          label_smoothing: Label smoothing
          num_iters: Number of iterations to run the method.
        Returns:
          A loss tensor.
        """
        internal_activations = torch.stack([activations,
            torch.zeros_like(activations)],
            dim=-1)
        internal_labels = torch.stack([labels.to(activations.dtype),
            1.0 - labels.to(activations.dtype)],
            dim=-1)
        return bi_tempered_logistic_loss(internal_activations, 
                internal_labels,
                t1,
                t2,
                label_smoothing = label_smoothing,
                num_iters = num_iters,
                reduction = reduction)

    def bi_tempered_logistic_loss(activations,
            labels,
            t1,
            t2,
            label_smoothing=0.0,
            num_iters=5,
            reduction = 'mean'):

        """Bi-Tempered Logistic Loss.
        Args:
          activations: A multi-dimensional tensor with last dimension `num_classes`.
          labels: A tensor with shape and dtype as activations (onehot), 
            or a long tensor of one dimension less than activations (pytorch standard)
          t1: Temperature 1 (< 1.0 for boundedness).
          t2: Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).
          label_smoothing: Label smoothing parameter between [0, 1). Default 0.0.
          num_iters: Number of iterations to run the method. Default 5.
          reduction: ``'none'`` | ``'mean'`` | ``'sum'``. Default ``'mean'``.
            ``'none'``: No reduction is applied, return shape is shape of
            activations without the last dimension.
            ``'mean'``: Loss is averaged over minibatch. Return shape (1,)
            ``'sum'``: Loss is summed over minibatch. Return shape (1,)
        Returns:
          A loss tensor.
        """

        if len(labels.shape)<len(activations.shape): #not one-hot
            labels_onehot = torch.zeros_like(activations)
            labels_onehot.scatter_(1, labels[..., None], 1)
        else:
            labels_onehot = labels

        if label_smoothing > 0:
            num_classes = labels_onehot.shape[-1]
            labels_onehot = ( 1 - label_smoothing * num_classes / (num_classes - 1) ) \
                    * labels_onehot + \
                    label_smoothing / (num_classes - 1)

        probabilities = tempered_softmax(activations, t2, num_iters)

        loss_values = labels_onehot * log_t(labels_onehot + 1e-10, t1) \
                - labels_onehot * log_t(probabilities, t1) \
                - labels_onehot.pow(2.0 - t1) / (2.0 - t1) \
                + probabilities.pow(2.0 - t1) / (2.0 - t1)
        loss_values = loss_values.sum(dim = -1) #sum over classes

        if reduction == 'none':
            return loss_values
        if reduction == 'sum':
            return loss_values.sum()
        if reduction == 'mean':
            return loss_values.mean()
    
    class BiTemperedLogisticLoss(nn.Module): 
        def __init__(self, t1, t2, smoothing=0.0): 
            super(BiTemperedLogisticLoss, self).__init__() 
            self.t1 = t1
            self.t2 = t2
            self.smoothing = smoothing
        def forward(self, logit_label, truth_label):
            loss_label = bi_tempered_logistic_loss(
                logit_label, truth_label,
                t1=self.t1, t2=self.t2,
                label_smoothing=self.smoothing,
                reduction='none'
            )

            loss_label = loss_label.mean()
            return loss_label
    # ====================================================
    # TaylorCrossEntropyLoss
    # ====================================================
    class TaylorSoftmax(nn.Module):
        '''
        This is the autograd version
        '''
        def __init__(self, dim=1, n=2):
            super(TaylorSoftmax, self).__init__()
            assert n % 2 == 0
            self.dim = dim
            self.n = n

        def forward(self, x):
            '''
            usage similar to nn.Softmax:
                >>> mod = TaylorSoftmax(dim=1, n=4)
                >>> inten = torch.randn(1, 32, 64, 64)
                >>> out = mod(inten)
            '''
            fn = torch.ones_like(x)
            denor = 1.
            for i in range(1, self.n+1):
                denor *= i
                fn = fn + x.pow(i) / denor
            out = fn / fn.sum(dim=self.dim, keepdims=True)
            return out


    class TaylorCrossEntropyLoss(nn.Module):
        def __init__(self, n=2, ignore_index=-1, reduction='mean', smoothing=0.05):
            super(TaylorCrossEntropyLoss, self).__init__()
            assert n % 2 == 0
            self.taylor_softmax = TaylorSoftmax(dim=1, n=n)
            self.reduction = reduction
            self.ignore_index = ignore_index
            self.lab_smooth = LabelSmoothingLoss(cfg['target_size'], smoothing=smoothing)

        def forward(self, logits, labels):
            log_probs = self.taylor_softmax(logits).log()
            #loss = F.nll_loss(log_probs, labels, reduction=self.reduction,
            #        ignore_index=self.ignore_index)
            loss = self.lab_smooth(log_probs, labels)
            return loss
    
    # ====================================================
    class AverageMeter(object):
        """Computes and stores the average and current value"""
        def __init__(self):
            self.reset()

        def reset(self):
            self.val = 0
            self.avg = 0
            self.sum = 0
            self.count = 0

        def update(self, val, n=1):
            self.val = val
            self.sum += val * n
            self.count += n
            self.avg = self.sum / self.count


    def asMinutes(s):
        m = math.floor(s / 60)
        s -= m * 60
        return '%dm %ds' % (m, s)


    def timeSince(since, percent):
        now = time.time()
        s = now - since
        es = s / (percent)
        rs = es - s
        return '%s (remain %s)' % (asMinutes(s), asMinutes(rs))
    
    
    # TRAIN
    def train_fn(train_loader, model, criterion, optimizer, epoch, scheduler, device):
        batch_time = AverageMeter()
        data_time = AverageMeter()
        losses = AverageMeter()
        scores = AverageMeter()
        # switch to train mode
        model.train()
        start = end = time.time()
        global_step = 0
        for step, (images, labels) in enumerate(train_loader):
            # measure data loading time
            data_time.update(time.time() - end)
            images = images.to(device)
            labels = labels.to(device)
            batch_size = labels.size(0)
            y_preds = model(images)
            loss = criterion(y_preds, labels)
            # record loss
            losses.update(loss.item(), batch_size)
            if cfg['gradient_accumulation_steps'] > 1:
                loss = loss / cfg['gradient_accumulation_steps']
            if cfg['apex']:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg['max_grad_norm'])
            if (step + 1) % cfg['gradient_accumulation_steps'] == 0:
                optimizer.step()
                optimizer.zero_grad()
                global_step += 1
            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()
            if step % cfg['print_freq'] == 0 or step == (len(train_loader)-1):
                print('Epoch: [{0}][{1}/{2}] '
                    'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
                    'Elapsed {remain:s} '
                    'Loss: {loss.val:.4f}({loss.avg:.4f}) '
                    'Grad: {grad_norm:.4f}  '
                    #'LR: {lr:.6f}  '
                    .format(
                    epoch+1, step, len(train_loader), batch_time=batch_time,
                    data_time=data_time, loss=losses,
                    remain=timeSince(start, float(step+1)/len(train_loader)),
                    grad_norm=grad_norm,
                    #lr=scheduler.get_lr()[0],
                    ))
        return losses.avg

    def valid_fn(valid_loader, model, criterion, device):
        batch_time = AverageMeter()
        data_time = AverageMeter()
        losses = AverageMeter()
        scores = AverageMeter()
        # switch to evaluation mode
        model.eval()
        preds = []
        start = end = time.time()
        for step, (images, labels) in enumerate(valid_loader):
            # measure data loading time
            data_time.update(time.time() - end)
            images = images.to(device)
            labels = labels.to(device)
            batch_size = labels.size(0)
            # compute loss
            with torch.no_grad():
                y_preds = model(images)
            loss = criterion(y_preds, labels)
            losses.update(loss.item(), batch_size)
            # record accuracy
            preds.append(y_preds.softmax(1).to('cpu').numpy())
            if cfg['gradient_accumulation_steps'] > 1:
                loss = loss / cfg['gradient_accumulation_steps']
            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()
            if step % cfg['print_freq'] == 0 or step == (len(valid_loader)-1):
                print('EVAL: [{0}/{1}] '
                    'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
                    'Elapsed {remain:s} '
                    'Loss: {loss.val:.4f}({loss.avg:.4f}) '
                    .format(
                    step, len(valid_loader), batch_time=batch_time,
                    data_time=data_time, loss=losses,
                    remain=timeSince(start, float(step+1)/len(valid_loader)),
                    ))
        predictions = np.concatenate(preds)
        return losses.avg, predictions
    
    ##### Loop
    def train_loop(folds, fold):
    
        LOGGER.info(f"========== fold: {fold} training ==========")

        # ====================================================
        # loader
        # ====================================================
        trn_idx = folds[folds['fold'] != fold].index
        val_idx = folds[folds['fold'] == fold].index

        train_folds = folds.loc[trn_idx].reset_index(drop=True)
        valid_folds = folds.loc[val_idx].reset_index(drop=True)

        train_dataset = dataset.TrainDataset(train_folds, 
                                    transform=get_transforms(data='train'))
        valid_dataset = dataset.TrainDataset(valid_folds, 
                                    transform=get_transforms(data='valid'))

        train_loader = DataLoader(train_dataset, 
                                batch_size=cfg['batch_size'], 
                                shuffle=True, 
                                num_workers=cfg['num_workers'], pin_memory=True, drop_last=True)
        valid_loader = DataLoader(valid_dataset, 
                                batch_size=cfg['batch_size'], 
                                shuffle=False, 
                                num_workers=cfg['num_workers'], pin_memory=True, drop_last=False)
        
        # ====================================================
        # scheduler 
        # ====================================================
        def get_scheduler(optimizer):
            if cfg['scheduler']=='ReduceLROnPlateau':
                scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=cfg['factor'], patience=cfg['patience'], verbose=True, eps=cfg['eps'])
            elif cfg['scheduler']=='CosineAnnealingLR':
                scheduler = CosineAnnealingLR(optimizer, T_max=cfg['T_max'], eta_min=cfg['min_lr'], last_epoch=-1)
            elif cfg['scheduler']=='CosineAnnealingWarmRestarts':
                scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=cfg['T_0'], T_mult=1, eta_min=cfg['min_lr'], last_epoch=-1)
            return scheduler
        
        
        # ====================================================
        # model & optimizer
        # ====================================================
        # ResNext
        if cfg['model_name'][:7] == 'resnext':
            model = models.CustomResNext(cfg['model_name'], pretrained=True)
        # ViT
        elif cfg['model_name'][:3] == 'vit':
            model = models.VITModel(num_classes=cfg['target_size'], model_name=cfg['model_name'])
        # EfficientNet
        elif cfg['model_name'][:12] == 'efficientnet' or cfg['model_name'][3:12+3] == 'efficientnet':
            model = models.EfficientNetModel(num_classes=cfg['target_size'], model_name=cfg['model_name'])
        # Seresnext
        elif cfg['model_name'][:9] == 'seresnext':
            model = models.SeResNext(num_classes=cfg['target_size'], model_name=cfg['model_name'])
            
        # Check parallel
        if torch.cuda.device_count() > 1:
            print("Let's use", torch.cuda.device_count(), "GPUs!")
            # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
            model = nn.DataParallel(model)    
        model.to(device)

        optimizer = Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'], amsgrad=False)
        scheduler = get_scheduler(optimizer)
        
        ####################################################
        # Add loss
        def get_criterion():
            if cfg['criterion']=='CrossEntropyLoss':
                criterion = nn.CrossEntropyLoss()
            elif cfg['criterion']=='LabelSmoothing':
                criterion = LabelSmoothingLoss(classes=cfg['target_size'], smoothing=cfg['smoothing'])
            elif cfg['criterion']=='FocalLoss':
                criterion = FocalLoss().to(device)
            elif cfg['criterion']=='FocalCosineLoss':
                criterion = FocalCosineLoss()
            elif cfg['criterion']=='SymmetricCrossEntropyLoss':
                criterion = SymmetricCrossEntropy().to(device)
            elif cfg['criterion']=='BiTemperedLoss': 
                criterion = BiTemperedLogisticLoss(t1=cfg['t1'], t2=cfg['t2'], smoothing=cfg['smoothing'])
            elif cfg['criterion']=='TaylorCrossEntropyLoss':
                criterion = TaylorCrossEntropyLoss(smoothing=cfg['smoothing'])
            return criterion
        
        criterion = get_criterion()
        LOGGER.info(f'Criterion: {criterion}')
        ####################################################

        best_score = 0.
        best_loss = np.inf
        
        for epoch in range(cfg['epochs']):
            
            start_time = time.time()
            
            # train
            avg_loss = train_fn(train_loader, model, criterion, optimizer, epoch, scheduler, device)

            # eval
            avg_val_loss, preds = valid_fn(valid_loader, model, criterion, device)
            valid_labels = valid_folds[cfg['target_col']].values
            
            if isinstance(scheduler, ReduceLROnPlateau):
                scheduler.step(avg_val_loss)
            elif isinstance(scheduler, CosineAnnealingLR):
                scheduler.step()
            elif isinstance(scheduler, CosineAnnealingWarmRestarts):
                scheduler.step()

            # scoring
            score = get_score(valid_labels, preds.argmax(1))

            elapsed = time.time() - start_time

            LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  avg_val_loss: {avg_val_loss:.4f}  time: {elapsed:.0f}s')
            LOGGER.info(f'Epoch {epoch+1} - Accuracy: {score}')

            if score > best_score:
                best_score = score
                LOGGER.info(f'Epoch {epoch+1} - Save Best Score: {best_score:.4f} Model')
                
                m_name = cfg['model_name']
                torch.save({'model': model.state_dict(), 
                            'preds': preds},
                            OUTPUT_DIR+f'{m_name}_fold{fold}_best.pth')
        
        check_point = torch.load(OUTPUT_DIR+f'{m_name}_fold{fold}_best.pth')
        # CHANGE THIS TI CFG FOLD
        valid_folds[[str(c) for c in range(cfg['target_size'])]] = check_point['preds']
        valid_folds['preds'] = check_point['preds'].argmax(1)

        return valid_folds
    
    def main():

        def get_result(result_df):
            preds = result_df['preds'].values
            labels = result_df[cfg['target_col']].values
            score = get_score(labels, preds)
            LOGGER.info(f'Score: {score:<.5f}')
        
        if cfg['train']:
            # train 
            oof_df = pd.DataFrame()
            for fold in range(cfg['n_fold']):
                if fold in cfg['trn_fold']:
                    print(fold)
                    _oof_df = train_loop(folds, fold)
                    oof_df = pd.concat([oof_df, _oof_df])
                    LOGGER.info(f"========== fold: {fold} result ==========")
                    get_result(_oof_df)
            # CV result
            LOGGER.info(f"========== CV ==========")
            get_result(oof_df)
            # save result
            oof_df.to_csv(OUTPUT_DIR+'oof_df.csv', index=False)
    
    # START TRAINING
    main()