import asyncio
import torch
from client import MyEnvClient
from models import MyAction
from train import DQN

async def main():
    client = MyEnvClient("http://127.0.0.1:8000")

    model = DQN()
    model.load_state_dict(torch.load("dqn.pt", map_location="cpu"))
    model.eval()

    obs = await client.reset()
    state = obs["observation"]["state"]
    print("Initial state:", state)

    while True:
        with torch.no_grad():
            q = model(torch.tensor([[float(state)]], dtype=torch.float32))
            action = torch.argmax(q).item()

        obs = await client.step(MyAction(move=action))
        state = obs["observation"]["state"]
        print(f"State: {state}, Action: {action}, Reward: {obs['reward']}, Done: {obs['done']}")

        if obs["done"]:
            print("finished", obs)
            break

if __name__ == "__main__":
    asyncio.run(main())