import pygame
import numpy as np
import tensorflow as tf

# Инициализация Pygame
pygame.init()
screen_width = 800
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
clock = pygame.time.Clock()

# Создание модели искусственного интеллекта
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(32, activation='relu', input_shape=(2,)),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(2, activation='softmax')
])
model.compile(optimizer='adam', loss='categorical_crossentropy')

# Задание трассы
track_points = [(100, 200), (300, 100), (500, 300), (700, 200)]

# Начальное положение и скорость машинки
car_pos = np.array([track_points[0][0], track_points[0][1]])
car_speed = np.array([5, 0])

# Обучение модели
learning_rate = 0.001
discount_factor = 0.99
epsilon = 0.1

for episode in range(1000):
    state = car_pos.copy()
    total_reward = 0

    while True:
        # Выбор действия на основе epsilon-greedy
        if np.random.rand() < epsilon:
            action = np.random.randint(2)
        else:
            q_values = model.predict(state.reshape(1, 2))
            action = np.argmax(q_values[0])

        # Исполнение действия
        if action == 0:
            car_speed[1] -= 0.5
        else:
            car_speed[1] += 0.5

        # Ограничение скорости машинки
        car_speed[1] = np.clip(car_speed[1], -5, 5)

        # Обновление положения машинки
        car_pos += car_speed

        # Расчет вознаграждения
        if car_pos[0] > track_points[-1][0]:
            reward = 1
        else:
            reward = 0

        # Обновление модели с помощью обучения с подкреплением
        target = reward + discount_factor * np.max(model.predict(car_pos.reshape(1, 2))[0])
        q_values_target = model.predict(state.reshape(1, 2))
        q_values_target[0][action] = target
        model.fit(state.reshape(1, 2), q_values_target, epochs=1, verbose=0)

        # Обновление состояния и суммарного вознаграждения
        state = car_pos.copy()
        total_reward += reward

        # Отрисовка игры
        screen.fill((255, 255, 255))
        pygame.draw.lines(screen, (0, 0, 0), False, track_points, 5)
        pygame.draw.circle(screen, (255, 0, 0), car_pos.astype(int), 10)
        pygame.display.flip()
        clock.tick(60)

        if car_pos[0] > screen_width:
            break

    print("Episode:", episode, "Total Reward:", total_reward)

pygame.quit()