# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/experiment_functions/01_run_experiment.ipynb.

# %% auto 0
__all__ = ['experiment_stepwise', 'experiment_stepwise_no_log']

# %% ../../nbs/experiment_functions/01_run_experiment.ipynb 4
# MushroomRL:
from mushroom_rl.core import Core, Logger
from mushroom_rl.utils.dataset import compute_J,  parse_dataset

# Pytorch
import torch

# General libraries:
import numpy as np
import pickle
import os
from tqdm import tqdm, trange

# %% ../../nbs/experiment_functions/01_run_experiment.ipynb 6
def experiment_stepwise(mdp,
                        run_params,
                        n_epochs,
                        n_steps,
                        n_episodes_test,
                        rand_seed = None,
                        logger_dir = "./results/loggers",
                        results_dir = "./results",
                        dataset_log_freq = None):

    """
    TODO 1: Add function to (deep)save the entire agent to be able to reloead it later and potentially continue training.
    TODO 2: Provide functionalita to turn of exploration during evaluation (e.g., epsilon=0 for DQN, sample_mean for SAC)

    Function to run an experiment with a given agent and mdp.
    Designed specifically for algorithms that learn on each step (e.g., SAC).

    Agents that have the fit function directly on data and are not trained
    stepwise via the Core class of Mushroom must have the attribute train_directly = True.

    The function assumes that that the mpd has a a demand attribute containing
    a list of historical demands. This is used as input for the directly trained agents.

    It will save the logger and save the logger containing the path. From the path it is 
    then possible to retrieve the results of J and the dataset.

    Args:
        mdp (object): MDP to be solved.
        run_params (dict): Dictionary with the agent and run number.
        n_epochs (int): Number of epochs to train the agent.
        n_steps (int): Number of steps per epoch.
        n_episodes_test (int): Number of episodes to evaluate the agent.
    """

    agent = run_params['agent']
    run = run_params['run']

    if rand_seed is not None:
        np.random.seed(rand_seed)
        torch.manual_seed(rand_seed)

    # Set-up logger
    logger = Logger(agent.name+f"_{run}", results_dir=f"{results_dir}/run_{run}")
    filename = f"{logger_dir}/logger_{agent.name}_{run}.pkl"
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    pickle.dump(logger, open(f"{logger_dir}/logger_{agent.name}_{run}.pkl", "wb"))

    # ensure mdp starts from the same state
    mdp.reset()

    # Start training

    logger.strong_line()
    logger.info('Experiment Algorithm: ' + agent.name)

    try:
        if agent.train_directly:
            agent.fit(mdp.demand)
            core = Core(agent, mdp)

    
    except:
   
        core = Core(agent, mdp)

        ### Initial evaluation
        dataset = core.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
        s, *_ = parse_dataset(dataset)

        J = np.mean(compute_J(dataset, mdp.info.gamma))
        R = np.mean(compute_J(dataset))
        try:
            E = agent.policy.entropy(s)
        except:
            E = None

        logger.epoch_info(epoch=0, J=J, R=R, entropy=E)
        logger.log_numpy(epoch=0, J=J, R=R, entropy=E)
        logger.log_dataset(dataset, name_addition = "_epoch_0")

        ### Training
        initial_replay_size = agent._replay_memory._initial_size
        core.learn(n_steps=initial_replay_size, n_steps_per_fit=initial_replay_size, quiet=True)

        for n in trange(n_epochs, leave=False):
            core.learn(n_steps=n_steps, n_steps_per_fit=1, quiet = True)
            dataset = core.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
            s, *_ = parse_dataset(dataset)

            J = np.mean(compute_J(dataset, mdp.info.gamma))
            R = np.mean(compute_J(dataset))
            try:
                E = agent.policy.entropy(s)
            except:
                E = None

            # logger.epoch_info(epoch = n+1, J=J, R=R, entropy=E)
            logger.log_numpy(epoch = n+1, J=J, R=R, entropy=E)

            if dataset_log_freq is not None:
                if n % dataset_log_freq == 0:
                    logger.log_dataset(dataset, name_addition = f"_epoch_{n+1}")
        
        ### Final evaluation
    
    dataset = core.evaluate(n_episodes=n_episodes_test, render=False, quiet=True)
    
    s, *_ = parse_dataset(dataset)

    J = np.mean(compute_J(dataset, mdp.info.gamma))
    R = np.mean(compute_J(dataset))
    try:
        E = agent.policy.entropy(s)
    except:
        E = None

    logger.epoch_info(epoch = "final", J=J, R=R, E=E)
    logger.log_numpy(J=J, R=R, E=E)
    logger.log_dataset(dataset, name_addition = "_final")
        
    return logger

# %% ../../nbs/experiment_functions/01_run_experiment.ipynb 7
def experiment_stepwise_no_log(mdp,
                        agent,
                        n_epochs,
                        n_steps,
                        n_episodes_test,
                        rand_seed = None,
                        dataset_log_freq = None):

    """
    TODO 1: Add function to (deep)save the entire agent to be able to reloead it later and potentially continue training.
    TODO 2: Provide functionalita to turn of exploration during evaluation (e.g., epsilon=0 for DQN, sample_mean for SAC)

    Function to run an experiment with a given agent and mdp.
    Designed specifically for algorithms that learn on each step (e.g., SAC) to tune hyperparameters.
    Does not use logger but simply returns the results.

    The function assumes that that the mpd has a a demand attribute containing
    a list of historical demands. This is used as input for the directly trained agents.

    Args:
        mdp (object): MDP to be solved.
        agent (object): Agent to solve the MDP.
        n_epochs (int): Number of epochs to train the agent.
        n_steps (int): Number of steps per epoch.
        n_episodes_test (int): Number of episodes to evaluate the agent.
    """

    if rand_seed is not None:
        np.random.seed(rand_seed)
        torch.manual_seed(rand_seed)

    # ensure mdp starts from the same state
    mdp.reset()

    core = Core(agent, mdp)

    ### Training
    initial_replay_size = agent._replay_memory._initial_size
    core.learn(n_steps=initial_replay_size, n_steps_per_fit=initial_replay_size, quiet=True)

    R_hist = []

    for n in trange(n_epochs, leave=False):
        core.learn(n_steps=n_steps, n_steps_per_fit=1, quiet = True)
        dataset = core.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)

        R = np.mean(compute_J(dataset))

        R_hist.append(R)
        
    return R_hist
