from smac_plus.starcraft2.starcraft2 import StarCraft2Env
import numpy as np


def main():
    env = StarCraft2Env(
    continuing_episode=False,
    difficulty="7",
    #game_version=null,
    #map_name="bane_vs_hM",
    move_amount=2,
    #obs_all_health=True,
    obs_instead_of_state=False,
    obs_last_action=False,
    obs_own_health=True,
    obs_pathing_grid=False,
    obs_terrain_height=False,
    obs_timestep_number=False,
    reward_death_value=10,
    reward_defeat=0,
    reward_negative_scale=0.5,
    reward_only_positive=True,
    reward_scale=True,
    reward_scale_rate=20,
    reward_sparse=False,
    reward_win=200,
    replay_dir="",
    replay_prefix="",
    state_last_action=True,
    state_timestep_number=False,
    step_mul=8,
    debug=True,
    print_rew=False,
    is_print=False,
    print_steps=1000,
    map_name='1o_1r2h_vs_2r1b',#'1o_5b_vs_1h_plain', #'1o_5b_vs_3h',#'1o_8b_vs_2h', #1o_10b_vs_1r #MMM #5z_vs_1ul #1o_10b_vs_1r #bane_vs_bane
    sight_range=2,
    shoot_range=2,
    obs_all_health=False,
    obs_enemy_health=False)
    env_info = env.get_env_info()

    n_actions = env_info["n_actions"]
    n_agents = env_info["n_agents"]

    n_episodes = 1#10

    for e in range(n_episodes):
        env.reset()
        terminated = False
        episode_reward = 0
        t=0

        while not terminated:
            t+=1
            obs = env.get_obs()
            state = env.get_state()
            # env.render()  # Uncomment for rendering

            actions = []
            for agent_id in range(n_agents):
                avail_actions = env.get_avail_agent_actions(agent_id)
                avail_actions_ind = np.nonzero(avail_actions)[0]
                action = np.random.choice(avail_actions_ind)
                actions.append(action)

            reward, terminated, _ = env.step(actions)
            episode_reward += reward
            break

        print('total steps',t)
        print("Total reward in episode {} = {}".format(e, episode_reward))

    env.close()

main()