from typing import TYPE_CHECKING, Optional, Union
from easydict import EasyDict
import os
import numpy as np

from ding.utils import save_file
from ding.policy import Policy
from ding.framework import task

if TYPE_CHECKING:
    from ding.framework import OnlineRLContext, OfflineRLContext


class CkptSaver:
    """
        Overview:
            The class used to save checkpoint data.
    """

    def __new__(cls, *args, **kwargs):
        if task.router.is_active and not (task.has_role(task.role.LEARNER) or task.has_role(task.role.EVALUATOR)):
            return task.void()
        return super(CkptSaver, cls).__new__(cls)

    def __init__(self, policy: Policy, save_dir: str, train_freq: Optional[int] = None, save_finish: bool = True):
        """
        Overview:
            Initialize the `CkptSaver`.
        Arguments:
            - policy (:obj:`Policy`): Policy used to save the checkpoint.
            - save_dir (:obj:`str`): The directory path to save ckpt.
            - train_freq (:obj:`int`): Number of training iterations between each saving checkpoint data.
            - save_finish (:obj:`bool`): Whether save final ckpt when ``task.finish = True``.
        """
        self.policy = policy
        self.train_freq = train_freq
        if str(os.path.basename(os.path.normpath(save_dir))) != "ckpt":
            self.prefix = '{}/ckpt'.format(os.path.normpath(save_dir))
        else:
            self.prefix = '{}/'.format(os.path.normpath(save_dir))
        if not os.path.exists(self.prefix):
            os.makedirs(self.prefix)
        self.last_save_iter = 0
        self.max_eval_value = -np.inf
        self.save_finish = save_finish

    def __call__(self, ctx: Union["OnlineRLContext", "OfflineRLContext"]) -> None:
        """
        Overview:
            The method used to save checkpoint data. \
            The checkpoint data will be saved in a file in following 3 cases: \
                - When a multiple of `self.train_freq` iterations have elapsed since the beginning of training; \
                - When the evaluation episode return is the best so far; \
                - When `task.finish` is True.
        Input of ctx:
            - train_iter (:obj:`int`): Number of training iteration, i.e. the number of updating policy related network.
            - eval_value (:obj:`float`): The episode return of current iteration.
        """
        # train enough iteration
        if self.train_freq:
            if ctx.train_iter == 0 or ctx.train_iter - self.last_save_iter >= self.train_freq:
                save_file(
                    "{}/iteration_{}.pth.tar".format(self.prefix, ctx.train_iter), self.policy.learn_mode.state_dict()
                )
                self.last_save_iter = ctx.train_iter

        # best episode return so far
        if ctx.eval_value is not None and ctx.eval_value > self.max_eval_value:
            save_file("{}/eval.pth.tar".format(self.prefix), self.policy.learn_mode.state_dict())
            self.max_eval_value = ctx.eval_value

        # finish
        if task.finish and self.save_finish:
            save_file("{}/final.pth.tar".format(self.prefix), self.policy.learn_mode.state_dict())