from pathlib import Path
import torch
from celeste import Celeste
import numpy as np
import matplotlib.pyplot as plt
from collections import namedtuple


compute_device = torch.device(
	"cuda" if torch.cuda.is_available() else "cpu"
)


# Celeste env properties
n_observations = len(Celeste.state_number_map)
n_actions = len(Celeste.action_space)


# Outline our network
class DQN(torch.nn.Module):
	def __init__(self, n_observations: int, n_actions: int):
		super(DQN, self).__init__()

		self.layers = torch.nn.Sequential(
			torch.nn.Linear(n_observations, 128),
			torch.nn.ReLU(),

			torch.nn.Linear(128, 128),
			torch.nn.ReLU(),

			torch.nn.Linear(128, 128),
			torch.nn.ReLU(),

			torch.torch.nn.Linear(128, n_actions)
		)

	# Can be called with one input, or with a batch.
	#
	# Returns tensor(
	#	[ Q(s, left), Q(s, right) ], ...
	# )
	#
	# Recall that Q(s, a) is the (expected) return of taking
	# action `a` at state `s`
	def forward(self, x):
		return self.layers(x)


policy_net = DQN(
	n_observations,
	n_actions
).to(compute_device)

target_net = DQN(
	n_observations,
	n_actions
).to(compute_device)

optimizer = torch.optim.AdamW(
	policy_net.parameters(),
	lr = 0.01, # Hyperparameter: learning rate
	amsgrad = True
)

Transition = namedtuple(
	"Transition",
	(
		"state",
		"action",
		"next_state",
		"reward"
	)
)




def makeplt(i, net):
	p = np.zeros((128, 128), dtype=np.float32)

	for r in range(len(p)):
		for c in range(len(p[r])):
			with torch.no_grad():
				k = net(
					torch.tensor(
						[c, r, 60, 80],
						dtype = torch.float32,
						device = compute_device
					).unsqueeze(0)
				)[0][i].item()

				p[r][c] = k
	return p


for i in Path("out/model_images").iterdir():


	checkpoint = torch.load(i)
	policy_net.load_state_dict(checkpoint["policy_state_dict"])


	fig, axs = plt.subplots(2, 4, figsize = (15, 10))

	for a in range(len(axs.ravel())):
		ax = axs.ravel()[a]
		ax.set(adjustable="box", aspect="equal")
		plot = ax.pcolor(
			makeplt(a, policy_net),
			cmap = "Greens_r",
			vmin = 0,
			vmax = 20
		)
		ax.set_title(Celeste.action_space[a])
		ax.invert_yaxis()
		fig.colorbar(plot)
	print(i)
	fig.savefig(f"out/{i.stem}.png")
	plt.close()