import pandas as pd
from dmc_gen.config import Args
from params_proto.neo_hyper import Sweep
from model_free_analysis.baselines import RUN
import numpy as np

import sys; sys.path.append('.')
from sweep_helper import set_args

soda_envs = ['Walker-walk', 'Walker-stand', 'Cartpole-swingup', 'Ball_in_cup-catch', 'Finger-spin']
extra_envs = ['Reacher-easy', 'Cheetah-run', 'Cartpole-balance', 'Finger-turn_easy']
# extra_envs = ['Reacher-easy', 'Cheetah-run', 'Cartpole-balance', 'Hopper-hop']

envs = soda_envs + extra_envs

with Sweep(RUN, Args) as sweep:
    with sweep.product:
        Args.algorithm = ['sac', 'soda', 'pad', 'svea']
        # Args.algorithm = ['soda', 'svea', 'pad']
        Args.seed = [(i + 1) * 100 for i in range(5)]

    with sweep.zip:
        Args.env_name = [f'dmc:{env_name}-v1' for env_name in envs]
        # Args.eval_env = [f'dmc:{env_name}-v1' for env_name in envs]
        Args.eval_env_name = [f'dmc:{env_name}-v1' for env_name in envs]


@sweep.each
def tail(RUN, Args: Args):
    RUN.job_name = "debug/" if RUN.debug else ""
    env_name = Args.env_name.split(':')[1][:-3].lower()
    RUN.job_name += f"{Args.algorithm}/{env_name}/{Args.seed}"
    Args.checkpoint_freq = 50_000

    set_args(Args.algorithm, Args.env_name, Args)


if RUN.debug:
    sweep.save("debug_sweep.jsonl")
else:
    sweep.save("sweep.jsonl")
