import sys
import torch
import torch.nn as nn
from models import FeedForward
from toy_dataset import ToyDataset, plot_data
import matplotlib.pyplot as plt
import numpy as np

fname = sys.argv[1]

n_samples = 10000
hidden_size = 512

model = FeedForward(input_size=2, hidden_size=hidden_size, output_size=4)
model.load_state_dict(torch.load(fname), strict=True)
model.eval()

dataset_test = ToyDataset(n_samples)

# For testing just do everything in one giant batch
testloader = torch.utils.data.DataLoader(
    dataset_test,
    batch_size=len(dataset_test),
    shuffle=False,
    num_workers=0,
)

criterion = nn.CrossEntropyLoss()

with torch.no_grad():
    # Everything is in one batch, so this loop will only happen once
    # Need to scale to SSP coordinates
    # Env is 0 to 13, SSP is -5 to 5
    x = ((x_env - 0) / coarse_size) * limit_range + xs[0]
    y = ((y_env - 0) / coarse_size) * limit_range + ys[0]

    item_memory += vocab[sp_name] * encode_point(x, y, x_axis_sp, y_axis_sp)
item_memory.normalize()

# Component functions of the full system

cleanup_network = FeedForward(input_size=ssp_dim,
                              hidden_size=512,
                              output_size=ssp_dim)
cleanup_network.load_state_dict(torch.load(args.cleanup_network), strict=True)
cleanup_network.eval()

# Input is x and y velocity plus the distance sensor measurements, plus map ID
localization_network = LocalizationModel(
    input_size=2 + n_sensors + n_maps,
    unroll_length=1,  #rollout_length,
    sp_dim=ssp_dim)
localization_network.load_state_dict(torch.load(args.localization_network),
                                     strict=True)
localization_network.eval()

if args.n_hidden_layers_policy == 1:
    policy_network = FeedForward(input_size=id_size + ssp_dim * 2,
                                 output_size=2)
else:
    policy_network = MLP(input_size=id_size + ssp_dim * 2,