#include <SDL2/SDL.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>  // For random seed

#define WIDTH 800
#define HEIGHT 800
#define SCALE 4  // Size of fluid cells
#define IX(x, y) ((x) + (y) * N)
#define SWAP(x0, x) {float *tmp = x0; x0 = x; x = tmp;}

// Simulation parameters
#define N (WIDTH / SCALE)
#define SIZE (N * N)
#define ITER 8        // Reduced iterations for stability
#define DIFFUSION 0.00001  // Lower diffusion
#define VISCOSITY 0.00001  // Lower viscosity
#define DT 0.02       // Much smaller time step for stability
#define FORCE 200.0   // Much lower force to prevent instability
#define VORTICITY 2.0  // Much lower vorticity
#define COLORFUL 1     // Enable colorful dye
#define CONTINUOUS_SOURCE 1  // Add continuous source

// Utility function to check for NaN
#define CHECK_NAN(x) \
    if (isnan(x)) { \
        printf("NaN detected! Fixing...\n"); \
        x = 0.0f; \
    }

// Fluid cell data
typedef struct {
    float *density;
    float *density_prev;

    float *u;
    float *v;
    float *u_prev;
    float *v_prev;

    // For color dye
    float *r;
    float *g;
    float *b;
    float *r_prev;
    float *g_prev;
    float *b_prev;
} fluid_grid;

// Function prototypes
void fluid_step(fluid_grid *grid);
void add_density(fluid_grid *grid, int x, int y, float amount, int color);
void add_velocity(fluid_grid *grid, int x, int y, float amount_x, float amount_y);
void diffuse(int b, float *x, float *x0, float diff, float dt);
void project(float *u, float *v, float *p, float *div);
void advect(int b, float *d, float *d0, float *u, float *v, float dt);
void render(SDL_Renderer *renderer, fluid_grid *grid);
void set_boundaries(int b, float *x);
void vorticity_confinement(fluid_grid *grid);

// Main function
int main(int argc, char *argv[]) {
    // Initialize SDL
    if (SDL_Init(SDL_INIT_VIDEO) != 0) {
        printf("SDL_Init Error: %s\n", SDL_GetError());
        return 1;
    }

    SDL_Window *window = SDL_CreateWindow("Fluid Simulation",
                                          SDL_WINDOWPOS_CENTERED,
                                          SDL_WINDOWPOS_CENTERED,
                                          WIDTH, HEIGHT,
                                          SDL_WINDOW_SHOWN);
    if (window == NULL) {
        printf("SDL_CreateWindow Error: %s\n", SDL_GetError());
        SDL_Quit();
        return 1;
    }

    SDL_Renderer *renderer = SDL_CreateRenderer(window, -1,
                                               SDL_RENDERER_ACCELERATED |
                                               SDL_RENDERER_PRESENTVSYNC);
    if (renderer == NULL) {
        SDL_DestroyWindow(window);
        printf("SDL_CreateRenderer Error: %s\n", SDL_GetError());
        SDL_Quit();
        return 1;
    }

    printf("Simulation started. Window size: %d x %d, Grid size: %d x %d\n", WIDTH, HEIGHT, N, N);

    // Create fluid grid
    fluid_grid grid;
    grid.density = calloc(SIZE, sizeof(float));
    grid.density_prev = calloc(SIZE, sizeof(float));
    grid.u = calloc(SIZE, sizeof(float));
    grid.v = calloc(SIZE, sizeof(float));
    grid.u_prev = calloc(SIZE, sizeof(float));
    grid.v_prev = calloc(SIZE, sizeof(float));

    // For colorful dye
    grid.r = calloc(SIZE, sizeof(float));
    grid.g = calloc(SIZE, sizeof(float));
    grid.b = calloc(SIZE, sizeof(float));
    grid.r_prev = calloc(SIZE, sizeof(float));
    grid.g_prev = calloc(SIZE, sizeof(float));
    grid.b_prev = calloc(SIZE, sizeof(float));

    // Add some initial fluid to make things visible right away
    srand(time(NULL));  // Initialize random seed

    printf("Adding initial fluid...\n");

    // Create a strong vortex in the center
    int cx = N/2;
    int cy = N/2;

    // Add a central density blob
    for (int i = cx-10; i <= cx+10; i++) {
        for (int j = cy-10; j <= cy+10; j++) {
            if (i >= 0 && i < N && j >= 0 && j < N) {
                float dist = sqrtf((i-cx)*(i-cx) + (j-cy)*(j-cy));
                if (dist < 10) {
                    float amount = 20 * (1 - dist/10);  // Reduced amount
                    add_density(&grid, i, j, amount, 1);

                    // Add vivid colors
                    grid.r[IX(i, j)] = 50 + rand() % 50;  // Lower values
                    grid.g[IX(i, j)] = 50 + rand() % 50;
                    grid.b[IX(i, j)] = 50 + rand() % 50;

                    // Create swirling velocity - much lower values
                    float angle = atan2f(j-cy, i-cx);
                    float speed = 20 * (1 - dist/10);  // Much smaller velocity
                    add_velocity(&grid, i, j, -speed * sinf(angle), speed * cosf(angle));
                }
            }
        }
    }

    // Add some random splotches too
    for (int n = 0; n < 5; n++) {  // Fewer splotches
        int x = N/4 + rand() % (N/2);
        int y = N/4 + rand() % (N/2);

        for (int i = x-3; i <= x+3; i++) {  // Smaller radius
            for (int j = y-3; j <= y+3; j++) {
                if (i >= 0 && i < N && j >= 0 && j < N) {
                    float dist = sqrtf((i-x)*(i-x) + (j-y)*(j-y));
                    if (dist < 3) {
                        add_density(&grid, i, j, 10 * (1 - dist/3), 1);  // Lower density
                        grid.r[IX(i, j)] = 20 + rand() % 30;  // Lower colors
                        grid.g[IX(i, j)] = 20 + rand() % 30;
                        grid.b[IX(i, j)] = 20 + rand() % 30;
                    }
                }
            }
        }
    }

    printf("Initial fluid added.\n");

    // Mouse state
    int mouse_x = 0, mouse_y = 0;
    int prev_mouse_x = 0, prev_mouse_y = 0;
    int mouse_down = 0;
    int right_mouse_down = 0;

    // Main loop
    int quit = 0;
    int frame_count = 0;
    while (!quit) {
        SDL_Event event;
        while (SDL_PollEvent(&event)) {
            switch (event.type) {
                case SDL_QUIT:
                    quit = 1;
                    break;

                case SDL_MOUSEMOTION:
                    mouse_x = event.motion.x / SCALE;
                    mouse_y = event.motion.y / SCALE;
                    break;

                case SDL_MOUSEBUTTONDOWN:
                    if (event.button.button == SDL_BUTTON_LEFT) {
                        mouse_down = 1;
                        prev_mouse_x = mouse_x;
                        prev_mouse_y = mouse_y;
                    } else if (event.button.button == SDL_BUTTON_RIGHT) {
                        right_mouse_down = 1;
                    }
                    break;

                case SDL_MOUSEBUTTONUP:
                    if (event.button.button == SDL_BUTTON_LEFT) {
                        mouse_down = 0;
                    } else if (event.button.button == SDL_BUTTON_RIGHT) {
                        right_mouse_down = 0;
                    }
                    break;
            }
        }

        // Add velocity with mouse drag
        if (mouse_down) {
            float dx = mouse_x - prev_mouse_x;
            float dy = mouse_y - prev_mouse_y;

            if (dx != 0 || dy != 0) {
                add_velocity(&grid, mouse_x, mouse_y, dx * FORCE, dy * FORCE);
            }

            // Cycle through colors based on frame count for fun visual effect
            if (COLORFUL) {
                float r = 0.5f + 0.5f * sin(0.01f * frame_count);
                float g = 0.5f + 0.5f * sin(0.01f * frame_count + 2.0f);
                float b = 0.5f + 0.5f * sin(0.01f * frame_count + 4.0f);

                add_density(&grid, mouse_x, mouse_y, 100, 1);
                grid.r[IX(mouse_x, mouse_y)] += 100 * r;
                grid.g[IX(mouse_x, mouse_y)] += 100 * g;
                grid.b[IX(mouse_x, mouse_y)] += 100 * b;
            } else {
                add_density(&grid, mouse_x, mouse_y, 100, 0);
            }
        }

        // Add density with right mouse button
        if (right_mouse_down) {
            add_density(&grid, mouse_x, mouse_y, 200, 1);

            // Create a little vortex
            for (int i = -3; i <= 3; i++) {
                for (int j = -3; j <= 3; j++) {
                    int x = mouse_x + i;
                    int y = mouse_y + j;

                    if (x >= 1 && x < N-1 && y >= 1 && y < N-1) {
                        // Create a swirling effect around the cursor
                        float dx = i;
                        float dy = j;
                        float dist = sqrt(dx*dx + dy*dy);

                        if (dist > 0) {
                            // Perpendicular vector for swirling
                            float vx = -dy / dist * 50;
                            float vy = dx / dist * 50;

                            add_velocity(&grid, x, y, vx, vy);
                        }
                    }
                }
            }
        }

        // Update fluid simulation
        fluid_step(&grid);

        // Render to screen
        SDL_SetRenderDrawColor(renderer, 0, 0, 0, 255);
        SDL_RenderClear(renderer);
        render(renderer, &grid);
        SDL_RenderPresent(renderer);

        // Store prev mouse position
        prev_mouse_x = mouse_x;
        prev_mouse_y = mouse_y;

        frame_count++;
    }

    // Clean up
    free(grid.density);
    free(grid.density_prev);
    free(grid.u);
    free(grid.v);
    free(grid.u_prev);
    free(grid.v_prev);
    free(grid.r);
    free(grid.g);
    free(grid.b);
    free(grid.r_prev);
    free(grid.g_prev);
    free(grid.b_prev);

    SDL_DestroyRenderer(renderer);
    SDL_DestroyWindow(window);
    SDL_Quit();

    return 0;
}

// Update fluid simulation for one time step
void fluid_step(fluid_grid *grid) {
    // Add continuous source in the center to keep the simulation active
    if (CONTINUOUS_SOURCE) {
        static int frame = 0;
        frame++;

        // Center of the grid
        int cx = N/2;
        int cy = N/2;

        // Add a pulsing source with changing colors
        float strength = 5.0f + 3.0f * sin(frame * 0.05f);  // Much lower strength

        // Vary the position slightly for more interesting effects
        int x = cx + 2 * sin(frame * 0.02f);  // Smaller offset
        int y = cy + 2 * cos(frame * 0.03f);

        if (frame % 5 == 0) {  // Only add every few frames to avoid overwhelming
            add_density(grid, x, y, strength, 1);
            grid->r[IX(x, y)] += 20 * (0.5f + 0.5f * sin(frame * 0.01f));  // Lower color values
            grid->g[IX(x, y)] += 20 * (0.5f + 0.5f * sin(frame * 0.01f + 2.0f));
            grid->b[IX(x, y)] += 20 * (0.5f + 0.5f * sin(frame * 0.01f + 4.0f));

            // Add some swirling velocity too - much lower values
            float angle = frame * 0.01f;
            add_velocity(grid, x, y, 5 * cos(angle), 5 * sin(angle));  // Much smaller velocity
        }
    }

    // Check grid for NaN or extreme values before processing
    for (int i = 0; i < SIZE; i++) {
        CHECK_NAN(grid->u[i]);
        CHECK_NAN(grid->v[i]);
        CHECK_NAN(grid->density[i]);
        CHECK_NAN(grid->r[i]);
        CHECK_NAN(grid->g[i]);
        CHECK_NAN(grid->b[i]);

        // Clamp extreme values
        if (fabs(grid->u[i]) > 100.0f) grid->u[i] = grid->u[i] > 0 ? 100.0f : -100.0f;
        if (fabs(grid->v[i]) > 100.0f) grid->v[i] = grid->v[i] > 0 ? 100.0f : -100.0f;
        if (grid->density[i] > 100.0f) grid->density[i] = 100.0f;
        if (grid->r[i] > 100.0f) grid->r[i] = 100.0f;
        if (grid->g[i] > 100.0f) grid->g[i] = 100.0f;
        if (grid->b[i] > 100.0f) grid->b[i] = 100.0f;
    }

    // Diffuse velocities
    diffuse(1, grid->u_prev, grid->u, VISCOSITY, DT);
    diffuse(2, grid->v_prev, grid->v, VISCOSITY, DT);

    // Project to enforce incompressibility
    project(grid->u_prev, grid->v_prev, grid->u, grid->v);

    // Advect velocities
    advect(1, grid->u, grid->u_prev, grid->u_prev, grid->v_prev, DT);
    advect(2, grid->v, grid->v_prev, grid->u_prev, grid->v_prev, DT);

    // Project again
    project(grid->u, grid->v, grid->u_prev, grid->v_prev);

    // Add vorticity confinement
    vorticity_confinement(grid);

    // Diffuse and advect density
    diffuse(0, grid->density_prev, grid->density, DIFFUSION, DT);
    advect(0, grid->density, grid->density_prev, grid->u, grid->v, DT);

    // For colorful dye
    if (COLORFUL) {
        diffuse(0, grid->r_prev, grid->r, DIFFUSION, DT);
        diffuse(0, grid->g_prev, grid->g, DIFFUSION, DT);
        diffuse(0, grid->b_prev, grid->b, DIFFUSION, DT);

        advect(0, grid->r, grid->r_prev, grid->u, grid->v, DT);
        advect(0, grid->g, grid->g_prev, grid->u, grid->v, DT);
        advect(0, grid->b, grid->b_prev, grid->u, grid->v, DT);
    }

    // Debug - Print max values for first 10 frames
    static int debug_frame = 0;
    if (debug_frame < 10) {
        float max_d = 0, max_u = 0, max_v = 0;
        float sum_d = 0;
        for (int i = 0; i < SIZE; i++) {
            sum_d += grid->density[i];
            max_d = fmaxf(max_d, grid->density[i]);
            max_u = fmaxf(max_u, fabsf(grid->u[i]));
            max_v = fmaxf(max_v, fabsf(grid->v[i]));

            // Check for NaN and fix if needed
            CHECK_NAN(grid->density[i]);
            CHECK_NAN(grid->u[i]);
            CHECK_NAN(grid->v[i]);
            CHECK_NAN(grid->r[i]);
            CHECK_NAN(grid->g[i]);
            CHECK_NAN(grid->b[i]);
        }

        if (isnan(sum_d)) {
            printf("ERROR: NaN detected in density sum! Fixing the grid...\n");
            // Reset the grid to fix NaN values
            for (int i = 0; i < SIZE; i++) {
                if (isnan(grid->density[i])) grid->density[i] = 0.0f;
                if (isnan(grid->u[i])) grid->u[i] = 0.0f;
                if (isnan(grid->v[i])) grid->v[i] = 0.0f;
                if (isnan(grid->r[i])) grid->r[i] = 0.0f;
                if (isnan(grid->g[i])) grid->g[i] = 0.0f;
                if (isnan(grid->b[i])) grid->b[i] = 0.0f;
            }
        }

        printf("Frame %d: max_density=%.2f, sum_density=%.2f, max_u=%.2f, max_v=%.2f\n",
               debug_frame, max_d, sum_d, max_u, max_v);
        debug_frame++;
    }
}

// Add density to a specific cell with optional color
void add_density(fluid_grid *grid, int x, int y, float amount, int color) {
    if (x < 1 || x > N-2 || y < 1 || y > N-2) return;

    // Clamp density to avoid extreme values
    if (amount > 100.0f) amount = 100.0f;

    int index = IX(x, y);
    grid->density[index] += amount;

    // Check for NaN
    CHECK_NAN(grid->density[index]);

    // Add density to surrounding cells for a smoother effect - only add half as much
    grid->density[IX(x+1, y)] += amount * 0.3f;
    grid->density[IX(x-1, y)] += amount * 0.3f;
    grid->density[IX(x, y+1)] += amount * 0.3f;
    grid->density[IX(x, y-1)] += amount * 0.3f;

    // For color dye
    if (color && COLORFUL) {
        // More controlled colors
        float r_amount = amount * ((float)rand() / RAND_MAX);
        float g_amount = amount * ((float)rand() / RAND_MAX);
        float b_amount = amount * ((float)rand() / RAND_MAX);

        // Clamp to reasonable values
        if (r_amount > 100.0f) r_amount = 100.0f;
        if (g_amount > 100.0f) g_amount = 100.0f;
        if (b_amount > 100.0f) b_amount = 100.0f;

        grid->r[index] += r_amount;
        grid->g[index] += g_amount;
        grid->b[index] += b_amount;

        // Check for NaN
        CHECK_NAN(grid->r[index]);
        CHECK_NAN(grid->g[index]);
        CHECK_NAN(grid->b[index]);
    }
}

// Add velocity to a specific cell
void add_velocity(fluid_grid *grid, int x, int y, float amount_x, float amount_y) {
    if (x < 1 || x > N-2 || y < 1 || y > N-2) return;

    // Clamp velocity to avoid extreme values
    if (amount_x > 50.0f) amount_x = 50.0f;
    if (amount_x < -50.0f) amount_x = -50.0f;
    if (amount_y > 50.0f) amount_y = 50.0f;
    if (amount_y < -50.0f) amount_y = -50.0f;

    int index = IX(x, y);
    grid->u[index] += amount_x;
    grid->v[index] += amount_y;

    // Check for NaN
    CHECK_NAN(grid->u[index]);
    CHECK_NAN(grid->v[index]);

    // Add velocity to surrounding cells for a smoother effect - less spread
    grid->u[IX(x+1, y)] += amount_x * 0.25f;
    grid->u[IX(x-1, y)] += amount_x * 0.25f;
    grid->u[IX(x, y+1)] += amount_x * 0.25f;
    grid->u[IX(x, y-1)] += amount_x * 0.25f;

    grid->v[IX(x+1, y)] += amount_y * 0.25f;
    grid->v[IX(x-1, y)] += amount_y * 0.25f;
    grid->v[IX(x, y+1)] += amount_y * 0.25f;
    grid->v[IX(x, y-1)] += amount_y * 0.25f;
}

// Diffusion step
void diffuse(int b, float *x, float *x0, float diff, float dt) {
    float a = dt * diff * (N - 2) * (N - 2);

    for (int k = 0; k < ITER; k++) {
        for (int i = 1; i < N - 1; i++) {
            for (int j = 1; j < N - 1; j++) {
                x[IX(i, j)] = (x0[IX(i, j)] + a * (
                    x[IX(i+1, j)] + x[IX(i-1, j)] +
                    x[IX(i, j+1)] + x[IX(i, j-1)]
                )) / (1 + 4 * a);

                // Prevent extreme values that can cause instability
                if (x[IX(i, j)] > 100.0f) x[IX(i, j)] = 100.0f;
                if (x[IX(i, j)] < -100.0f) x[IX(i, j)] = -100.0f;

                // Check for NaN
                CHECK_NAN(x[IX(i, j)]);
            }
        }
        set_boundaries(b, x);
    }

    // Debug first few diffusions
    static int diffuse_count = 0;
    if (diffuse_count < 5) {
        float sum = 0, max_val = 0;
        for (int i = 0; i < SIZE; i++) {
            sum += fabsf(x[i]);
            max_val = fmaxf(max_val, fabsf(x[i]));

            // Check for NaN
            CHECK_NAN(x[i]);
        }
        printf("Diffuse %d (type %d): sum=%.2f, max=%.2f\n",
               diffuse_count, b, sum, max_val);
        diffuse_count++;
    }
}

// Project step (enforce incompressibility)
void project(float *u, float *v, float *p, float *div) {
    float h = 1.0f / N;

    // Calculate divergence
    for (int i = 1; i < N - 1; i++) {
        for (int j = 1; j < N - 1; j++) {
            div[IX(i, j)] = -0.5f * h * (
                u[IX(i+1, j)] - u[IX(i-1, j)] +
                v[IX(i, j+1)] - v[IX(i, j-1)]
            );
            p[IX(i, j)] = 0;

            // Check for NaN
            CHECK_NAN(div[IX(i, j)]);

            // Clamp divergence to avoid instability
            if (div[IX(i, j)] > 100.0f) div[IX(i, j)] = 100.0f;
            if (div[IX(i, j)] < -100.0f) div[IX(i, j)] = -100.0f;
        }
    }
    set_boundaries(0, div);
    set_boundaries(0, p);

    // Solve pressure
    for (int k = 0; k < ITER; k++) {
        for (int i = 1; i < N - 1; i++) {
            for (int j = 1; j < N - 1; j++) {
                p[IX(i, j)] = (div[IX(i, j)] + (
                    p[IX(i+1, j)] + p[IX(i-1, j)] +
                    p[IX(i, j+1)] + p[IX(i, j-1)]
                )) / 4;

                // Check for NaN
                CHECK_NAN(p[IX(i, j)]);

                // Clamp pressure to avoid instability
                if (p[IX(i, j)] > 100.0f) p[IX(i, j)] = 100.0f;
                if (p[IX(i, j)] < -100.0f) p[IX(i, j)] = -100.0f;
            }
        }
        set_boundaries(0, p);
    }

    // Apply pressure gradient
    for (int i = 1; i < N - 1; i++) {
        for (int j = 1; j < N - 1; j++) {
            u[IX(i, j)] -= 0.5f * (p[IX(i+1, j)] - p[IX(i-1, j)]) / h;
            v[IX(i, j)] -= 0.5f * (p[IX(i, j+1)] - p[IX(i, j-1)]) / h;

            // Check for NaN
            CHECK_NAN(u[IX(i, j)]);
            CHECK_NAN(v[IX(i, j)]);

            // Clamp velocities to avoid instability
            if (u[IX(i, j)] > 100.0f) u[IX(i, j)] = 100.0f;
            if (u[IX(i, j)] < -100.0f) u[IX(i, j)] = -100.0f;
            if (v[IX(i, j)] > 100.0f) v[IX(i, j)] = 100.0f;
            if (v[IX(i, j)] < -100.0f) v[IX(i, j)] = -100.0f;
        }
    }
    set_boundaries(1, u);
    set_boundaries(2, v);

    // Debug - Check for NaN anywhere in the grid
    for (int i = 0; i < SIZE; i++) {
        CHECK_NAN(u[i]);
        CHECK_NAN(v[i]);
        CHECK_NAN(p[i]);
        CHECK_NAN(div[i]);
    }
}

// Advection step (semi-Lagrangian)
void advect(int b, float *d, float *d0, float *u, float *v, float dt) {
    float dt0 = dt * (N - 2);

    // Apply a tiny constant force to keep things moving
    static int advect_count = 0;

    for (int i = 1; i < N - 1; i++) {
        for (int j = 1; j < N - 1; j++) {
            // Trace particle back
            float x = i - dt0 * u[IX(i, j)];
            float y = j - dt0 * v[IX(i, j)];

            // Clamp to grid
            if (x < 0.5f) x = 0.5f;
            if (x > N - 1.5f) x = N - 1.5f;
            if (y < 0.5f) y = 0.5f;
            if (y > N - 1.5f) y = N - 1.5f;

            // Bilinear interpolation
            int i0 = (int)x;
            int i1 = i0 + 1;
            int j0 = (int)y;
            int j1 = j0 + 1;

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

            // Check for invalid indices
            if (i0 < 0) i0 = 0; if (i0 >= N) i0 = N-1;
            if (i1 < 0) i1 = 0; if (i1 >= N) i1 = N-1;
            if (j0 < 0) j0 = 0; if (j0 >= N) j0 = N-1;
            if (j1 < 0) j1 = 0; if (j1 >= N) j1 = N-1;

            d[IX(i, j)] =
                s0 * (t0 * d0[IX(i0, j0)] + t1 * d0[IX(i0, j1)]) +
                s1 * (t0 * d0[IX(i1, j0)] + t1 * d0[IX(i1, j1)]);

            // Check for NaN
            CHECK_NAN(d[IX(i, j)]);

            // Prevent too much dissipation for density
            if (b == 0 && d[IX(i, j)] > 0.01f) {
                // Apply a little damping but not too much
                d[IX(i, j)] *= 0.998f;

                // Ensure a minimum density if there was some density before
                if (d[IX(i, j)] < 0.1f && d0[IX(i, j)] > 0.1f) {
                    d[IX(i, j)] = 0.1f;
                }
            }

            // Clamp values to prevent instability
            if (d[IX(i, j)] > 100.0f) d[IX(i, j)] = 100.0f;
            if (b != 0 && d[IX(i, j)] < -100.0f) d[IX(i, j)] = -100.0f;
        }
    }

    // Debug for first few advections
    if (advect_count < 5) {
        float sum = 0, max_val = 0;
        for (int i = 0; i < SIZE; i++) {
            sum += fabsf(d[i]);
            max_val = fmaxf(max_val, fabsf(d[i]));

            // Check for NaN
            CHECK_NAN(d[i]);
        }
        printf("Advect %d (type %d): sum=%.2f, max=%.2f\n",
               advect_count, b, sum, max_val);
        advect_count++;
    }

    set_boundaries(b, d);
}

// Render the fluid simulation
void render(SDL_Renderer *renderer, fluid_grid *grid) {
    static int frame_counter = 0;
    frame_counter++;

    // Always print for the first 10 frames
    if (frame_counter <= 10 || frame_counter % 100 == 0) {
        // Debug: Print max density value to check if we have any density at all
        float max_density = 0;
        float sum_density = 0;
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                max_density = fmaxf(max_density, grid->density[IX(i, j)]);
                sum_density += grid->density[IX(i, j)];
            }
        }
        printf("Frame %d: Max density = %.2f, Sum density = %.2f\n",
               frame_counter, max_density, sum_density);
    }

    // First clear with a very dark blue background to show grid boundaries
    SDL_SetRenderDrawColor(renderer, 0, 0, 20, 255);
    SDL_RenderClear(renderer);

    // Count how many cells have visible content
    int visible_cells = 0;

    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            int idx = IX(i, j);
            float d = grid->density[idx];

            // Skip cells with basically no density - optimization
            if (d < 0.01f) continue;

            // Make small density values more visible (logarithmic scaling)
            d = 80.0f * log(1.0f + d * 0.2f);

            if (d > 255) d = 255;

            visible_cells++;

            Uint8 r, g, b;

            if (COLORFUL) {
                float r_val = grid->r[idx];
                float g_val = grid->g[idx];
                float b_val = grid->b[idx];

                // Ensure colors are visible even with low density
                float brightness = fmaxf(30.0f, d);

                // Normalize colors
                float max_color = fmaxf(r_val, fmaxf(g_val, b_val));
                if (max_color < 0.01f) {
                    // If no color, use white dye
                    r = g = b = (Uint8)brightness;
                } else {
                    // Normalize and apply brightness
                    r = (Uint8)(brightness * r_val / max_color);
                    g = (Uint8)(brightness * g_val / max_color);
                    b = (Uint8)(brightness * b_val / max_color);
                }

                // Ensure some minimum color to make it visible
                r = fmaxf(r, 20);
                g = fmaxf(g, 20);
                b = fmaxf(b, 20);
            } else {
                r = g = b = (Uint8)d;
            }

            SDL_SetRenderDrawColor(renderer, r, g, b, 255);
            SDL_Rect rect = {i * SCALE, j * SCALE, SCALE, SCALE};
            SDL_RenderFillRect(renderer, &rect);
        }
    }

    // Debug output for first 10 frames
    if (frame_counter <= 10) {
        printf("Frame %d: %d visible cells (%.1f%% of grid)\n",
               frame_counter, visible_cells, 100.0f * visible_cells / SIZE);
    }
}

// Set boundary conditions
void set_boundaries(int b, float *x) {
    // Edges
    for (int i = 1; i < N - 1; i++) {
        x[IX(i, 0)] = b == 2 ? -x[IX(i, 1)] : x[IX(i, 1)];
        x[IX(i, N-1)] = b == 2 ? -x[IX(i, N-2)] : x[IX(i, N-2)];
    }

    for (int j = 1; j < N - 1; j++) {
        x[IX(0, j)] = b == 1 ? -x[IX(1, j)] : x[IX(1, j)];
        x[IX(N-1, j)] = b == 1 ? -x[IX(N-2, j)] : x[IX(N-2, j)];
    }

    // Corners
    x[IX(0, 0)] = 0.5f * (x[IX(1, 0)] + x[IX(0, 1)]);
    x[IX(0, N-1)] = 0.5f * (x[IX(1, N-1)] + x[IX(0, N-2)]);
    x[IX(N-1, 0)] = 0.5f * (x[IX(N-2, 0)] + x[IX(N-1, 1)]);
    x[IX(N-1, N-1)] = 0.5f * (x[IX(N-2, N-1)] + x[IX(N-1, N-2)]);

    // Debug boundary check - Print max values every 500 frames
    static int counter = 0;
    if (++counter % 500 == 0) {
        float max_val = 0.0f;
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                max_val = fmaxf(max_val, fabsf(x[IX(i, j)]));
            }
        }
        printf("Max value in grid (type %d): %.2f\n", b, max_val);
    }
}

// Vorticity confinement to make the fluid more swirly
void vorticity_confinement(fluid_grid *grid) {
    float *curl = (float *)malloc(SIZE * sizeof(float));

    // Calculate curl (vorticity)
    for (int i = 1; i < N - 1; i++) {
        for (int j = 1; j < N - 1; j++) {
            float du_dy = (grid->u[IX(i, j+1)] - grid->u[IX(i, j-1)]) * 0.5f;
            float dv_dx = (grid->v[IX(i+1, j)] - grid->v[IX(i-1, j)]) * 0.5f;

            curl[IX(i, j)] = du_dy - dv_dx;
        }
    }

    // Calculate vorticity force
    for (int i = 1; i < N - 1; i++) {
        for (int j = 1; j < N - 1; j++) {
            float curl_ij = curl[IX(i, j)];

            // Find magnitude of curl gradient
            float dx = fabs(curl[IX(i+1, j)]) - fabs(curl[IX(i-1, j)]);
            float dy = fabs(curl[IX(i, j+1)]) - fabs(curl[IX(i, j-1)]);

            float grad_mag = sqrt(dx*dx + dy*dy) + 1e-5f;
            dx /= grad_mag;
            dy /= grad_mag;

            // Apply force
            grid->u[IX(i, j)] += VORTICITY * dy * curl_ij * DT;
            grid->v[IX(i, j)] -= VORTICITY * dx * curl_ij * DT;
        }
    }

    free(curl);
}