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

# %% auto 0
__all__ = ['experiment_stepwise', 'experiment_stepwise_no_log', 'test_agent', 'save_networks', 'load_networks', 'check_val_test',
           'experiment_wandb', 'experiment_wandb_optimize']

# %% ../../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
import wandb

# %% ../../nbs/experiment_functions/01_run_experiment.ipynb 6
def experiment_stepwise(mdp,
                        run_params,
                        n_epochs,
                        n_steps,
                        n_episodes_test,
                        mdp_test = None,
                        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 functionality 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()

    if mdp_test is not None:
        mdp_test.reset(state=0)

    # Start training

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

    core = Core(agent, mdp)
    if mdp_test is not None:
        core_test = Core(agent, mdp_test)

    try:
        if agent.train_directly:
            agent.fit(mdp.demand)

    except:
   
        if mdp_test is not None:
            core_test = Core(agent, mdp_test)
        

        ### Initial evaluation
        if mdp_test is not None:
            dataset = core_test.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
        else:
            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)
            if mdp_test is not None:
                dataset = core_test.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
            else:
                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
    
    if mdp_test is not None:
        dataset = core_test.evaluate(n_episodes=n_episodes_test, render=False, quiet=True)
    else:
        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,
                        mdp_test = None,
                        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
    if mdp_test is not None:
        mdp_test.reset()
    else:
        mdp.reset()

    core = Core(agent, mdp)
    if mdp_test is not None:
        core_test = Core(agent, mdp_test)

    ### 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)
        if mdp_test is not None:
            dataset = core_test.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
        else:
            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

# %% ../../nbs/experiment_functions/01_run_experiment.ipynb 8
def test_agent(agent, mdp, core, logger, artifact, path, run_id, n_episodes_test, n, log_wandb=True, val=False, return_J_R=False, dataset_log_freq = None):

    # TODO: Make this into a class to avoid passing the same arguments over and over again

    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))

    if hasattr(agent.policy, 'entropy'):
        E = agent.policy.entropy(s)

        if log_wandb:
            if val:
                wandb.log({"val/J": J, "val/R": R, "val/entropy": E})
            else:
                wandb.log({"test/J": J, "test/R": R, "test/entropy": E})
    
    else:
        E = None
        if log_wandb:
            if val:
                wandb.log({"val/J": J, "val/R": R})
            else:
                wandb.log({"test/J": J, "test/R": R})

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

    if dataset_log_freq is not None:
        # check if n is string:
        if isinstance(n, str):
            logger.log_dataset(dataset, name_addition = f"_{n}")
            if log_wandb:
                artifact.add_file(f"{path}/{agent.name}_{run_id}/dataset_{n}.pkl", name=f"dataset_{n}")

        elif (n) % dataset_log_freq == 0:
            logger.log_dataset(dataset, name_addition = f"_epoch_{n}")
            if log_wandb:
                print("adding dataset to artifact")
                artifact.add_file(f"{path}/{agent.name}_{run_id}/dataset_epoch_{n}.pkl", name=f"dataset_epoch_{n}")

    if return_J_R:
        return J, R
    
def save_networks(agent, path, run_id):

    # check class of agent

    agent_class = agent.agent.__class__.__name__

    if agent_class == "SAC" or agent_class == "TD3":

        ensemble_critic = agent._critic_approximator._impl.model
        networks = []

        for i, model in enumerate(ensemble_critic):
            networks.append(model.network)

    if agent_class == "SAC":
        networks.append(agent.policy._mu_approximator._impl.model.network)
        networks.append(agent.policy._sigma_approximator._impl.model.network)
    elif agent_class == "TD3":
        networks.append(agent.policy._approximator._impl.model.network)

    path_full = f"{path}/parameters"
    os.makedirs(path_full, exist_ok=True)

    for i, network in enumerate(networks):
        network_path = "{}/parameters/network_{}_run_{}.pt".format(path, i, run_id)
        torch.save(network.state_dict(), network_path)

def load_networks(agent, path, run_id):

    agent_class = agent.agent.__class__.__name__

    if agent_class == "SAC" or agent_class == "TD3":

        ensemble_critic = agent._critic_approximator._impl.model
        networks = []

        for i, model in enumerate(ensemble_critic):
            networks.append(model.network)
    
    if agent_class == "SAC":
        networks.append(agent.policy._mu_approximator._impl.model.network)
        networks.append(agent.policy._sigma_approximator._impl.model.network)   
    
    elif agent_class == "TD3":
        networks.append(agent.policy._approximator._impl.model.network)

    for i, network in enumerate(networks):
        network_path = "{}/parameters/network_{}_run_{}.pt".format(path, i, run_id)
        network.load_state_dict(torch.load(network_path))

    # for i, network in enumerate(ensemble_critic):
    #     network_path = "{}/parameters/network_{}_run_{}.pt".format(path, i, run_id)
    #     network.network.load_state_dict(torch.load(network_path))

    # _mu_approximator = agent.policy._mu_approximator._impl.model.network
    # _sigma_approximator = agent.policy._sigma_approximator._impl.model.network

    # _mu_approximator_path = "{}/parameters/network_{}_run_{}.pt".format(path, i+1, run_id)
    # _sigma_approximator_path = "{}/parameters/network_{}_run_{}.pt".format(path, i+2, run_id)

    # _mu_approximator.load_state_dict(torch.load(_mu_approximator_path))
    # _sigma_approximator.load_state_dict(torch.load(_sigma_approximator_path))

def check_val_test(mdp_val, mdp_test):
    if mdp_val is not None:
        print("validation set provided")
        val_provided = True
    else:
        print("no validation set provided")
        val_provided = False

    if mdp_test is not None:
        print("test set provided")
        test_provided = True
    else:
        print("no test set provided")
        test_provided = False
    
    return val_provided, test_provided


def experiment_wandb(   agent,
                        mdp_train,
                        mdp_test,
                        run_id,

                        n_epochs,
                        n_steps,
                        n_episodes_test,
                        mdp_val = None,
                        rand_seed = None,
                        dataset_log_freq = None,
                        results_dir = "results",

                        early_stopping_patience = None,
                        early_stopping_warmup = 100,
                        save_best=True,
                        
                        ):

    """
    # TODO: update docstring
    """

    # check if validation and test sets are provided
    val_provided, test_provided = check_val_test(mdp_val, mdp_test)

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

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

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

    if test_provided:
        mdp_test.reset(state=0)

    if val_provided:
        mdp_val.reset(state=0)
    else:
        mdp_val = mdp_test # Same validation as test set
        raise Warning("No validation set provided. Using test set as validation set.")

    # Start training

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

    core_train = Core(agent, mdp_train)

    if test_provided:
        core_test = Core(agent, mdp_test)
        core_val = Core(agent, mdp_val) # will use same as test set if no validation set provided

    artifact = wandb.Artifact("transitions_dataset", type='dataset')

    ### Initial evaluation
    if test_provided:
        # will be implicitly test set, if no validation set provided
        test_agent(agent, mdp_val, core_val, logger, artifact, path, run_id, n_episodes_test, 0, log_wandb=True, val = val_provided, dataset_log_freq = dataset_log_freq)
    else:
        test_agent(agent, mdp_train, core_train, logger, artifact, path, run_id, n_episodes_test, 0, log_wandb=True, dataset_log_freq = dataset_log_freq) 
    
    stopped_early=False
    J_history = []

    ### Training loop outside mushroom_rl env:
    if hasattr(agent, 'train_directly') and agent.train_directly:
        if agent.train_mode == "direct":
            agent.fit(features=mdp_train.features, demand=mdp_train.demand)
            for n in trange(n_epochs):
                if n==0:
                    if test_provided:
                        # will be implicitly test set, if no validation set provided
                        J, R = test_agent(agent, mdp_val, core_val, logger, artifact, path, run_id, n_episodes_test, 1, log_wandb=True, val = val_provided, return_J_R=True, dataset_log_freq = dataset_log_freq)
                    else:
                        J, R = test_agent(agent, mdp_train, core_train, logger, artifact, path, run_id, n_episodes_test, 1, log_wandb=True, return_J_R=True, dataset_log_freq = dataset_log_freq) 
                else:
                    if val_provided:
                        wandb.log({"val/J": J, "val/R": R})
                    else:
                        wandb.log({"test/J": J, "test/R": R})

        elif agent.train_mode == "epochs":
            for n in trange(n_epochs):

                if stopped_early:
                    wandb.log({"val/J": J, "val/R": R})
                    continue

                agent.fit_epoch(mdp_train.features, mdp_train.demand)
                
                agent.eval()
                
                if test_provided:
                    # will be implicitly test set, if no validation set provided
                    J, R = test_agent(agent, mdp_val, core_val, logger, artifact, path, run_id, n_episodes_test, n+1, log_wandb=True, val = val_provided, dataset_log_freq = dataset_log_freq, return_J_R=True)
                else:
                    J, R = test_agent(agent, mdp_train, core_train, logger, artifact, path, run_id, n_episodes_test, n+1, log_wandb=True, dataset_log_freq = dataset_log_freq, return_J_R=True)   
                
                agent.train()
                
                J_history.append(J)

                if early_stopping_patience is not None:
                    if n > early_stopping_warmup and n > 2*early_stopping_patience:
                        if np.mean(J_history[-early_stopping_patience:]) <= np.mean(J_history[-2*early_stopping_patience:-early_stopping_patience]):
                            print("early stopping")
                            stopped_early = True
        else:
            raise ValueError("train_mode not recognized")
            
    else:
        ### Training
        initial_replay_size = agent._replay_memory._initial_size

        if hasattr(agent.policy, 'train'):
            agent.policy.train()

        core_train.learn(n_steps=initial_replay_size, n_steps_per_fit=initial_replay_size, quiet=True)

        for n in trange(n_epochs, leave=False):

            if stopped_early:
                wandb.log({"val/J": J, "val/R": R})
                continue

            core_train.learn(n_steps=n_steps, n_steps_per_fit=1, quiet = True)
            if hasattr(agent.policy, 'eval'):
                agent.policy.eval()
            if test_provided:
                # will be implicitly test set, if no validation set provided
                J, R = test_agent(agent, mdp_val, core_val, logger, artifact, path, run_id, n_episodes_test, n+1, log_wandb=True, val = val_provided, return_J_R=True, dataset_log_freq = dataset_log_freq)
            else:
                J, R = test_agent(agent, mdp_train, core_train, logger, artifact, path, run_id, n_episodes_test, n+1, log_wandb=True, return_J_R=True, dataset_log_freq = dataset_log_freq)  

            J_history.append(J)
            
            if hasattr(agent.policy, 'train'):
                agent.policy.train()
            
            if save_best:
                if J >= np.max(J_history):
                    save_networks(agent, path, run_id)

            if early_stopping_patience is not None:
                if n > early_stopping_warmup and n > 2*early_stopping_patience:
                    if np.mean(J_history[-early_stopping_patience:]) <= np.mean(J_history[-2*early_stopping_patience:-early_stopping_patience]):
                        print("early stopping")
                        stopped_early = True
            
        ### Final evaluation

    print("doing final evaluation with last model on test set")
    if hasattr(agent.policy, 'eval'):
        agent.policy.eval()

    if test_provided:
        # will always use test set
        test_agent(agent, mdp_test, core_test, logger, artifact, path, run_id, n_episodes_test, "test_last", log_wandb=True, val = False, return_J_R=True, dataset_log_freq = dataset_log_freq)
    else:
        test_agent(agent, mdp_train, core_train, logger, artifact, path, run_id, n_episodes_test, "train_last", log_wandb=True, return_J_R=True, dataset_log_freq = dataset_log_freq)   

    if save_best and not hasattr(agent, 'train_directly'):
        print("doing final evaluation with best model on test set")
        # load saved networks
        
        load_networks(agent, path, run_id)

        if hasattr(agent.policy, 'eval'):
            agent.policy.eval()
    
        if test_provided:
            # will always use test set
            test_agent(agent, mdp_test, core_test, logger, artifact, path, run_id, n_episodes_test, "test_best", log_wandb=True, val = False, return_J_R=True, dataset_log_freq = dataset_log_freq)
        else:
            test_agent(agent, mdp_train, core_train, logger, artifact, path, run_id, n_episodes_test, "train_best", log_wandb=True, return_J_R=True, dataset_log_freq = dataset_log_freq) 
     
    wandb.log_artifact(artifact)
 
    return agent

# %% ../../nbs/experiment_functions/01_run_experiment.ipynb 9
def experiment_wandb_optimize(   agent_l,
                        mdp_train_l,
                        mdp_val_l,
                        run_id,

                        weight_final,
                        weight_best,
                        eval_period,

                        n_epochs,
                        n_steps,
                        n_episodes_test,
                        rand_seed = None,
                        dataset_log_freq = None,
                        results_dir = "results",

                        early_stopping_patience = None,
                        early_stopping_warmup = 100,
                        save_best=False,
                        ):

    """
    # TODO: update docstring
    """

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

    weighted_j_list = list()

    for dataset_num in range(len(mdp_train_l)):

        mdp_train = mdp_train_l[dataset_num]
        mdp_val = mdp_val_l[dataset_num]
        agent = agent_l[dataset_num]

        J_history = list()
        stopped_early = False

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

        # ensure mdp starts from the same state
        mdp_train.reset()
        mdp_val.reset(state=0)

        artifact = None # no dataset tracking during hyperparameteroptimization

        # Start training

        logger.strong_line()
        logger.info('Experiment Algorithm: ' + agent.name + f" on dataset {dataset_num}")

        core_train = Core(agent, mdp_train)
        core_val = Core(agent, mdp_val)

        if hasattr(agent, 'train_directly') and agent.train_directly:
            if agent.train_mode == "direct":
                agent.fit(features=mdp_train.features, demand=mdp_train.demand)
                for n in trange(n_epochs):
                    if n==0:
                        J, R = test_agent(agent, mdp_val, core_val, logger, artifact, path, run_id, n_episodes_test, 1, log_wandb=True, val = True, return_J_R=True, dataset_log_freq = None)
                    else:
                        wandb.log({"val/J": J, "val/R": R})
                    J_history.append(J)

            else:
                for n in trange(n_epochs):
                    if stopped_early:
                        wandb.log({"val/J": J, "val/R": R})
                        continue
                    agent.fit_epoch(mdp_train.features, mdp_train.demand)
                    agent.eval()
                    J, R = test_agent(agent, mdp_val, core_val, logger, artifact, path, run_id, n_episodes_test, n+1, log_wandb=True, val = True, return_J_R=True, dataset_log_freq = None)
                    J_history.append(J)

                    if early_stopping_patience is not None:
                        if n > early_stopping_warmup and n > 2*early_stopping_patience:
                            if np.mean(J_history[-early_stopping_patience:]) <= np.mean(J_history[-2*early_stopping_patience:-early_stopping_patience]):
                                print("early stopping")
                                stopped_early = True

                    agent.train()   
                
        else:
            ### Training
            if hasattr(agent.policy, 'train'):
                agent.policy.train()

            initial_replay_size = agent._replay_memory._initial_size
            core_train.learn(n_steps=initial_replay_size, n_steps_per_fit=initial_replay_size, quiet=True)

            
            for n in trange(n_epochs, leave=False):

                if stopped_early:
                    wandb.log({"val/J": J, "val/R": R})
                    continue
                    
                core_train.learn(n_steps=n_steps, n_steps_per_fit=1, quiet = True)
                
                if hasattr(agent.policy, 'eval'):
                    agent.policy.eval()
                    
                J, R = test_agent(agent, mdp_val, core_val, logger, artifact, path, run_id, n_episodes_test, n+1, log_wandb=True, val = True, return_J_R=True, dataset_log_freq = None)
                
                J_history.append(J)

                if hasattr(agent.policy, 'train'):
                    agent.policy.train()


                if early_stopping_patience is not None:
                    if n > early_stopping_warmup and n > 2*early_stopping_patience:
                        if np.mean(J_history[-early_stopping_patience:]) <= np.mean(J_history[-2*early_stopping_patience:-early_stopping_patience]):
                            print("early stopping")
                            stopped_early = True

        
        final_model_res = weight_final * J_history[-1]
        best_model_res = weight_best * np.max(J_history)
        eval_period_res = (1-weight_final-weight_best) * np.mean(J_history[-eval_period:])

        # print("evaluating")
        # print(J_history[-1])
        # print(np.max(J_history))
        # print(np.mean(J_history[-eval_period:]))

        weighted_j = final_model_res + best_model_res + eval_period_res
        weighted_j_list.append(weighted_j)
        
    # ### Final evaluation
    # print("ealculating final eval")
    # print(weighted_j_list)
    weighted_j_total = np.mean(weighted_j_list)

    # print("logging final eval to wandeb")
    wandb.log({f"val/J_weighted": weighted_j_total})
    
    return agent
