import os
import torch
from src.olgen.ol_generator import VecOnlineGenerator
from src.olgen.olg_game import MarioOnlineGenGame
from src.olgen.olg_policy import RLGenPolicy
from src.smb.level import save_batch
from src.utils.filesys import getpath
from src.utils.img import make_img_sheet

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

if __name__ == '__main__':
    path = 'models/example_policy'
    # Generate with example policy model
    N, L = 8, 10
    plc = RLGenPolicy.from_path(path, device=device)
    generator = VecOnlineGenerator(plc, g_device=device)
    fd, _ = os.path.split(getpath(path))
    os.makedirs(fd, exist_ok=True)

    lvls = generator.generate(N, L)
    save_batch(lvls, f'{path}/samples.lvls')
    imgs = [lvl.to_img() for lvl in lvls]
    make_img_sheet(imgs, 1, save_path=f'{path}/samples.png')

    # # Play with the example policy model
    # 请保证您的电脑上已经安装了jvm, 并且在命令行中输入java可以看到Java的信息
    game = MarioOnlineGenGame(path)
    game.play()
    pass