#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
"""
The Trainer for Cross-Entropy Method. The idea is that an ensemble of
 world models are fitted to predict transitions and reward functions.
A cross entropy method-based planner will then plan the best next action
based on simulation data generated by the fitted world models.

The idea is inspired by: https://arxiv.org/abs/1805.12114
"""
import logging
from typing import List

import reagent.core.types as rlt
import torch.nn as nn
from reagent.core.parameters import CEMTrainerParameters
from reagent.models.cem_planner import CEMPlannerNetwork
from reagent.training.reagent_lightning_module import ReAgentLightningModule
from reagent.training.world_model.mdnrnn_trainer import MDNRNNTrainer


logger = logging.getLogger(__name__)


def print_mdnrnn_losses(minibatch, model_index, losses) -> None:
    logger.info(
        f"{minibatch}-th minibatch {model_index}-th model: \n"
        f'loss={losses["loss"]}, bce={losses["bce"]}, '
        f'gmm={losses["gmm"]}, mse={losses["mse"]}\n'
    )


class CEMTrainer(ReAgentLightningModule):
    def __init__(
        self,
        cem_planner_network: CEMPlannerNetwork,
        world_model_trainers: List[MDNRNNTrainer],
        parameters: CEMTrainerParameters,
    ) -> None:
        super().__init__()
        self.cem_planner_network = cem_planner_network
        self.world_model_trainers = nn.ModuleList(world_model_trainers)

    def configure_optimizers(self):
        return [o for t in self.world_model_trainers for o in t.configure_optimizers()]

    def train_step_gen(self, training_batch: rlt.MemoryNetworkInput, batch_idx: int):
        for t in self.world_model_trainers:
            yield from t.train_step_gen(training_batch, batch_idx)
