"""
A very basic arena for showing off our requirements from the 2D env.
"""
import sys
import numpy as np
import os
from tqdm import tqdm
import argparse
from PIL import Image

from bringbackshapes.twod_playground.arena import Arena
from moviepy.editor import ImageSequenceClip


def main(args):
    arena = Arena(
        render_game=args.render,
        add_brownian_distractor=args.distractor,)
    once = False
    obss = []
    pbar = tqdm()
    time_limit = args.time_limit
    step = 0
    while arena.running:
        arena.check_running_status()
        if arena.add_brownian_distractor:
            assert(arena.add_brownian_distractor == args.distractor)
            arena.move_distractor()
        if not args.overide:
            arena.apply_action()
        else:
            arena.apply_action((0, 0))
        arena.update()
        obs = arena.get_observation()
        obss.append(obs)
        if not once:
            obs = np.array(obs)
            print(obs.shape)
            print(obs.dtype)
            print(obs.min(), obs.max())
            once = True
        pbar.update(1)
        step += 1
        if time_limit and step > time_limit:
            break
    print(arena.get_states(state_type="relative"))
    arena.close()
    pbar.close()
    dir_name = "demos/cache/demo_2d_playground/"
    if not os.path.exists(dir_name):
        os.makedirs(dir_name, exist_ok=True)
    save_f = 100
    for i in range(len(obss)):
        if i % save_f == 0:
            im = Image.fromarray(obss[i])
            im.save(os.path.join(dir_name, f"{i}.jpg"))
    clip = ImageSequenceClip(obss[::save_f], fps=60)
    clip.write_gif(
        os.path.join(dir_name, "pygame_obs.gif"), fps=60, logger=None
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--time_limit", type=int, default=None)
    parser.add_argument("--overide", type=int, default=0)
    parser.add_argument("--distractor", action='store_true')
    parser.add_argument("--render", action='store_true')
    args = parser.parse_args()
    sys.exit(main(args))
