#!/usr/bin/env python3
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
import bisect
import datetime
import math
import os
import pickle
import time
from typing import Optional

import megengine as mge
import megengine.distributed as dist
import megengine.module as M
from basecore.config import ConfigDict
from basecore.engine import BaseHook, BaseTrainer
from basecore.utils import (
    Checkpoint,
    MeterBuffer,
    cached_property,
    ensure_dir,
    get_last_call_deltatime,
)
from loguru import logger
from tensorboardX import SummaryWriter

from basecls.layers import compute_precise_bn_stats
from basecls.models import sync_model
from basecls.utils import default_logging, registers

from .tester import ClsTester

__all__ = [
    "CheckpointHook",
    "EvalHook",
    "LoggerHook",
    "LRSchedulerHook",
    "PreciseBNHook",
    "ResumeHook",
    "TensorboardHook",
]


def _create_checkpoint(trainer: BaseTrainer, save_dir: str) -> Checkpoint:
    """Create a checkpoint for save and resume"""
    model = trainer.model
    ema = trainer.ema
    ckpt_kws = {"ema": ema} if ema is not None else {}
    optim = trainer.solver.optimizer
    scaler = trainer.solver.grad_scaler
    progress = trainer.progress
    ckpt = Checkpoint(
        save_dir,
        model,
        tag_file=None,
        optimizer=optim,
        scaler=scaler,
        progress=progress,
        **ckpt_kws,
    )
    return ckpt


class CheckpointHook(BaseHook):
    """Hook for managing checkpoints during training.

    Effect during ``after_epoch`` and ``after_train`` procedure.

    Args:
        save_dir: checkpoint directory.
        save_every_n_epoch: interval for saving checkpoint. Default: ``1``
    """

    def __init__(self, save_dir: str = None, save_every_n_epoch: int = 1):
        super().__init__()
        ensure_dir(save_dir)
        self.save_dir = save_dir
        self.save_every_n_epoch = save_every_n_epoch

    def after_epoch(self):
        progress = self.trainer.progress
        ckpt = _create_checkpoint(self.trainer, self.save_dir)
        ckpt.save("latest.pkl")

        if progress.epoch % self.save_every_n_epoch == 0:
            progress_str = progress.progress_str_list()
            save_name = "_".join(progress_str[:-1]) + ".pkl"
            ckpt.save(save_name)
            logger.info(f"Save checkpoint {save_name} to {self.save_dir}")

    def after_train(self):
        # NOTE: usually final ema is not the best so we dont save it
        mge.save(
            {"state_dict": self.trainer.model.state_dict()},
            os.path.join(self.save_dir, "dumped_model.pkl"),
            pickle_protocol=pickle.DEFAULT_PROTOCOL,
        )


class EvalHook(BaseHook):
    """Hook for evaluating during training.

    Effect during ``after_epoch`` and ``after_train`` procedure.

    Args:
        save_dir: checkpoint directory.
        eval_every_n_epoch: interval for evaluating. Default: ``1``
    """

    def __init__(self, save_dir: str = None, eval_every_n_epoch: int = 1):
        super().__init__()
        ensure_dir(save_dir)
        self.save_dir = save_dir
        self.eval_every_n_epoch = eval_every_n_epoch
        self.best_acc1 = 0
        self.best_ema_acc1 = 0

    def after_epoch(self):
        trainer = self.trainer
        cfg = trainer.cfg
        model = trainer.model
        ema = trainer.ema
        progress = trainer.progress
        if progress.epoch % self.eval_every_n_epoch == 0 and progress.epoch != progress.max_epoch:
            self.test(cfg, model, ema)

    def after_train(self):
        trainer = self.trainer
        cfg = trainer.cfg
        model = trainer.model
        ema = trainer.ema
        # TODO: actually useless maybe when precise_bn is on
        sync_model(model)
        if ema is not None:
            sync_model(ema)
        self.test(cfg, model, ema)

    def test(self, cfg: ConfigDict, model: M.Module, ema: Optional[M.Module] = None):
        dataloader = registers.dataloaders.get(cfg.data.name).build(cfg, False)
        # FIXME: need atomic user_pop, maybe in MegEngine 1.5?
        # tester = BaseTester(model, dataloader, AccEvaluator())
        tester = ClsTester(cfg, model, dataloader)
        acc1, _ = tester.test()

        if acc1 > self.best_acc1:
            self.best_acc1 = acc1
            if dist.get_rank() == 0:
                mge.save(
                    {"state_dict": model.state_dict(), "acc1": self.best_acc1},
                    os.path.join(self.save_dir, "best_model.pkl"),
                    pickle_protocol=pickle.DEFAULT_PROTOCOL,
                )

        logger.info(
            f"Epoch: {self.trainer.progress.epoch}, Test Acc@1: {acc1:.3f}, "
            f"Best Test Acc@1: {self.best_acc1:.3f}"
        )

        if ema is None:
            return

        tester_ema = ClsTester(cfg, ema, dataloader)
        ema_acc1, _ = tester_ema.test()

        if ema_acc1 > self.best_ema_acc1:
            self.best_ema_acc1 = ema_acc1
            if 
dist.get_rank()