from typing import TYPE_CHECKING, Callable, Union
from easydict import EasyDict
import treetensor.torch as ttorch
from ditk import logging
import numpy as np
from ding.policy import Policy
from ding.framework import task, OfflineRLContext, OnlineRLContext


def trainer(cfg: EasyDict, policy: Policy, log_freq: int = 100) -> Callable:
    """
    Overview:
        The middleware that executes a single training process.
    Arguments:
        - cfg (:obj:`EasyDict`): Config.
        - policy (:obj:`Policy`): The policy to be trained in step-by-step mode.
        - log_freq (:obj:`int`): The frequency (iteration) of showing log.
    """
    if task.router.is_active and not task.has_role(task.role.LEARNER):
        return task.void()

    def _train(ctx: Union["OnlineRLContext", "OfflineRLContext"]):
        """
        Input of ctx:
            - train_data (:obj:`Dict`): The data used to update the network. It will train only if \
                the data is not empty.
            - train_iter: (:obj:`int`): The training iteration count. The log will be printed once \
                it reachs certain values.
        Output of ctx:
            - train_output (:obj:`Dict`): The training output in the Dict format, including loss info.
        """

        if ctx.train_data is None:
            return
        train_output = policy.forward(ctx.train_data)
        if ctx.train_iter % log_freq == 0:
            if isinstance(train_output, list):
                train_output_loss = np.mean([item['total_loss'] for item in train_output])
            else:
                train_output_loss = train_output['total_loss']
            if isinstance(ctx, OnlineRLContext):
                logging.info(
                    'Training: Train Iter({})\tEnv Step({})\tLoss({:.3f})'.format(
                        ctx.train_iter, ctx.env_step, train_output_loss
                    )
                )
            elif isinstance(ctx, OfflineRLContext):
                logging.info('Training: Train Iter({})\tLoss({:.3f})'.format(ctx.train_iter, train_output_loss))
            else:
                raise TypeError("not supported ctx type: {}".format(type(ctx)))
        ctx.train_iter += 1
        ctx.train_output = train_output

    return _train


def multistep_trainer(policy: Policy, log_freq: int = 100) -> Callable:
    """
    Overview:
        The middleware that executes training for a target num of steps.
    Arguments:
        - policy (:obj:`Policy`): The policy specialized for multi-step training.
        - log_freq (:obj:`int`): The frequency (iteration) of showing log.
    """
    if task.router.is_active and not task.has_role(task.role.LEARNER):
        return task.void()
    last_log_iter = -1

    def _train(ctx: Union["OnlineRLContext", "OfflineRLContext"]):
        """
        Input of ctx:
            - train_data: The data used to update the network.
                It will train only if the data is not empty.
            - train_iter: (:obj:`int`): The training iteration count.
                The log will be printed if it reachs certain values.
        Output of ctx:
            - train_output (:obj:`List[Dict]`): The training output listed by steps.
        """

        if ctx.train_data is None:  # no enough data from data fetcher
            return
        if hasattr(policy, "_device"):  # For ppof policy
            data = ctx.train_data.to(policy._device)
        elif hasattr(policy, "get_attribute"):  # For other policy
            data = ctx.train_data.to(policy.get_attribute("device"))
        else:
            assert AttributeError("Policy should have attribution '_device'.")
        train_output = policy.forward(data)
        nonlocal last_log_iter
        if ctx.train_iter - last_log_iter >= log_freq:
            loss = np.mean([o['total_loss'] for o in train_output])
            if isinstance(ctx, OfflineRLContext):
                logging.info('Training: Train Iter({})\tLoss({:.3f})'.format(ctx.train_iter, loss))
            else:
                logging.info(
                    'Training: Train Iter({})\tEnv Step({})\tLoss({:.3f})'.format(ctx.train_iter, ctx.env_step, loss)
                )
            last_log_iter = ctx.train_iter
        ctx.train_iter += len(train_output)
        ctx.train_output = train_output

    return _train


# TODO reward model