import argparse
import pathlib
import time

import numpy as np

import crafter


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('--seed', type=int, default=None)
  parser.add_argument('--area', nargs=2, type=int, default=(64, 64))
  parser.add_argument('--length', type=int, default=10000)
  parser.add_argument('--health', type=int, default=9)
  parser.add_argument('--record', type=pathlib.Path, default=None)
  args = parser.parse_args()

  random = np.random.RandomState(args.seed)
  crafter.constants.items['health']['max'] = args.health
  crafter.constants.items['health']['initial'] = args.health
  env = crafter.Env(area=args.area, length=args.length, seed=args.seed)
  if args.record:
    env = crafter.Recorder(env)

  start = time.time()
  obs = env.reset()
  print(f'Reset time: {1000*(time.time()-start):.2f}ms')
  print('Coal exist:    ', env._world.count('coal'))
  print('Iron exist:    ', env._world.count('iron'))
  print('Diamonds exist:', env._world.count('diamond'))

  start = time.time()
  done = False
  while not done:
    action = random.randint(0, env.action_space.n)
    obs, reward, done, info = env.step(action)
  duration = time.time() - start
  step = env._step
  print(f'Step time: {1000*duration/step:.2f}ms ({int(step/duration)} FPS)')
  print('Episode length:', step)

  if args.record:
    env.save(args.record)


if __name__ == '__main__':
  main()
