# -*- coding: utf-8 -*-
# This repo is licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import os
import pickle
import boto3
import megengine as mge

from collections import defaultdict
from termcolor import colored

from common import utils


class Manager():
    def __init__(self, model, optimizer, params, dataloaders, writer, logger, scheduler):
        # params status
        self.model = model
        self.writer = writer
        self.logger = logger
        self.params = params
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.dataloaders = dataloaders

        # metric_rule should be either Descende or Ascende
        self.metric_rule = params.metric_rule

        self.epoch = 0
        self.step = 0

        # 越低越好
        if self.metric_rule == "Descende":
            self.best_val_score = 100
            self.best_test_score = 100
        # 越高越好
        elif self.metric_rule == "Ascende":
            self.best_val_score = 0
            self.best_test_score = 0

        self.cur_val_score = 0
        self.cur_test_score = 0

        # train status
        self.train_status = defaultdict(utils.AverageMeter)

        # val status
        self.val_status = defaultdict(utils.AverageMeter)

        # test status
        self.test_status = defaultdict(utils.AverageMeter)

        # model status
        self.loss_status = defaultdict(utils.AverageMeter)

        # client init
        self.s3_client = boto3.client('s3', endpoint_url='http://oss.i.brainpp.cn')
        self.bucket_name = params.bucket_name

    def update_step(self):
        self.step += 1

    def update_epoch(self):
        self.epoch += 1

    def update_loss_status(self, loss, split, bs=None):
        if split == "train":
            for k, v in loss.items():
                bs = self.params.train_batch_size
                self.loss_status[k].update(val=v.item(), num=bs)
        elif split == "val":
            for k, v in loss.items():
                self.loss_status[k].update(val=v.item(), num=bs)
        elif split == "test":
            for k, v in loss.items():
                self.loss_status[k].update(val=v.item(), num=bs)
        else:
            raise ValueError("Wrong eval type: {}".format(split))

    def update_metric_status(self, metrics, split, bs):
        if split == "val":
            for k, v in metrics.items():
                self.val_status[k].update(val=v.item(), num=bs)
                self.cur_val_score = self.val_status[self.params.major_metric].avg
        elif split == "test":
            for k, v in metrics.items():
                self.test_status[k].update(val=v.item(), num=bs)
                self.cur_test_score = self.test_status[self.params.major_metric].avg
        else:
            raise ValueError("Wrong eval type: {}".format(split))

    def reset_loss_status(self):
        for k, v in self.loss_status.items():
            self.loss_status[k].reset()

    def reset_metric_status(self, split):
        if split == "val":
            for k, v in self.val_status.items():
                self.val_status[k].reset()
        elif split == "test":
            for k, v in self.test_status.items():
                self.test_status[k].reset()
        else:
            raise ValueError("Wrong eval type: {}".format(split))

    def print_train_info(self):
        exp_name = self.params.model_dir.split('/')[-1]
        print_str = "{} Epoch: {:4d}, lr={:.6f} ".format(exp_name, self.epoch, self.scheduler.get_lr()[0])
        print_str += "total loss: %.4f(%.4f)" % (self.loss_status['total'].val, self.loss_status['total'].avg)
        return print_str

    def print_metrics(self, split, title="Eval", color="red"):
        if split == "val":
            metric_status = self.val_status
        elif split == "test":
            metric_status = self.test_status
        else:
            raise ValueError("Wrong eval type: {}".format(split))
        if split == "val":
            print_str = " | ".join("{}: {:.4f} - avg: {:.4f}".format(k, v.val, v.avg) for k, v in metric_status.items())
            print_str += " | lastest: {:.4f} | pre_best: {:.4f}".format(self.cur_val_score, self.best_val_score)
        elif split == "test":
            print_str = " | ".join("{}: {:.4f} - avg: {:.4f}".format(k, v.val, v.avg) for k, v in metric_status.items())
        self.logger.info(colored("{} Results: {}".format(title, print_str), color, attrs=["bold"]))

    def check_best_save_last_checkpoints(self, latest_freq=5):

        state = {
            "state_dict": self.model.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            "scheduler": self.scheduler.state_dict(),
            "step": self.step,
            "epoch": self.epoch,
        }
        if "val" in self.dataloaders:
            state["best_val_score"] = self.best_val_score
        if "test" in self.dataloaders:
            state["best_test_score"] = self.best_test_score

        # save latest checkpoint
        if self.epoch % latest_freq == 0:
            latest_ckpt_name = os.path.join(self.params.model_dir, "model_latest.pth")
            if self.params.save_mode == "local":
                # print(state)
                
mge.save(state, latest_ckpt_name)