# 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
import sys,traceback

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

        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

    print("testing agent")
    got_nan = False

    dataset, dataset_info = core.evaluate(n_episodes=n_episodes_test, render=False, quiet = True, get_env_info=True)
    info_keys = dataset_info.keys()

    s, *_ = parse_dataset(dataset)

    J = np.mean(compute_J(dataset, mdp.info.gamma))
    R = np.mean(compute_J(dataset))

    if np.isnan(J):
        J = -1e6
        got_nan = True

    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:

        for i in range(len(dataset)):
            step_info = dict()

            for key in info_keys:
                step_info[key] = dataset_info[key][i]

            dataset[i] = dataset[i] + (step_info,)

        # 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, got_nan
    
    else:
        return got_nan
    
def save_networks(agent, path, run_id):

    # check class of agent

    if hasattr(agent, "agent"):
        agent_class = agent.agent.__class__.__name__
    else:
        agent_class = agent.__class__.__name__

    networks = []

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

        ensemble_critic = agent._critic_approximator._impl.model

        for i, model in enumerate(ensemble_critic):
            networks.append(model.network)
    
    elif agent_class == "PPO":
        critic = agent._V._impl.model
        networks.append(critic.network)
        networks.append(agent.policy._mu._impl.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)
    
    if agent_class == "MLPsgdAgent":
        networks.append(agent.model)
 
    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):

    if hasattr(agent, "agent"):
        agent_class = agent.agent.__class__.__name__
    else:
        agent_class = agent.__class__.__name__

    networks = []

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

        ensemble_critic = agent._critic_approximator._impl.model

        for i, model in enumerate(ensemble_critic):
            networks.append(model.network)
    
    elif agent_class == "PPO":
        critic = agent._V._impl.model
        networks.append(critic.network)
        networks.append(agent.policy._mu._impl.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)
    
    if agent_class == "MLPsgdAgent":
        networks.append(agent.model)

    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))
        print("loaded parameters")

    # 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,

                        num_reruns_nan = 10,

                        create_agent_function = None,
                        
                        ):

    """
    # TODO: update docstring
    """

    print("starting run experiment function")

    # if agnent has attribut "n_steps_per_fit" set n_steps_per_fit
    if hasattr(agent, "n_steps_per_fit"):
        n_steps_per_fit = agent.n_steps_per_fit
    else:
        n_steps_per_fit = 1

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

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

    
    num_runs = 0

    while num_runs <= num_reruns_nan:

        if num_runs > 0:
            agent, _ = create_agent_function[0](*create_agent_function[1])

        num_runs += 1
        nan = False

        # 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

        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
            nan = 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:
            nan = 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) 

        if nan:
            print("got nan, rerunning with num_re-runs = ", num_runs)
            continue

        
        stopped_early=False
        J_history = []

        ### Training loop outside mushroom_rl env:

        print("starting training")
        if hasattr(agent, 'train_directly') and agent.train_directly:
            if agent.train_mode == "direct":
                print("fitting")
                agent.fit(features=mdp_train.features, demand=mdp_train.demand)
                print("fitted")
                if not hasattr(agent, 'skip_val') or not agent.skip_val:
                    for n in trange(n_epochs):
                        if n==0:
                            if test_provided:
                                # will be implicitly test set, if no validation set provided
                                J, R, nan = 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, nan  = 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, nan = 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, nan = 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 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
            else:
                raise ValueError("train_mode not recognized")
                
        else:
            print("starting training")
            ### Training
            if hasattr(agent.policy, 'train'):
                agent.policy.train()
                print("setting train flag")

            if hasattr(agent, "_replay_memory"):
                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

                try:

                    core_train.learn(n_steps=n_steps, n_steps_per_fit=n_steps_per_fit, 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, nan = 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, nan = 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)  

                except Exception as e:

                            # exit gracefully, so wandb logs the problem
                    print(traceback.print_exc(), file=sys.stderr)
                    exit(1)
                    
                if nan:
                    break

                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
        
        if nan:
            print("got nan, rerunning with num_re-runs = ", num_runs)
            continue

        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
            nan = 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=False, dataset_log_freq = dataset_log_freq)
        else:
            nan = test_agent(agent, mdp_train, core_train, logger, artifact, path, run_id, n_episodes_test, "train_last", log_wandb=True, return_J_R=False, dataset_log_freq = dataset_log_freq)   

        if save_best and (not hasattr(agent, 'train_directly') or agent.train_mode == "epochs"):
            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
                nan = 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=False, dataset_log_freq = dataset_log_freq)
            else:
                nan = test_agent(agent, mdp_train, core_train, logger, artifact, path, run_id, n_episodes_test, "train_best", log_wandb=True, return_J_R=False, dataset_log_freq = dataset_log_freq) 
        
        if nan:
            print("got nan, rerunning with num_re-runs = ", num_runs)
        else:
            break
        
    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,

                        num_reruns_nan = 10,

                        create_agent_function = None,
                        
                        ):

    """
    # TODO: update docstring
    """

    print("in experiment function")
    if hasattr(agent_l[0], "n_steps_per_fit"):
        n_steps_per_fit = agent_l[0].n_steps_per_fit
    else:
        n_steps_per_fit = 1

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

        # print("in dataset loop")

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

        # print("got agent and mdp")

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

        # print("created logger")

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

        num_runs = 0

        print("starting training")


        while num_runs <= num_reruns_nan:
            
            
            nan = False

            num_runs += 1

            if num_runs > 0:
                agent, _ = create_agent_function[0](mdp_train, *create_agent_function[1][1:])

            J_history = list()
            stopped_early = False

            # Set-up logger


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

            artifact = None # no dataset tracking during hyperparameteroptimization

        # Start training


            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, nan = 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, nan = 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'):
                    print("setting train flag")
                    agent.policy.train()


                if hasattr(agent, "_replay_memory"):
                    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=n_steps_per_fit, quiet = True)
                    if hasattr(agent.policy, 'eval'):
                        agent.policy.eval()
                        
                    J, R, nan = 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 nan:
                        break


                    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
            
            if nan:
                print("got nan, rerunning with num_re-runs = ", num_runs)
                continue
                
            else:
                break

        
        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
