"""
This is a fine-tuning notebook which uses EfficientNet-b0 imagenet pretrained model as backbone. You can use any model from the `timm` library. Available models can be found via the `timm.list_models()` function.

The models can be used as a backbone in the ongoing [Shopee - Price Match Guarantee](https://www.kaggle.com/c/shopee-product-matching/) Challange.
"""
!pip install timm
"""
# Imports
"""
import os
import random
from pathlib import Path
from tqdm import tqdm

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import torch

import cv2
import albumentations

import timm
BASE_DATA_DIR = Path("../input/shopee-product-detection/")

df_train = pd.read_csv(BASE_DATA_DIR / "train.csv")
df_test = pd.read_csv(BASE_DATA_DIR / "test.csv")

df_train = df_train.loc[~df_train.filename.isin(["64faf0b221af4767ba8c167b228fde00.jpg", 
                                                 "d946ee19ac1d2997bac5f18ce75656cb.jpg"])].reset_index(drop=True)
counts = df_train.category.value_counts()
df_train.category.max(), df_train.category.min()
plt.figure(figsize=(16, 10))
plt.bar(counts.index, counts)
plt.xticks(range(42));

plt.show()
"""
# Utilities
"""
import time
from contextlib import contextmanager

LOGS_PATH = Path("logs")
LOGS_PATH.mkdir(exist_ok=True)


def init_logger(log_file=LOGS_PATH / '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


LOGGER = init_logger()


@contextmanager
def timer(name):
    t0 = time.time()
    LOGGER.info(f'[{name}] start')
    yield
    LOGGER.info(f'[{name}] done in {time.time() - t0:.0f} s.')
"""
# Simple Visualization
"""
BASE_IMG_DIR = Path("../input/shopee-product-detection/train/train/")

def read_img_and_cvt_format(img_path, clr_format=cv2.COLOR_BGR2RGB):
    return cv2.cvtColor(cv2.imread(img_path), clr_format)

def visualize_batch(img_ids, labels):
    
    plt.figure(figsize=(16, 12))
    
    for idx, (img_id, label) in enumerate(zip(img_ids, labels)):
        plt.subplot(3, 3, idx + 1)
        img_fn = str(BASE_IMG_DIR / img_id)
        img = read_img_and_cvt_format(img_fn)
        plt.imshow(img)
        plt.title(f"Class: {label}", fontsize=9)
        plt.axis("off")
        
    plt.show()
df_train.columns
sampled_df = df_train.sample(9)
img_ids = sampled_df["filename"].values
labels = sampled_df["category"].values

visualize_batch(img_ids, labels)
"""
# Dataset
"""
import torch
from torch.utils.data import Dataset, DataLoader

class ShopeeDataset(Dataset):
    
    def __init__(self, image_paths, labels=None, transform=None):
        
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform
        
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        
        img_filepath = self.image_paths[idx]
        img = read_img_and_cvt_format(img_filepath)
        if self.transform:
            img = self.transform(image=img)["image"]
        
        label = 0
        if self.labels is not None:
            label = torch.tensor(self.labels[idx]).long()
        return img, label

train_img_paths = [f"{BASE_IMG_DIR}/{img_id}" for img_id in df_train["filename"].values]
train_dataset = ShopeeDataset(image_paths=train_img_paths, 
                               labels=df_train["category"].values,
                               transform=None)

for i in range(1):
    img, label = train_dataset[i]
    
    plt.title(f"Label: {label}")
    plt.imshow(img)

plt.show()
len(df_train.category.unique())
"""
# Config
"""
class Config:
    
    model_name = "efficientnet_b0" # resnet34
    n_epochs = 10
    batch_size = 32
    img_size = 512
    n_classes = len(df_train.category.unique())
    lr = 1e-3
    weight_decay = 1e-6
    gradient_accumulation_steps = 1
    max_grad_norm = 1000
    seed = 42
    scheduler = ""
    n_fold = 1
    train_fold = [0, 1, 2, 3, 4]
    train = True
    print_every = 100
    num_workers = 4
    

def seed_torch(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

seed_torch(seed=Config.seed)
"""
# Classifier
"""
import torch.nn as nn

class Classifier(nn.Module):
    
    def __init__(self, model_name, pretrained=False):
        super(Classifier, self).__init__()
        
        self.model = timm.create_model(model_name, pretrained=pretrained)
        if model_name.startswith("eff"):
            n_features = self.model.classifier.in_features
            self.model.classifier = nn.Linear(n_features, Config.n_classes)
        else:    
            n_features = self.model.fc.in_features
            self.model.fc = nn.Linear(n_features, Config.n_classes)
        
    def forward(self, x):
        return self.model(x)
"""
# Augmentations
"""
from albumentations.pytorch import ToTensorV2
from torchvision import transforms as T

def get_train_transforms():
    return albumentations.Compose([
        albumentations.Resize(
            Config.img_size, Config.img_size),
        albumentations.Transpose(),
        albumentations.HorizontalFlip(),
        albumentations.VerticalFlip(),
        albumentations.ShiftScaleRotate(),
        albumentations.Normalize(
            mean=[0.485, 0.456, 0.406], 
            std=[0.229, 0.224, 0.225]),
        albumentations.Cutout(num_holes=8, max_h_size=32, max_w_size=32, fill_value=0, p=0.5),
        ToTensorV2(),
    ])


def get_test_transforms():
    
    return albumentations.Compose([
        albumentations.Resize(Config.img_size, Config.img_size),
        albumentations.Normalize(mean=[0.485, 0.456, 0.406], 
                  std=[0.229, 0.224, 0.225]),
        ToTensorV2()
    ])

"""
# Metric Tracking
"""
import math
import time


class AverageMeter:
    
    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 as_minutes(s):
    m = math.floor(s / 60)
    s -= m * 60
    return f"{m}m {s}s"


def time_since(since, percent):
    now = time.time()
    s = now - since
    es = s / percent
    rs = es - s
    return f"{as_minutes(s)} (remain {as_minutes(rs)})"
def train_step(model, data_loader, criterion, optimizer, epoch, scheduler, device):
    """
    There is no scheduler update currently.
    """
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    # scores = AverageMeter()
    
    model.train()
    start = end = time.time()
    # global_step = 0
    total_len = len(data_loader)
    
    for step, (images, labels) in enumerate(data_loader):
        
        data_time.update(time.time() - end)
        images = images.to(device)
        labels = labels.to(device)
        batch_size = labels.size(0)
        preds = model(images)
        loss = criterion(preds, labels)
        losses.update(loss.item(), batch_size)
        
        if Config.gradient_accumulation_steps > 1:
            loss = loss / Config.gradient_accumulation_steps
        
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 
                                                   Config.max_grad_norm)
        if (step + 1) % Config.gradient_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
            scheduler.step()
            # global_step += 1
        
        batch_time.update(time.time() - end)
        end = time.time()
        if step % Config.print_every == 0 or step == (total_len - 1):
            print(f"Epoch: [{epoch+1}][{step}/{total_len}] "
                  f"Data: {data_time.val:.3f} ({data_time.avg:.3f}) "
                  f"Batch: {batch_time.val:.3f} ({batch_time.avg:.3f}) "
                  f"Elapsed: {time_since(start, float(step + 1) / (total_len))} "
                  f"Loss: {losses.val:.5f}({losses.avg:.5f}) "
                  f"Grad: {grad_norm:.4f}" # LR: {lr:.6f}
                 )
        
    return losses.avg
            

def valid_step(model, data_loader, criterion, device):
    
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    scores = AverageMeter()
    
    model.eval()
    start = end = time.time()
    total_len = len(data_loader)
    predictions = []
    
    for step, (images, labels) in enumerate(data_loader):
        data_time.update(time.time() - end)
        images = images.to(device)
        labels = labels.to(device)
        batch_size = labels.size(0)
        
        with torch.no_grad():
            preds = model(images)
        
        loss = criterion(preds, labels)
        losses.update(loss.item(), batch_size)
        predictions.append(preds.softmax(1).cpu().numpy())
        
        if Config.gradient_accumulation_steps > 1:
            loss = loss / Config.gradient_accumulation_steps
            
        batch_time.update(time.time() - end)
        end = time.time()
        
        if step % Config.print_every == 0 or step == (total_len - 1):
            print(f"Eval: [{step}/{total_len}] "
                  f"Data: {data_time.val:.3f} ({data_time.avg:.3f}) "
                  f"Batch: {batch_time.val:.3f} ({batch_time.avg:.3f}) "
                  f"Elapsed: {time_since(start, float(step + 1) / total_len)} "
                  f"Loss: {losses.val:.5f} ({losses.avg:.5f})"
                 )
    
    predictions = np.concatenate(predictions)
    return losses.avg, predictions
# !rm -rf models
import torch.optim as optim
from sklearn.metrics import accuracy_score, classification_report

device = "cuda" if torch.cuda.is_available() else "cpu"

MODELS_DIR = Path("models")
MODELS_DIR.mkdir(exist_ok=False)
"""
# Train Loop
"""
def train_loop(df_tr, df_val):

    train_img_paths = [f"{BASE_IMG_DIR}/{img_id}" for img_id in df_tr["filename"].values]
    valid_img_paths = [f"{BASE_IMG_DIR}/{img_id}" for img_id in df_val["filename"].values]
    
    train_dataset = ShopeeDataset(
        train_img_paths, 
        labels=df_tr["category"].values, 
        transform=get_train_transforms()
    )
    
    valid_dataset = ShopeeDataset(
        valid_img_paths,
        labels=df_val["category"].values,
        transform=get_test_transforms()
    )
    
    train_data_loader = DataLoader(
        train_dataset, batch_size=Config.batch_size, 
        shuffle=True, num_workers=Config.num_workers
    )
    valid_data_loader = DataLoader(
        valid_dataset, batch_size=Config.batch_size, 
        shuffle=False, num_workers=Config.num_workers
    )
    
    model = Classifier(Config.model_name, pretrained=True)
    model.to(device)
    # amsgrad = False
    optimizer = optim.Adam(model.parameters(), 
                           lr=Config.lr, 
                           weight_decay=Config.weight_decay)
    criterion = nn.CrossEntropyLoss()

        
    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 
                                                                     T_0=10, 
                                                                     T_mult=1, 
                                                                     eta_min=1e-6, 
                                                                     last_epoch=-1)
    
    best_score = 0.0
    best_loss = np.inf

    for epoch in range(Config.n_epochs):
        
        start_time = time.time()
        avg_epoch_loss = train_step(model, 
                                    train_data_loader, 
                                    criterion, 
                                    optimizer, 
                                    epoch, 
                                    scheduler=scheduler, 
                                    device=device)

        avg_valid_loss, valid_preds = valid_step(model, 
                                                 valid_data_loader, 
                                                 criterion, 
                                                 device)
        valid_labels = df_val["category"].values
        accuracy = accuracy_score(valid_labels, valid_preds.argmax(1))
        classification_result = classification_report(valid_labels, 
                                                      valid_preds.argmax(1))
        elapsed = time.time() - start_time
        LOGGER.info(f"Epoch: {epoch+1} - avg_epoch_loss: {avg_epoch_loss:.5f} - avg_val_loss: {avg_valid_loss:.5f} - time: {elapsed:.0f}s")
        LOGGER.info(f"Epoch: {epoch+1} - Accuracy: {accuracy}")
        print(classification_result)
        
        if accuracy > best_score:
            best_score = accuracy
            LOGGER.info(f"Epoch: {epoch+1} - Save best score: {best_score:.4f} Model")
            torch.save({
                "model": model.state_dict(),
                "preds": valid_preds
            }, str(MODELS_DIR / f"{Config.model_name}_best.pth"))
            
#     check_point = torch.load(str(MODELS_DIR / f"{Config.model_name}_fold_{fold}_best.pth"))
#     valid_folds[[str(c) for c in range(5)]] = check_point["preds"]
#     valid_folds["preds"] = check_point["preds"].argmax(1)
#     return valid_folds

from sklearn.model_selection import train_test_split
df_train_, df_valid_ = train_test_split(df_train, test_size=0.2, random_state=42)
len(df_train_.category.unique()), len(df_valid_.category.unique())
train_loop(df_train_, df_valid_)