# pylint:disable=missing-module-docstring
from abc import ABC, abstractmethod
from typing import Dict, List, Tuple

from ray.rllib import SampleBatch

from raylab.options import option
from raylab.utils.timer import TimerStat
from raylab.utils.types import StatDict

from ..stats import learner_stats


def model_based_options(cls: type) -> type:
    """Decorator to add default model-based options used by MBPolicyMixin."""
    model_update_interval = option(
        "model_update_interval",
        default=1,
        help="""Number of calls to `learn_on_batch` between each model update run.

        Example:
            With a 'rollout_fragment_length' of 1 and 'model_update_interval' of 25,
            will collect 25 environment transitions between each model optimization
            loop.
        """,
    )
    for opt in [model_update_interval]:
        cls = opt(cls)

    return cls


class MBPolicyMixin(ABC):
    """Off-policy mixin with dynamics model learning."""

    timers: Dict[str, TimerStat]
    _learn_calls: int = 0
    _info: dict

    def build_timers(self):
        """Create timers for model and policy training."""
        self.timers = {"model": TimerStat(), "policy": TimerStat()}
        self._info = {}

    @learner_stats
    def learn_on_batch(self, samples: SampleBatch) -> dict:
        # pylint:disable=missing-function-docstring
        self.add_to_buffer(samples)
        self._learn_calls += 1

        warmup = self._learn_calls == 1
        if (self._learn_calls % self.config["model_update_interval"] == 0) or warmup:
            with self.timers["model"] as timer:
                _, model_info = self.train_dynamics_model(warmup=warmup)
                timer.push_units_processed(model_info["model_epochs"])
                self._info.update(model_info)

        with self.timers["policy"] as timer:
            times = self.config["improvement_steps"]
            policy_info = self.update_policy(times=times)
            timer.push_units_processed(times)
            self._info.update(policy_info)

        self._info.update(self.timer_stats())
        return self._info.copy()

    @abstractmethod
    def train_dynamics_model(
        self, warmup: bool = False
    ) -> Tuple[List[float], StatDict]:
        """Implements the model training step.

        Calls the policy to optimize the model on the environment replay buffer.

        Args:
            warmup: Whether the optimization is being done on data collected
                via :meth:`sample_until_learning_starts`.

        Returns:
            A tuple containing the list of evaluation losses for each model and
            a dictionary of training statistics
        """

    def update_policy(self, times: int) -> StatDict:
        """Improve the policy on previously collected environment data.

        Improves the policy on batches sampled from the replay buffer.

        Args:
            times: number of times to call :meth:`improve_policy`

        Returns:
            A dictionary of training statistics
        """
        for _ in range(times):
            batch = self.replay.sample(self.config["batch_size"])
            batch = self.lazy_tensor_dict(batch)
            info = self.improve_policy(batch)

        return info

    def timer_stats(self) -> dict:
        """Returns the timer statistics."""
        model_timer = self.timers["model"]
        policy_timer = self.timers["policy"]
        return dict(
            model_time_s=round(model_timer.mean, 3),
            policy_time_s=round(policy_timer.mean, 3),
            # Get mean number of model epochs per second spent updating the model
            model_update_throughput=round(model_timer.mean_throughput, 3),
            # Get mean number of policy updates per second spent updating the policy
            policy_update_throughput=round(policy_timer.mean_throughput, 3),
        )

    @staticmethod
    def add_options(policy_cls: type) -> type:
        """Decorator to add default model-based options used by MBPolicyMixin."""
        return model_based_options(policy_cls)
