# %%
import sys
import os
sys.path.append(os.path.join(os.getcwd(), '..', '..'))

import pickle
import jax
import wandb 
import tqdm 

from jaxmarl.environments.jaxnav import JaxNav


# %% [markdown]
# ### Constants

# %%
SEED = 0

NUM_ENVS_TO_SAMPLE = 2500
EVAL_ITERS = 4
TOTAL_ENVS = NUM_ENVS_TO_SAMPLE * EVAL_ITERS

N_PARALLEL = 5
ROLLOUT_LEN = 1000

PATH = f'sfl/data/eval/jaxnav/eval_multi_agent_{TOTAL_ENVS}e.pkl'
os.makedirs(os.path.dirname(PATH), exist_ok=True)
print('Total envs sampled: ', TOTAL_ENVS)

# %% [markdown]
# ### Run ID

# %%
rng = jax.random.PRNGKey(SEED)

api = wandb.Api()

runs = api.runs("alex-plus/multi_robot_ued",
                filters={"group": "MA-Learnability"})

print('runs', runs)
for run in runs:
    print(run.name, run.id)


first_config = runs[0].config
first_config["env"]["env_params"]['map_params']['valid_path_check'] = True

env = JaxNav(num_agents=first_config["env"]["num_agents"], 
                    **first_config["env"]["env_params"])

state_set = []
for _ in tqdm.tqdm(range(EVAL_ITERS)):
    rng, _rng = jax.random.split(rng)
    reset_rngs = jax.random.split(_rng, NUM_ENVS_TO_SAMPLE)
    _, states = jax.vmap(env.reset)(reset_rngs)
    state_set.append(states)


with open(PATH, 'wb') as f:
    pickle.dump(state_set, f)
    
print('saved to ', PATH)