import pygame
import math
import random

# Настройки экрана
WIDTH, HEIGHT = 800, 600
FPS = 60

# Цвета
BLUE = (0, 105, 148)
DARK_BLUE = (0, 64, 128)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# Настройки жидкости
GRID_SIZE = 10
DAMPING = 0.98

# Инициализация Pygame
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("2D Liquid Simulation")
clock = pygame.time.Clock()

# Создание сетки для жидкости
grid_width = WIDTH // GRID_SIZE
grid_height = HEIGHT // GRID_SIZE
grid = [[0 for _ in range(grid_width)] for _ in range(grid_height)]
velocity = [[0 for _ in range(grid_width)] for _ in range(grid_height)]

# Функция для обновления сетки жидкости
def update_liquid():
    for y in range(1, grid_height - 1):
        for x in range(1, grid_width - 1):
            avg = (
                grid[y - 1][x] +
                grid[y + 1][x] +
                grid[y][x - 1] +
                grid[y][x + 1]
            ) / 4
            velocity[y][x] += avg - grid[y][x]
            velocity[y][x] *= DAMPING

    for y in range(grid_height):
        for x in range(grid_width):
            grid[y][x] += velocity[y][x]

# Функция для добавления возмущений в жидкость
def disturb_liquid(cx, cy, radius, strength):
    for y in range(max(0, cy - radius), min(grid_height, cy + radius)):
        for x in range(max(0, cx - radius), min(grid_width, cx + radius)):
            dx = x - cx
            dy = y - cy
            if dx * dx + dy * dy <= radius * radius:
                grid[y][x] += strength

# Главный игровой цикл
running = True
ball = {
    "x": WIDTH // 2,
    "y": 0,
    "radius": 15,
    "vy": 0,
    "gravity": 0.5
}

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Обновление жидкости
    update_liquid()

    # Обновление позиции шара
    ball["vy"] += ball["gravity"]
    ball["y"] += ball["vy"]

    # Проверка столкновения шара с жидкостью
    if ball["y"] // GRID_SIZE < grid_height:
        disturb_liquid(
            ball["x"] // GRID_SIZE,
            int(ball["y"] // GRID_SIZE),
            ball["radius"] // GRID_SIZE,
            strength=5
        )

    # Ограничение шара снизу
    if ball["y"] > HEIGHT:
        ball["y"] = HEIGHT
        ball["vy"] = 0

    # Отрисовка экрана
    screen.fill(BLACK)

    # Отрисовка жидкости
    for y in range(grid_height):
        for x in range(grid_width):
            value = grid[y][x]
            color = (0, min(255, max(0, int(128 + value * 128))), min(255, max(0, int(255 + value * 255))))
            pygame.draw.rect(screen, color, (x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE))

    # Отрисовка шара
    pygame.draw.circle(screen, WHITE, (ball["x"], int(ball["y"])), ball["radius"])

    pygame.display.flip()
    clock.tick(FPS)

pygame.quit()