import torch
from pathlib import Path

import celeste_ai.plotting as plotting
from multiprocessing import Pool


m = Path("model_data/current")


def plot_pred(src_model):
	plotting.predicted_reward(
		src_model,
		m / f"plots/predicted/{src_model.stem}.png",

		device = torch.device("cpu")
	)

def plot_best(src_model):
	plotting.best_action(
		src_model,
		m / f"plots/best_action/{src_model.stem}.png",

		device = torch.device("cpu")
	)


for k, v in {
	#"prediction":	plot_pred,
	"best_action":	plot_best,
}.items():
	print(f"Making {k} plots...")
	with Pool(5) as p:
		p.map(
			v,
			list((m / "model_archive").iterdir())
		)