import cv2
import numpy as np
import pyautogui
import pygetwindow as gw
import pytesseract
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import random
import time
import ctypes
from ctypes import wintypes

pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"

PROCESS_VM_READ = 0x0010
PROCESS_QUERY_INFORMATION = 0x0400

def read_memory(process_id, address):
    process_handle = ctypes.windll.kernel32.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, False, process_id)
    if not process_handle:
        raise Exception("Cannot open process: {}".format(process_id))
    value = ctypes.c_int()
    bytes_read = wintypes.SIZE()
    ctypes.windll.kernel32.ReadProcessMemory(process_handle, ctypes.c_void_p(address), ctypes.byref(value), ctypes.sizeof(value), ctypes.byref(bytes_read))
    ctypes.windll.kernel32.CloseHandle(process_handle)
    return value.value

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(6480, 50)
        self.fc2 = nn.Linear(50, 5)  # Actions: w, a, s, d, space

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 6480)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

def get_game_image(game_window):
    screenshot = pyautogui.screenshot(region=(game_window.left, game_window.top, game_window.width, game_window.height))
    screenshot = np.array(screenshot)
    screenshot = cv2.cvtColor(screenshot, cv2.COLOR_BGR2GRAY)
    return cv2.resize(screenshot, (84, 84))

def perform_action(action_index, game_window):
    actions = ["w", "a", "s", "d", "space"]
    action = actions[action_index]
    print(f"Performing action: {action}")
    if action == "space":
        x = random.randint(game_window.left, game_window.left + game_window.width)
        y = random.randint(game_window.top, game_window.top + game_window.height)
        pyautogui.moveTo(x, y)
        pyautogui.keyDown(action)
        time.sleep(0.1)
        pyautogui.keyUp(action)
    else:
        pyautogui.keyDown(action)
        time.sleep(0.1)
        pyautogui.keyUp(action)
    return action

def train(model, device, game_window, optimizer, previous_score, previous_hp):
    model.train()
    current_score = previous_score
    current_hp = previous_hp
    for i in range(10000):
        image = get_game_image(game_window)
        image_tensor = torch.from_numpy(image).float().unsqueeze(0).unsqueeze(0).to(device)
        optimizer.zero_grad()
        output = model(image_tensor)
        action_index = output.max(1)[1].view(1, 1).item()
        action = perform_action(action_index, game_window)
        time.sleep(0.1)
        new_hp = read_memory(process_id, hp_address)
        new_score = read_memory(process_id, score_address)
        print(f"Previous HP: {previous_hp}, Current HP: {new_hp}")
        print(f"Previous Score: {previous_score}, Current Score: {new_score}")
        hp_diff = new_hp - previous_hp
        score_diff = new_score - previous_score
        reward = hp_diff + score_diff
        print(f"HP Diff: {hp_diff}, Score Diff: {score_diff}, Reward: {reward}")
        loss = -torch.tensor([reward], dtype=torch.float, device=device) * output[0][action_index]
        loss.backward()
        optimizer.step()
        previous_hp = new_hp
        previous_score = new_score
        print(f"Loss: {loss.item()}")

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = Net().to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    game_window = gw.getWindowsWithTitle("Yet Another Zombie Defense 2")[0]
    game_window.activate()
    initial_image = get_game_image(game_window)
    initial_hp = read_memory(process_id, hp_address)
    initial_score = read_memory(process_id, score_address)
    train(model, device, game_window, optimizer, initial_score, initial_hp)

if __name__ == "__main__":
    main()