#!/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()