#game.py
import os

# Set RENDER flag
RENDER = False  # Set to True to enable rendering

if not RENDER:
    os.environ["SDL_VIDEODRIVER"] = "dummy"

import pygame
import random
import torch
import sys
import numpy as np
import matplotlib.pyplot as plt

from collections import deque
from agent_cuda import DQNAgent
from parameters import (
    lr, gamma, epsilon, epsilon_decay, buffer_size,
    penalty, DEVICE, batch_size, target_update
)

# Initialize Pygame
pygame.init()

# Screen dimensions
SCREEN_WIDTH = 400
SCREEN_HEIGHT = 600

# Colors
RED = (255, 0, 0)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# Game parameters
RECT_WIDTH = 35
RECT_HEIGHT = 35
PIPE_WIDTH = 70
PIPE_GAP = 220
GRAVITY = 1
JUMP_STRENGTH = -13
PIPE_SPEED = 5

class Rectangle:
    def __init__(self):
        self.x = SCREEN_WIDTH // 4
        self.y = SCREEN_HEIGHT // 2
        self.y_speed = 0

    def draw(self, screen):
        pygame.draw.rect(screen, RED, (self.x, self.y, RECT_WIDTH, RECT_HEIGHT))

    def distance(self, pipes):
        # Find the first pipe that is to the right of the rectangle
        for pipe in pipes:
            if pipe.top.left > self.x:
                # Calculate the horizontal distance to the pipe
                dist_horizontal = pipe.top.left - self.x

                # Calculate the vertical distance to the center of the gap in the pipe
                gap_y_center = pipe.top.height + PIPE_GAP // 2
                dist_vertically = self.y + RECT_HEIGHT // 2 - gap_y_center

                return dist_horizontal, dist_vertically

        return None, None  # If no pipe is to the right

    def jump(self):
        self.y_speed = JUMP_STRENGTH

    def apply_gravity(self):
        self.y_speed += GRAVITY
        self.y += self.y_speed

class Pipe:
    def __init__(self):
        self.top = pygame.Rect(SCREEN_WIDTH, 0, PIPE_WIDTH, random.randint(15, SCREEN_HEIGHT - PIPE_GAP - 15))
        self.bottom = pygame.Rect(SCREEN_WIDTH, self.top.height + PIPE_GAP, PIPE_WIDTH, SCREEN_HEIGHT - self.top.height - PIPE_GAP)

    def draw(self, screen):
        pygame.draw.rect(screen, WHITE, self.top)
        pygame.draw.rect(screen, WHITE, self.bottom)

    def move(self):
        self.top.x -= PIPE_SPEED
        self.bottom.x -= PIPE_SPEED

    def off_screen(self):
        return self.top.right < 0

def get_observation():
    dist_horizontal, dist_vertically = rectangle.distance(pipes)
    if dist_horizontal is None or dist_vertically is None:
        dist_horizontal = SCREEN_WIDTH
        dist_vertically = SCREEN_HEIGHT // 2

    # Normalize the distances
    normalized_horizontal = dist_horizontal / SCREEN_WIDTH
    normalized_vertical = dist_vertically / SCREEN_HEIGHT

    # Include the rectangle's y-position and velocity
    rect_y_normalized = rectangle.y / SCREEN_HEIGHT
    rect_y_speed_normalized = rectangle.y_speed / JUMP_STRENGTH

    observation = np.array([
        normalized_horizontal,
        normalized_vertical,
        rect_y_normalized,
        rect_y_speed_normalized
    ])
    return observation

def take_action(action):
    if action == 1:  # Jump
        rectangle.jump()

# Initialize the game components
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Flappy Rectangle")

rectangle = Rectangle()
pipes = []
pipe_timer = 80
score = 0
font = pygame.font.SysFont(None, 36)
game_active = True  # Start the game immediately

# Initialize the DQN Agent
agent = DQNAgent(
    state_dim=4,
    action_dim=2,
    lr=lr,
    gamma=gamma,
    epsilon=epsilon,
    epsilon_decay=epsilon_decay,
    epsilon_min=0.01,          # Ensuring epsilon_min is set
    buffer_size=buffer_size
)

# Tracking rewards and scores
episode_rewards = []
episode_scores = []
num_episodes = int(sys.argv[1]) if len(sys.argv) > 1 else 5000

# Ensure the models directory exists
os.makedirs('models', exist_ok=True)

for episode in range(num_episodes):
    rectangle = Rectangle()  # Reset rectangle position and speed
    pipes = []
    pipe_timer = 0
    score = 0
    game_active = True

    episode_reward = 0

    while game_active:
        if RENDER:
            screen.fill(BLACK)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

        obs = get_observation()
        action = agent.act(obs)
        take_action(action)
        rectangle.apply_gravity()

        pipe_timer += 1
        if pipe_timer > 90:
            pipes.append(Pipe())
            pipe_timer = 0
        pipes_to_remove = []
        for pipe in pipes:
            pipe.move()
            if pipe.off_screen():
                pipes_to_remove.append(pipe)
                score += 1
                episode_reward += 1  # Reward for passing a pipe
        for pipe in pipes_to_remove:
            pipes.remove(pipe)

        if RENDER:
            for pipe in pipes:
                pipe.draw(screen)

            rectangle.draw(screen)

            # Display the score
            score_text = font.render(f"Score: {score}", True, WHITE)
            screen.blit(score_text, (10, 10))

            pygame.display.flip()

        # Check for collisions
        done = False
        reward = 0.1  # Small reward for each step survived
        for pipe in pipes:
            rect = pygame.Rect(rectangle.x, rectangle.y, RECT_WIDTH, RECT_HEIGHT)
            if rect.colliderect(pipe.top) or rect.colliderect(pipe.bottom):
                game_active = False
                done = True
                reward = penalty
                break
        if rectangle.y <= 0 or rectangle.y + RECT_HEIGHT >= SCREEN_HEIGHT:
            game_active = False
            done = True
            reward = penalty

        next_obs = get_observation()

        agent.remember(obs, action, reward, next_obs, done)
        episode_reward += reward

        # Update the agent
        agent.replay()

        # Update target network periodically
        if (episode + 1) % target_update == 0:
            agent.update_target_network()

        if RENDER:
            # Remove frame limitation to speed up training
            pygame.time.Clock().tick(30)

    # Decay epsilon after each episode
    agent.decay_epsilon()

    # Track scores
    episode_rewards.append(episode_reward)
    episode_scores.append(score)

    # Print episode results
    print(f"Episode {episode+1}, Reward: {episode_reward:.2f}, Score: {score}, Epsilon: {agent.epsilon:.4f}")

    # Optionally save the model every few episodes
    if (episode + 1) % 50 == 0:
        number = len(os.listdir('models/')) + 1
        torch.save(agent.policy_net.state_dict(), f"models/flappy{number}.pth")

# Save the final trained model
number = len(os.listdir('models/')) + 1
torch.save(agent.policy_net.state_dict(), f"models/flappy{number}.pth")

# Plot the rewards and scores
fig, ax1 = plt.subplots(figsize=(12, 6))

color = 'tab:blue'
ax1.set_xlabel('Episode')
ax1.set_ylabel('Reward', color=color)
ax1.plot(range(len(episode_rewards)), episode_rewards, color=color, label='Reward')
ax1.tick_params(axis='y', labelcolor=color)
ax1.legend(loc='upper left')

ax2 = ax1.twinx()  # Instantiate a second axes that shares the same x-axis

color = 'tab:red'
ax2.set_ylabel('Score', color=color)  # We already handled the x-label with ax1
ax2.plot(range(len(episode_scores)), episode_scores, color=color, label='Score')
ax2.tick_params(axis='y', labelcolor=color)
ax2.legend(loc='upper right')

plt.title("Episode Rewards and Scores")
fig.tight_layout()  # Prevents labels from overlapping
plt.show()

pygame.quit()

# model.py
import torch
import torch.nn as nn
import torch.nn.functional as F

class DQN(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(DQN, self).__init__()
        self.fc1 = nn.Linear(state_dim, 64)
        self.fc2 = nn.Linear(64, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, action_dim)

    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        return x

#agent_cuda.py
import random
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from collections import deque
from model import DQN
from parameters import DEVICE, lr, gamma, epsilon, epsilon_min, epsilon_decay, buffer_size, batch_size, target_update, state_dim, action_dim

class DQNAgent:
    def __init__(self, state_dim, action_dim, lr, gamma, epsilon, epsilon_decay, epsilon_min, buffer_size):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma

        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min

        self.memory = deque(maxlen=buffer_size)
        self.batch_size = batch_size

        self.policy_net = DQN(state_dim, action_dim).to(DEVICE)
        self.target_net = DQN(state_dim, action_dim).to(DEVICE)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()

        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=lr)
        self.loss_fn = nn.MSELoss()

        self.steps_done = 0

    def remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def act(self, state):
        if random.random() < self.epsilon:
            return random.choices([0, 1], weights=[0.60, 0.40])[0]
        state = torch.FloatTensor(state).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            q_values = self.policy_net(state)
        return q_values.argmax().item()

    def replay(self):
        if len(self.memory) < self.batch_size:
            return
        
        batch = random.sample(self.memory, self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)

        states = torch.FloatTensor(np.array(states)).to(DEVICE)
        actions = torch.LongTensor(actions).unsqueeze(1).to(DEVICE)
        rewards = torch.FloatTensor(rewards).unsqueeze(1).to(DEVICE)
        next_states = torch.FloatTensor(np.array(next_states)).to(DEVICE)
        dones = torch.FloatTensor(dones).unsqueeze(1).to(DEVICE)

        # Current Q values
        current_q = self.policy_net(states).gather(1, actions)

        # Next Q values from target network
        with torch.no_grad():
            max_next_q = self.target_net(next_states).max(1)[0].unsqueeze(1)
            target_q = rewards + (self.gamma * max_next_q * (1 - dones))
        
        # Compute loss
        loss = self.loss_fn(current_q, target_q)

        # Optimize the model
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def decay_epsilon(self):
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay
            self.epsilon = max(self.epsilon_min, self.epsilon)

    def update_target_network(self):
        self.target_net.load_state_dict(self.policy_net.state_dict())

#parameters.py
import torch

# Device configuration
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Hyperparameters
lr = 0.001          # Learning rate
gamma = 0.99        # Discount factor
epsilon = 1.0       # Initial exploration rate
epsilon_min = 0.01  # Minimum exploration rate
epsilon_decay = 0.998  # Decay rate for exploration
buffer_size = 100000     # Replay buffer size
batch_size = 64        # Mini-batch size for training
target_update = 5      # How often to update the target network
penalty = -10          # Penalty for losing the game

# Other parameters
state_dim = 4    # Number of state dimensions
action_dim = 2   # Number of actions (jump or do nothing)