#!/usr/bin/env python3
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
import copy
import time
from typing import Iterable

import megengine as mge
import megengine.amp as amp
import megengine.distributed as dist
import megengine.functional as F
import megengine.module as M
import megengine.optimizer as optim
from basecore.config import ConfigDict
from basecore.engine import BaseHook, BaseTrainer
from basecore.utils import MeterBuffer
from megengine import jit

from basecls.data import DataLoaderType
from basecls.layers import Preprocess, build_loss
from basecls.solver import Solver
from basecls.utils import registers

__all__ = ["ClsTrainer"]


@registers.trainers.register()
class ClsTrainer(BaseTrainer):
    """Classification trainer.

    Args:
        cfg: config for training.
        model: model for training.
        dataloader: dataloader for training.
        solver: solver for training.
        hooks: hooks for training.

    Attributes:
        cfg: config for training.
        model: model for training.
        ema: model exponential moving average.
        dataloader: dataloader for training.
        solver: solver for training.
        progress: object for recording training process.
        loss: loss function for training.
        meter : object for recording metrics.
    """

    def __init__(
        self,
        cfg: ConfigDict,
        model: M.Module,
        dataloader: DataLoaderType,
        solver: Solver,
        hooks: Iterable[BaseHook] = None,
    ):
        super().__init__(model, dataloader, solver, hooks)
        self.cfg = cfg
        self.ema = copy.deepcopy(model) if cfg.model_ema.enabled else None
        self.preprocess = Preprocess(cfg.preprocess.img_mean, cfg.preprocess.img_std)
        self.loss = build_loss(cfg)
        self.meter = MeterBuffer(cfg.log_every_n_iter)
        if cfg.trace:
            # FIXME: tracing makes the training slower than before, why?
            self.model_step = jit.trace(self.model_step, symbolic=True)

    def train(self):
        start_training_info = (1, 1)
        max_iter = len(self.dataloader)
        max_training_info = (self.cfg.solver.max_epoch, max_iter)
        super().train(start_training_info, max_training_info)

    def before_train(self):
        super().before_train()

    def before_epoch(self):
        super().before_epoch()
        self.dataloader_iter = iter(self.dataloader)

    def after_epoch(self):
        del self.dataloader_iter
        super().after_epoch()

    def train_one_iter(self):
        """Basic logic of training one iteration."""
        data_tik = time.perf_counter()
        data = next(self.dataloader_iter)
        samples, targets = self.preprocess(data)
        mge._full_sync()  # use full_sync func to sync launch queue for dynamic execution
        data_tok = time.perf_counter()

        train_tik = time.perf_counter()
        losses, accs = self.model_step(samples, targets)
        mge._full_sync()  # use full_sync func to sync launch queue for dynamic execution
        train_tok = time.perf_counter()

        # TODO: stats and accs
        loss_meters = {"loss": losses.item()}
        stat_meters = {"stat_acc@1": accs[0].item() * 100, "stat_acc@5": accs[1].item() * 100}
        time_meters = {"train_time": train_tok - train_tik, "data_time": data_tok - data_tik}
        self.meter.update(**loss_meters, **stat_meters, **time_meters)

    def model_step(self, samples, targets):
        optimizer = self.solver.optimizer
        grad_manager = self.solver.grad_manager
        grad_scaler = self.solver.grad_scaler

        with grad_manager:
            with amp.autocast(enabled=self.cfg.amp.enabled):
                outputs = self.model(samples)
                losses = self.loss(outputs, targets)

            if isinstance(losses, mge.Tensor):
                total_loss = losses
            elif isinstance(losses, dict):
                if "total_loss" in losses:
                    total_loss = losses["total_loss"]
                else:
                    # only key contains "loss" will be calculated.
                    total_loss = sum([v for k, v in losses.items() if "loss" in k])
                    losses["total_loss"] = total_loss
            else:
                # list or tuple
                total_loss = sum(losses)

            total_loss = total_loss / self.cfg.solver.accumulation_steps

            # this is made compatible with one hot labels
            if targets.ndim == 2:
                targets = F.argmax(targets, axis=1)
            accs = F.metric.topk_accuracy(outputs, targets, (1, 5))

            if self.cfg.amp.enabled:
                grad_scaler.backward(grad_manager, total_loss)
            else:
                grad_manager.backward(total_loss)

        if self.progress.iter % self.cfg.solver.accumulation_steps == 0:
            self.modify_grad()

            optimizer.step().clear_grad()

            self.model_ema_step()

        return losses, accs

    def modify_grad(self):
        grad_cfg = self.cfg.solver.grad_clip
        # TODO: support advanced params for grad clip in the future
        params = self.model.parameters()
        if grad_cfg.name is None:
            return
        elif grad_cfg.name == "norm":
            
optim.clip_grad_norm(params, grad_cfg.max_norm)