"""
Stable 2D fluid simulation (Jos Stam) using Pygame + PyOpenGL + NumPy.

- Single-file demo.
- Left mouse: add density (smoke).
- Right mouse: add velocity (drag to create flow).
- Keys: C = clear, D = toggle display velocity, Esc or close window to quit.

Dependencies:
    pip install pygame PyOpenGL numpy

Run:
    python stable_fluid_pygame_opengl.py

Notes: This is a CPU solver on a grid (not GPU accelerated). Grid size defaults to 128x128 which is interactive on modern machines.
"""

import sys
import math
import numpy as np
import pygame
from pygame.locals import *
from OpenGL.GL import *

# ----------------------- Fluid solver (Stable Fluids, Stam) -----------------------

class StableFluid:
    def __init__(self, N, diffusion, viscosity, dt):
        self.N = N  # grid size (N x N)
        self.size = N + 2  # include boundary
        self.dt = dt
        self.diff = diffusion
        self.visc = viscosity

        # fields (with boundary padding)
        self.s = np.zeros((self.size, self.size), dtype=np.float32)        # temp
        self.density = np.zeros((self.size, self.size), dtype=np.float32)  # smoke

        self.Vx = np.zeros((self.size, self.size), dtype=np.float32)
        self.Vy = np.zeros((self.size, self.size), dtype=np.float32)
        self.Vx0 = np.zeros((self.size, self.size), dtype=np.float32)
        self.Vy0 = np.zeros((self.size, self.size), dtype=np.float32)

    def add_density(self, x, y, amount):
        self.density[y, x] += amount

    def add_velocity(self, x, y, amount_x, amount_y):
        self.Vx[y, x] += amount_x
        self.Vy[y, x] += amount_y

    def step(self):
        N = self.N
        diff = self.diff
        visc = self.visc
        dt = self.dt

        self.diffuse(1, self.Vx0, self.Vx, visc, dt)
        self.diffuse(2, self.Vy0, self.Vy, visc, dt)

        self.project(self.Vx0, self.Vy0, self.Vx, self.Vy)

        self.advect(1, self.Vx, self.Vx0, self.Vx0, self.Vy0, dt)
        self.advect(2, self.Vy, self.Vy0, self.Vx0, self.Vy0, dt)

        self.project(self.Vx, self.Vy, self.Vx0, self.Vy0)

        self.diffuse(0, self.s, self.density, diff, dt)
        self.advect(0, self.density, self.s, self.Vx, self.Vy, dt)

    # helper routines
    def set_bnd(self, b, x):
        N = self.N
        # edges
        x[0, 1:N+1] = -x[1, 1:N+1] if b == 2 else x[1, 1:N+1]
        x[N+1, 1:N+1] = -x[N, 1:N+1] if b == 2 else x[N, 1:N+1]

        x[1:N+1, 0] = -x[1:N+1, 1] if b == 1 else x[1:N+1, 1]
        x[1:N+1, N+1] = -x[1:N+1, N] if b == 1 else x[1:N+1, N]

        # corners
        x[0, 0] = 0.5 * (x[1, 0] + x[0, 1])
        x[0, N+1] = 0.5 * (x[1, N+1] + x[0, N])
        x[N+1, 0] = 0.5 * (x[N, 0] + x[N+1, 1])
        x[N+1, N+1] = 0.5 * (x[N, N+1] + x[N+1, N])

    def lin_solve(self, b, x, x0, a, c):
        N = self.N
        for _ in range(20):  # Gauss-Seidel iterations
            x[1:N+1, 1:N+1] = (
                x0[1:N+1, 1:N+1] + a * (
                    x[0:N, 1:N+1] + x[2:N+2, 1:N+1] + x[1:N+1, 0:N] + x[1:N+1, 2:N+2]
                )
            ) / c
            self.set_bnd(b, x)

    def diffuse(self, b, x, x0, diff, dt):
        N = self.N
        a = dt * diff * N * N
        self.lin_solve(b, x, x0, a, 1 + 4 * a)

    def advect(self, b, d, d0, velocX, velocY, dt):
        N = self.N
        dt0 = dt * N
        for i in range(1, N+1):
            for j in range(1, N+1):
                x = i - dt0 * velocX[j, i]
                y = j - dt0 * velocY[j, i]

                if x < 0.5:
                    x = 0.5
                if x > N + 0.5:
                    x = N + 0.5
                i0 = int(math.floor(x))
                i1 = i0 + 1

                if y < 0.5:
                    y = 0.5
                if y > N + 0.5:
                    y = N + 0.5
                j0 = int(math.floor(y))
                j1 = j0 + 1

                s1 = x - i0
                s0 = 1 - s1
                t1 = y - j0
                t0 = 1 - t1

                d[j, i] = (
                    s0 * (t0 * d0[j0, i0] + t1 * d0[j1, i0]) +
                    s1 * (t0 * d0[j0, i1] + t1 * d0[j1, i1])
                )
        self.set_bnd(b, d)

    def project(self, velocX, velocY, p, div):
        N = self.N
        div[1:N+1, 1:N+1] = -0.5 * (
            velocX[1:N+1, 2:N+2] - velocX[1:N+1, 0:N] + velocY[2:N+2, 1:N+1] - velocY[0:N, 1:N+1]
        ) / N
        p[1:N+1, 1:N+1] = 0
        self.set_bnd(0, div)
        self.set_bnd(0, p)

        self.lin_solve(0, p, div, 1, 4)

        velocX[1:N+1, 1:N+1] -= 0.5 * (p[1:N+1, 2:N+2] - p[1:N+1, 0:N]) * N
        velocY[1:N+1, 1:N+1] -= 0.5 * (p[2:N+2, 1:N+1] - p[0:N, 1:N+1]) * N

        self.set_bnd(1, velocX)
        self.set_bnd(2, velocY)

# ----------------------- OpenGL texture renderer -----------------------

class GLRenderer:
    def __init__(self, width, height, gridN):
        self.width = width
        self.height = height
        self.gridN = gridN
        self.tex_id = glGenTextures(1)
        self.init_texture()

    def init_texture(self):
        glBindTexture(GL_TEXTURE_2D, self.tex_id)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)

        # initial empty texture (RGBA)
        data = np.zeros((self.gridN, self.gridN, 4), dtype=np.uint8)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, self.gridN, self.gridN, 0, GL_RGBA, GL_UNSIGNED_BYTE, data)

    def update_texture_from_density(self, density):
        # density is (N+2)x(N+2), we upload only 1..N
        N = self.gridN
        den = density[1:N+1, 1:N+1]
        # map density to 0-255
        img = np.empty((N, N, 4), dtype=np.uint8)
        # grayscale smoke map; store in rgb and alpha
        clipped = np.clip(den * 255.0, 0, 255).astype(np.uint8)
        img[:, :, 0] = clipped
        img[:, :, 1] = clipped
        img[:, :, 2] = clipped
        img[:, :, 3] = clipped

        glBindTexture(GL_TEXTURE_2D, self.tex_id)
        glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, N, N, GL_RGBA, GL_UNSIGNED_BYTE, img)

    def draw_fullscreen(self):
        glEnable(GL_TEXTURE_2D)
        glBindTexture(GL_TEXTURE_2D, self.tex_id)
        glBegin(GL_QUADS)
        glTexCoord2f(0.0, 0.0); glVertex2f(-1.0, -1.0)
        glTexCoord2f(1.0, 0.0); glVertex2f(1.0, -1.0)
        glTexCoord2f(1.0, 1.0); glVertex2f(1.0, 1.0)
        glTexCoord2f(0.0, 1.0); glVertex2f(-1.0, 1.0)
        glEnd()
        glDisable(GL_TEXTURE_2D)

    def draw_velocity(self, fluid):
        # draw simple vector field lines
        N = fluid.N
        glPushAttrib(GL_CURRENT_BIT | GL_POINT_BIT | GL_LINE_BIT)
        glPointSize(2.0)
        glBegin(GL_LINES)
        step = max(1, N // 32)
        for i in range(1, N+1, step):
            for j in range(1, N+1, step):
                x = (i - 0.5) / N * 2 - 1
                y = (j - 0.5) / N * 2 - 1
                vx = fluid.Vx[j, i]
                vy = fluid.Vy[j, i]
                glVertex2f(x, y)
                glVertex2f(x + vx * 0.02, y + vy * 0.02)
        glEnd()
        glPopAttrib()

# ----------------------- Main app -----------------------

def main():
    # settings
    WIN_W, WIN_H = 800, 800
    GRID_N = 128
    DIFFUSION = 0.0001
    VISCOSITY = 0.0001
    DT = 0.1

    pygame.init()
    pygame.display.set_mode((WIN_W, WIN_H), DOUBLEBUF | OPENGL)
    pygame.display.set_caption('Stable Fluid (Stam) - Pygame + PyOpenGL')

    # setup orthographic projection
    glViewport(0, 0, WIN_W, WIN_H)
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    glOrtho(-1, 1, -1, 1, -1, 1)
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()

    fluid = StableFluid(GRID_N, DIFFUSION, VISCOSITY, DT)
    renderer = GLRenderer(WIN_W, WIN_H, GRID_N)

    clock = pygame.time.Clock()
    running = True
    show_velocity = False

    is_down_left = False
    is_down_right = False
    last_mouse = None

    while running:
        for event in pygame.event.get():
            if event.type == QUIT:
                running = False
            elif event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    running = False
                elif event.key == K_c:
                    fluid.density.fill(0)
                    fluid.Vx.fill(0)
                    fluid.Vy.fill(0)
                elif event.key == K_d:
                    show_velocity = not show_velocity
            elif event.type == MOUSEBUTTONDOWN:
                if event.button == 1:
                    is_down_left = True
                elif event.button == 3:
                    is_down_right = True
                last_mouse = pygame.mouse.get_pos()
            elif event.type == MOUSEBUTTONUP:
                if event.button == 1:
                    is_down_left = False
                elif event.button == 3:
                    is_down_right = False
                last_mouse = None

        mx, my = pygame.mouse.get_pos()
        # convert to grid coords (i,x) with padding
        def screen_to_grid(px, py):
            gx = int((px / WIN_W) * GRID_N) + 1
            gy = int(((WIN_H - py) / WIN_H) * GRID_N) + 1
            gx = max(1, min(GRID_N, gx))
            gy = max(1, min(GRID_N, gy))
            return gx, gy

        if last_mouse is not None:
            lx, ly = last_mouse
            if is_down_left or is_down_right:
                g1 = screen_to_grid(lx, ly)
                g2 = screen_to_grid(mx, my)
                gx1, gy1 = g1
                gx2, gy2 = g2
                # draw along line between points
                steps = max(abs(gx2 - gx1), abs(gy2 - gy1), 1)
                for s in range(steps + 1):
                    t = s / steps
                    gx = int(round(gx1 + (gx2 - gx1) * t))
                    gy = int(round(gy1 + (gy2 - gy1) * t))
                    if is_down_left:
                        fluid.add_density(gx, gy, 100.0 * DT)
                    if is_down_right:
                        vx = (gx2 - gx1) * 5.0
                        vy = (gy2 - gy1) * 5.0
                        fluid.add_velocity(gx, gy, vx, vy)
                last_mouse = (mx, my)

        # step simulation
        fluid.step()

        # render
        glClear(GL_COLOR_BUFFER_BIT)
        renderer.update_texture_from_density(fluid.density)
        renderer.draw_fullscreen()

        if show_velocity:
            glColor3f(1.0, 0.0, 0.0)
            renderer.draw_velocity(fluid)
            glColor3f(1.0, 1.0, 1.0)

        pygame.display.flip()
#
        clock.tick(10)

    pygame.quit()

if __name__ == '__main__':
    main()