# Copyright (c) Alibaba, Inc. and its affiliates.

import os
from collections import OrderedDict
from typing import Callable, Dict, Optional

import numpy as np
import torch
import torch.nn as nn
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from sklearn.metrics import confusion_matrix
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset

from modelscope.metainfo import Trainers
from modelscope.models.base import Model
from modelscope.msdatasets import MsDataset
from modelscope.trainers.base import BaseTrainer
from modelscope.trainers.builder import TRAINERS
from modelscope.trainers.multi_modal.team.team_trainer_utils import (
    get_optimizer, train_mapping, val_mapping)
from modelscope.utils.config import Config
from modelscope.utils.constant import DownloadMode, ModeKeys
from modelscope.utils.logger import get_logger

logger = get_logger()


@TRAINERS.register_module(module_name=Trainers.image_classification_team)
class TEAMImgClsTrainer(BaseTrainer):

    def __init__(self, cfg_file: str, model: str, device_id: int,
                 data_collator: Callable, train_dataset: Dataset,
                 val_dataset: Dataset, *args, **kwargs):
        super().__init__(cfg_file)

        self.cfg = Config.from_file(cfg_file)
        team_model = Model.from_pretrained(model)
        image_model = team_model.model.image_model.vision_transformer
        classification_model = nn.Sequential(
            OrderedDict([('encoder', image_model),
                         ('classifier',
                          nn.Linear(768, self.cfg.dataset.class_num))]))
        self.model = classification_model

        for pname, param in self.model.named_parameters():
            if 'encoder' in pname:
                param.requires_grad = False

        self.device_id = device_id
        self.total_epoch = self.cfg.train.epoch
        self.train_batch_size = self.cfg.train.batch_size
        self.val_batch_size = self.cfg.evaluation.batch_size
        self.ckpt_dir = self.cfg.train.ckpt_dir

        self.collate_fn = data_collator
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset

        self.criterion = nn.CrossEntropyLoss().to(self.device_id)

    def train(self, *args, **kwargs):
        self.model.train()
        self.model.to(self.device_id)

        optimizer = get_optimizer(self.model)

        for epoch in range(self.total_epoch):
            train_params = {
                'pin_memory': True,
                'collate_fn': self.collate_fn,
                'batch_size': self.train_batch_size,
                'shuffle': True,
                'drop_last': True,
                'num_workers': 8
            }

            train_loader = DataLoader(self.train_dataset, **train_params)

            for batch_idx, data in enumerate(train_loader):
                img_tensor, label_tensor = data['pixel_values'], data['labels']
                img_tensor = img_tensor.to(self.device_id, non_blocking=True)
                label_tensor = label_tensor.to(
                    self.device_id, non_blocking=True)

                pred_logits = self.model(img_tensor)
                loss = self.criterion(pred_logits, label_tensor)

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                if batch_idx % 10 == 0:
                    logger.info(
                        'epoch: {}, train batch {}/{}, loss={:.5f}'.format(
                            epoch, batch_idx, len(train_loader), loss.item()))

            os.makedirs(self.ckpt_dir, exist_ok=True)
            torch.save(self.model.state_dict(),
                       '{}/epoch{}.pth'.format(self.ckpt_dir, epoch))
            self.evaluate()

    def evaluate(self,
                 checkpoint_path: Optional[str] = None,
                 *args,
                 **kwargs) -> Dict[str, float]:
        if checkpoint_path is not None:
            checkpoint_params = torch.load(checkpoint_path, 'cpu')
            self.model.load_state_dict(checkpoint_params)
        self.model.eval()
        self.model.to(self.device_id)

        val_params = {
            'collate_fn': self.collate_fn,
            'batch_size': self.val_batch_size,
            'shuffle': False,
            'drop_last': False,
            'num_workers': 8
        }
        val_loader = DataLoader(self.val_dataset, **val_params)

        tp_cnt, processed_cnt = 0, 0
        all_pred_labels, all_gt_labels = [], []
        with torch.no_grad():
            for batch_idx, data in enumerate(val_loader):
                img_tensor, label_tensor = data['pixel_values'], data['labels']
                img_tensor = img_tensor.to(self.device_id, non_blocking=True)
                label_tensor = label_tensor.to(
                    self.device_id, non_blocking=True)

                pred_logits = self.model(img_tensor)
                pred_labels = torch.max(pred_logits, dim=1)[1]
                tp_cnt += torch.sum(pred_labels == label_tensor).item()
                processed_cnt += img_tensor.shape[0]
                logger.info('Accuracy: {:.3f}'.format(tp_cnt / processed_cnt))

                all_pred_labels.extend(pred_labels.tolist())
                all_gt_labels.extend(label_tensor.tolist())
        conf_mat = confusion_matrix(all_gt_labels, all_pred_labels)
        acc_mean_per_class = np.mean(conf_mat.diagonal()
                                     / conf_mat.sum(axis=1))
        logger.info(
            'Accuracy mean per class: {:.3f}'.format(acc_mean_per_class))
