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

# %% auto 0
__all__ = ['experiment_stepwise', 'experiment_stepwise_no_log', 'test', 'experiment_wandb', 'Weighted_Average',
           '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, mdp_train, mdp_test, core_train, core_test, logger, n_episodes_test, n, log_wandb=True):

    print("testing")
    
    if mdp_test is not None:
        dataset = core_test.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
    else:
        dataset = core_train.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
    s, *_ = parse_dataset(dataset)

    J = np.mean(compute_J(dataset, mdp_train.info.gamma))
    R = np.mean(compute_J(dataset))
    
    print("check log wandb")

    try:
        E = agent.policy.entropy(s)
        if log_wandb:
            wandb.log({"test/J": J, "test/R": R, "test/entropy": E})
    except:
        E = None
        if log_wandb:
            wandb.log({"test/J": J, "test/R": R})
    print("wandb done")

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

    print("testing done")
    
    return J, R


def experiment_wandb(   agent,
                        mdp_train,
                        mdp_test,
                        run_id,

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

    """
    # TODO: update docstring
    """

    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 mdp_test is not None:
        mdp_test.reset(state=0)

    # Start training

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

    core_train = Core(agent, mdp_train)
    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_train.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
    s, *_ = parse_dataset(dataset)

    J = np.mean(compute_J(dataset, mdp_train.info.gamma))
    R = np.mean(compute_J(dataset))
    wandb.log({"test/J": J, "test/R": R})
    
    try:
        E = agent.policy.entropy(s)
        wandb.log({"test/entropy": E})
    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")

    artifact = wandb.Artifact("transitions_dataset", type='dataset')
    artifact.add_file(f"{path}/{agent.name}_{run_id}/dataset_epoch_0.pkl", name="dataset_epoch_0")

    try:   
        if agent.train_directly:
            if agent.train_mode == "direct":
                agent.fit(features=mdp_train.features, demand=mdp_train.demand)
                for i in trange(n_epochs):
                    if i==0:
                        J, R = test(agent, mdp_train, mdp_test, core_train, core_test, logger, n_episodes_test, i)
                    else:
                        wandb.log({"test/J": J, "test/R": R, "test/entropy": E})

            else:
                for i in trange(n_epochs):
                    agent.fit_epoch(mdp_train.features, mdp_train.demand)
                    agent.eval()
                    J, R = test(agent, mdp_train, mdp_test, core_train, core_test, logger, n_episodes_test, i)   
                    agent.train()
            
    except:
        ### Training
        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):
            core_train.learn(n_steps=n_steps, n_steps_per_fit=1, quiet = True)
            try:
                agent.policy.eval()
            except:
                pass
            if mdp_test is not None:
                dataset = core_test.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
            else:
                dataset = core_train.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
            s, *_ = parse_dataset(dataset)

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

            try:
                agent.policy.train()
            except:
                pass
            
            try:
                E = agent.policy.entropy(s)
                wandb.log({"test/J": J, "test/R": R, "test/entropy": E})
            except:
                E = None
                wandb.log({"test/J": J, "test/R": R})

            # 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+1) % dataset_log_freq == 0:
                    logger.log_dataset(dataset, name_addition = f"_epoch_{n+1}")
                    artifact.add_file(f"{path}/{agent.name}_{run_id}/dataset_epoch_{n+1}.pkl", name=f"dataset_epoch_{n+1}")
        
        ### Final evaluation
    
    try:
        agent.policy.eval()
    except:
        pass
    
    if mdp_test is not None:
        dataset = core_test.evaluate(n_episodes=n_episodes_test, render=False, quiet=True)
    else:
        dataset = core_train.evaluate(n_episodes=n_episodes_test, render=False, quiet=True)
    
    s, *_ = parse_dataset(dataset)

    J = np.mean(compute_J(dataset, mdp_train.info.gamma))
    R = np.mean(compute_J(dataset))
    wandb.log({"test/J": J, "test/R": R})
    try:
        E = agent.policy.entropy(s)
        wandb.log({"test/entropy": E})
    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")

    artifact.add_file(f"{path}/{agent.name}_{run_id}/dataset_final.pkl", name=f"dataset_final")

    wandb.log_artifact(artifact)
        
    return agent

# %% ../../nbs/experiment_functions/01_run_experiment.ipynb 9
class Weighted_Average:
    def __init__(self,
                 weight_final = 0.2,
                 start_tracking_epoch = 50, 
                 total_epochs = 200):
        self.weight_final = weight_final
        self.start_tracking_epoch = start_tracking_epoch
        self.totaol_epochs = total_epochs
        self.weight_datapoint = (1-weight_final)/(total_epochs-1-start_tracking_epoch)
        self.data = list()
        self.total_epochs = total_epochs

    def add_datapoint(self, datapoint, epoch):
        print("inside adding function")
        print("epoch", epoch)
        print("start_tracking_epoch", self.start_tracking_epoch)
        if epoch >= self.start_tracking_epoch:
            print("starting adding")
            print(epoch)
            print(self.total_epochs)
            if epoch < self.total_epochs-1:
                print("adding datapoint with weight")
                print("datapoing", datapoint)
                print("weight_datapoint", self.weight_datapoint)
                print("datapoint*weight_datapoint", datapoint*self.weight_datapoint)
                self.data.append(datapoint*self.weight_datapoint)
            elif epoch == self.total_epochs-1:
                print("adding datapoint with final weight")
                self.data.append(datapoint*self.weight_final)
            else:
                pass
        else:
            print("not adding datapoint")
        print("did")
    
    def get_result(self):
        print(self.data)
        return np.sum(self.data)

def experiment_wandb_optimize(   agent_l,
                        mdp_train_l,
                        mdp_val_l,
                        run_id,

                        weight_final,
                        start_tracking_epoch,

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

    """
    # TODO: update docstring
    """


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

    weighted_j_list = list()

    print("starting loop")

    print("len mdp_train_l", len(mdp_train_l))

    for dataset_num in range(len(mdp_train_l)):

        print("dataset_num", dataset_num)


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

        print("setting up tracker")
        j_tracker = Weighted_Average(weight_final = weight_final, start_tracking_epoch = start_tracking_epoch, total_epochs = n_epochs)

        # Set-up logger
        print("setting 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)

        # Start training

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

        print("going into training loop")

        try:   
            if agent.train_directly:
                if agent.train_mode == "direct":
                    agent.fit(features=mdp_train.features, demand=mdp_train.demand)
                    for i in trange(n_epochs):
                        if i==0:
                            J, R = test(agent, mdp_train, mdp_val, core_train, core_val, logger, n_episodes_test, i, log_wandb=False)
                        wandb.log({f"val/J_{dataset_num}": J, f"val/R_{dataset_num}": R})
                    print("did training of one model")

                else:
                    for i in trange(n_epochs):
                        agent.fit_epoch(mdp_train.features, mdp_train.demand)
                        print("done fitting")
                        J, R = test(agent, mdp_train, mdp_val, core_train, core_val, logger, n_episodes_test, i, log_wandb=False)
                        print("done testing")
                        wandb.log({f"val/J_{dataset_num}": J, f"val/R_{dataset_num}": R}) 
                        print(J)
                        j_tracker.add_datapoint(J, i)
                        print("added datapoint")
                
        except:
            ### Training
            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):
                core_train.learn(n_steps=n_steps, n_steps_per_fit=1, quiet = True)
                try:
                    agent.policy.eval()
                except:
                    pass
                dataset = core_val.evaluate(n_episodes=n_episodes_test, render=False, quiet = True)
                
                s, *_ = parse_dataset(dataset)

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

                j_tracker.add_datapoint(J, n)

                try:
                    agent.policy.train()
                except:
                    pass

                # logger.epoch_info(epoch = n+1, J=J, R=R, entropy=E)
                logger.log_numpy(epoch = n+1, J=J, R=R)
                wandb.log({f"val/J_{dataset_num}": J, f"val/R_{dataset_num}": R})
            
        weighted_j = j_tracker.get_result()
        print("weighted_j", weighted_j)

        weighted_j_list.append(weighted_j)
        

        
    ### Final evaluation
    print("ealculating final eval")
    print(weighted_j_list)
    weighed_j_total = np.mean(weighted_j_list)

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