import matplotlib.pyplot as plt
import pandas as pd

from game import run_game
from util import initial_q_fill


AMOUNT_RUNS = 5000
EPSILON = 0.5
ALPHA = 0.3
GAMMA = 0.5


"""
state: (x_distance_to_ghost, y_distance_to_ghost, next_cookie_Direction)
action: Direction
q_value: (state, action)
"""
q_values = {}
initial_q_fill(q_values)

cookies_per_run = []
# Amount of single runs
for x in range(AMOUNT_RUNS):
    amount_cookies_ate = run_game(q_values, EPSILON, ALPHA, GAMMA)
    cookies_per_run.append(amount_cookies_ate)
    # print(f"Run {x}: {amount_cookies_ate} cookies ate\n")

wins = 0
for element in cookies_per_run:
    if element == 20:
        wins += 1
    
print(f"Win percentage: {wins/AMOUNT_RUNS}%")


window_size = 100  # Adjust based on your needs
rolling_avg = pd.Series(cookies_per_run).rolling(window=window_size, center=True).mean()

plt.figure(figsize=(12, 6))
plt.plot(cookies_per_run, alpha=0.2, label='Raw Data', linewidth=0.5, color='gray')
plt.plot(rolling_avg, label=f'{window_size}-point Moving Average', 
            linewidth=2, color='blue')
plt.title("Data with Rolling Average")
plt.xlabel("Index")
plt.ylabel("Value")
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()