import pygame
from pygame.locals import *
from OpenGL.GL import *
import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import pycuda.gpuarray as gpuarray
import numpy as np

# ================== PARAMETERS ==================
NUM_TYPES          = 12
NUM_PARTICLES      = 3000
WORLD_SIZE         = 600.0
NUM_GLOBAL_REPULSIVE = 10   # User-defined: number of global repulsive force layers
NUM_GLOBAL_ATTRACTIVE = 0      # User-defined: number of global attractive force layers
NUM_TYPE_SPECIFIC    = 8  # User-defined: number of type-specific force layers
NUM_DENSITY_REPULSIVE = 1 # User-defined: number of density-dependent repulsive force layers
NUM_FORCES         = NUM_GLOBAL_REPULSIVE + NUM_GLOBAL_ATTRACTIVE + NUM_TYPE_SPECIFIC + NUM_DENSITY_REPULSIVE

# Parameters for global repulsive
FIRST_REPULSIVE_RANGE = 5
FIRST_REPULSIVE_STRENGTH = 0.02
REPULSIVE_RANGE_MULTIPLIER = 1.25    # Multiplier for range scaling between layers
REPULSIVE_STRENGTH_MULTIPLIER = 0.8 # Multiplier for strength scaling between layers
REPULSIVE_FALLOFF = 'quadratic'  # 'linear', 'quadratic', or 'none'

# Parameters for global attractive
FIRST_ATTRACTIVE_RANGE = 8
FIRST_ATTRACTIVE_STRENGTH = 0.04
ATTRACTIVE_RANGE_MULTIPLIER = 2.0   # Multiplier for range scaling between layers
ATTRACTIVE_STRENGTH_MULTIPLIER = 0.8# Multiplier for strength scaling between layers
ATTRACTIVE_FALLOFF = 'linear'  # 'linear', 'quadratic', or 'none'

# Parameters for type-specific
FIRST_TYPE_RANGE =3.5
FIRST_TYPE_STRENGTH = 0.5
TYPE_RANGE_MULTIPLIER = 1.5    # Multiplier for range scaling between layers
TYPE_STRENGTH_MULTIPLIER = 0.85      # Multiplier for strength scaling between layers
TYPE_FALLOFF = 'quadratic'  # 'linear', 'quadratic', or 'none'

# Parameters for density-dependent repulsive
FIRST_DENSITY_RANGE =7
FIRST_DENSITY_STRENGTH = 0.1
DENSITY_RANGE_MULTIPLIER = 1.8    # Multiplier for range scaling between layers
DENSITY_STRENGTH_MULTIPLIER = 0.25 # Multiplier for strength scaling between layers
DENSITY_FALLOFF = 'linear'  # 'linear', 'quadratic', or 'none'
FIRST_DENSITY_SCALING = 1   # Scaling factor per nearby particle (additional scaling per particle)
DENSITY_SCALING_MULTIPLIER = 1  # Multiplier for scaling factor between layers

# Density-dependent damping parameters - SMOOTH QUADRATIC VERSION
DENSITY_DAMPING_RADIUS = 25.0
MIN_NEIGHBORS          = 0.0      # below this → almost full damping (very little slowing)
MAX_NEIGHBORS          = 350.0    # above this → maximum damping (strong slowing)
MAX_DAMPING            = 0.9    # low density: very little velocity loss
MIN_DAMPING            = 0.8   # high density: strong velocity loss

DT                 = 0.3       # reduced time step for stability

# ================== CUDA KERNEL ==================

kernel_code = """
__global__ void update_particles(
    float *pos_x, float *pos_y,
    float *vel_x, float *vel_y,
    int   *types,
    float *interaction_matrices,   // shape: [num_forces * num_types * num_types]
    float *force_ranges,
    float *force_strengths,
    float *falloff_exponents,
    int   *force_types,
    float *density_scalings,
    int    n_particles,
    int    n_types,
    int    num_forces,
    float  dt,
    float  world_size,
    float  density_damping_radius,
    float  min_neighbors,
    float  max_neighbors,
    float  max_damping,
    float  min_damping
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n_particles) return;

    float x = pos_x[idx];
    float y = pos_y[idx];
    int   ta = types[idx];

    float fx = 0.0f;
    float fy = 0.0f;

    float damping_count = 0.0f;

    // First pass: compute neighbor counts for density-dependent forces and damping
    float counts[32];
    for (int i = 0; i < 32; i++) counts[i] = 0.0f;

    for (int j = 0; j < n_particles; ++j) {
        if (j == idx) continue;

        float dx = pos_x[j] - x;
        float dy = pos_y[j] - y;

        // periodic boundary
        dx -= world_size * roundf(dx / world_size);
        dy -= world_size * roundf(dy / world_size);

        float dist2 = dx*dx + dy*dy;
        if (dist2 < 1e-10f) continue;
        float dist = sqrtf(dist2);

        if (dist < density_damping_radius) {
            damping_count += 1.0f;
        }

        for (int f = 0; f < num_forces; ++f) {
            if (force_types[f] == 1 && dist < force_ranges[f]) {
                counts[f] += 1.0f;
            }
        }
    }

    // SMOOTH QUADRATIC DAMPING
    float t = (damping_count - min_neighbors) / (max_neighbors - min_neighbors);
    t = fmaxf(0.0f, fminf(1.0f, t));  // clamp to [0,1]
    float t2 = t * t;
    float local_damping = max_damping - (max_damping - min_damping) * t2;

    // Second pass: compute forces
    for (int j = 0; j < n_particles; ++j) {
        if (j == idx) continue;

        float dx = pos_x[j] - x;
        float dy = pos_y[j] - y;

        // periodic boundary
        dx -= world_size * roundf(dx / world_size);
        dy -= world_size * roundf(dy / world_size);

        float dist2 = dx*dx + dy*dy;
        if (dist2 < 1e-10f) continue;
        float dist = sqrtf(dist2);

        int tb = types[j];

        // Sum over all force layers
        for (int f = 0; f < num_forces; ++f) {
            float r_max = force_ranges[f];
            if (dist >= r_max) continue;

            // Get the interaction strength for this pair and this force
            float k = interaction_matrices[f * n_types * n_types + ta * n_types + tb];

            float norm_dist = dist / r_max;
            float base_fall = 1.0f - norm_dist;
            base_fall = fmaxf(0.0f, base_fall);

            float exp = falloff_exponents[f];
            float fall = (exp < 1e-6f) ? 1.0f : powf(base_fall, exp);

            float scale = 1.0f;
            if (force_types[f] == 1) {
                scale = density_scalings[f] * counts[f];
            }

            float f_mag = k * force_strengths[f] * fall * scale;
            float inv_dist = 1.0f / (dist + 1e-5f);

            fx += f_mag * dx * inv_dist;
            fy += f_mag * dy * inv_dist;
        }
    }

    // Update velocity and position
    vel_x[idx] += fx * dt;
    vel_y[idx] += fy * dt;

    vel_x[idx] *= local_damping;
    vel_y[idx] *= local_damping;

    pos_x[idx] += vel_x[idx] * dt;
    pos_y[idx] += vel_y[idx] * dt;

    // Wrap around (periodic boundaries)
    float hw = world_size * 0.5f;
    if (pos_x[idx] < -hw) pos_x[idx] += world_size;
    if (pos_x[idx] >  hw) pos_x[idx] -= world_size;
    if (pos_y[idx] < -hw) pos_y[idx] += world_size;
    if (pos_y[idx] >  hw) pos_y[idx] -= world_size;
}

__global__ void prepare_render(
    float *pos_x, float *pos_y,
    int   *types,
    float *type_colors,        // flattened [type * 3]
    float *rx, float *ry,
    float *colors,
    int    n_particles,
    float  world_size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n_particles) return;

    rx[idx] = (pos_x[idx] / world_size) * 2.0f;
    ry[idx] = (pos_y[idx] / world_size) * 2.0f;

    int t = types[idx];
    colors[idx*3 + 0] = type_colors[t*3 + 0];
    colors[idx*3 + 1] = type_colors[t*3 + 1];
    colors[idx*3 + 2] = type_colors[t*3 + 2];
}
"""

class MinimalParticleLife:
    def __init__(self):
        pygame.init()
        self.screen = pygame.display.set_mode((1280, 720), DOUBLEBUF | OPENGL)
        pygame.display.set_caption("Particle Life - Press R to reset with new matrices")
        glEnable(GL_POINT_SMOOTH)
        glPointSize(3.0)
        glEnable(GL_BLEND)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)

        self.mod = SourceModule(kernel_code)
        self.update_kernel = self.mod.get_function("update_particles")
        self.render_kernel = self.mod.get_function("prepare_render")

        # GPU arrays that persist
        self.d_rx     = gpuarray.zeros(NUM_PARTICLES, np.float32)
        self.d_ry     = gpuarray.zeros(NUM_PARTICLES, np.float32)
        self.d_colors = gpuarray.zeros(NUM_PARTICLES * 3, np.float32)

        # Initialize first simulation state
        self.reset_simulation()

    def reset_simulation(self):
        """Regenerate all random elements: positions, velocities, types, matrices, colors"""
        half = WORLD_SIZE / 2

        # Random particles
        pos_x = np.random.uniform(-half, half, NUM_PARTICLES).astype(np.float32)
        pos_y = np.random.uniform(-half, half, NUM_PARTICLES).astype(np.float32)
        vel_x = np.random.uniform(-0.3, 0.3, NUM_PARTICLES).astype(np.float32)
        vel_y = np.random.uniform(-0.3, 0.3, NUM_PARTICLES).astype(np.float32)
        types = np.random.randint(0, NUM_TYPES, NUM_PARTICLES).astype(np.int32)

        self.d_pos_x = gpuarray.to_gpu(pos_x)
        self.d_pos_y = gpuarray.to_gpu(pos_y)
        self.d_vel_x = gpuarray.to_gpu(vel_x)
        self.d_vel_y = gpuarray.to_gpu(vel_y)
        self.d_types = gpuarray.to_gpu(types)

        # === Per-force interaction matrices ===
        matrices = np.zeros((NUM_FORCES, NUM_TYPES, NUM_TYPES), dtype=np.float32)

        # Force parameters: ranges, strengths, exponents
        ranges    = np.zeros(NUM_FORCES, dtype=np.float32)
        strengths = np.zeros(NUM_FORCES, dtype=np.float32)
        falloff_exponents = np.zeros(NUM_FORCES, dtype=np.float32)
        force_types = np.zeros(NUM_FORCES, dtype=np.int32)  # 0: normal, 1: density-dependent
        density_scalings = np.zeros(NUM_FORCES, dtype=np.float32)

        falloff_map = {'none': 0.0, 'linear': 1.0, 'quadratic': 2.0, '8-dratic': 8.0, '16-dratic': 16.0}

        # Global repulsive forces
        force_index = 0
        r = FIRST_REPULSIVE_RANGE
        s = FIRST_REPULSIVE_STRENGTH
        falloff_exp = falloff_map[REPULSIVE_FALLOFF]
        for _ in range(NUM_GLOBAL_REPULSIVE):
            ranges[force_index] = r
            strengths[force_index] = s
            falloff_exponents[force_index] = falloff_exp
            matrices[force_index, :, :] = -1.0
            force_index += 1
            r *= REPULSIVE_RANGE_MULTIPLIER
            s *= REPULSIVE_STRENGTH_MULTIPLIER

        # Global attractive forces
        r = FIRST_ATTRACTIVE_RANGE
        s = FIRST_ATTRACTIVE_STRENGTH
        falloff_exp = falloff_map[ATTRACTIVE_FALLOFF]
        for _ in range(NUM_GLOBAL_ATTRACTIVE):
            ranges[force_index] = r
            strengths[force_index] = s
            falloff_exponents[force_index] = falloff_exp
            matrices[force_index, :, :] = 1.0
            force_index += 1
            r *= ATTRACTIVE_RANGE_MULTIPLIER
            s *= ATTRACTIVE_STRENGTH_MULTIPLIER

        # Type-specific forces
        r = FIRST_TYPE_RANGE
        s = FIRST_TYPE_STRENGTH
        falloff_exp = falloff_map[TYPE_FALLOFF]
        for _ in range(NUM_TYPE_SPECIFIC):
            ranges[force_index] = r
            strengths[force_index] = s
            falloff_exponents[force_index] = falloff_exp
            matrices[force_index, :, :] = np.random.uniform(-1.0, 1.0, (NUM_TYPES, NUM_TYPES))
            force_index += 1
            r *= TYPE_RANGE_MULTIPLIER
            s *= TYPE_STRENGTH_MULTIPLIER

        # Density-dependent repulsive forces
        r = FIRST_DENSITY_RANGE
        s = FIRST_DENSITY_STRENGTH
        sc = FIRST_DENSITY_SCALING
        falloff_exp = falloff_map[DENSITY_FALLOFF]
        for _ in range(NUM_DENSITY_REPULSIVE):
            ranges[force_index] = r
            strengths[force_index] = s
            falloff_exponents[force_index] = falloff_exp
            density_scalings[force_index] = sc
            matrices[force_index, :, :] = -1.0
            force_types[force_index] = 1
            force_index += 1
            r *= DENSITY_RANGE_MULTIPLIER
            s *= DENSITY_STRENGTH_MULTIPLIER
            sc *= DENSITY_SCALING_MULTIPLIER

        self.d_matrices = gpuarray.to_gpu(matrices.flatten())
        self.d_ranges    = gpuarray.to_gpu(ranges)
        self.d_strengths = gpuarray.to_gpu(strengths)
        self.d_falloff_exponents = gpuarray.to_gpu(falloff_exponents)
        self.d_force_types = gpuarray.to_gpu(force_types)
        self.d_density_scalings = gpuarray.to_gpu(density_scalings)

        # Random colors
        colors = np.random.rand(NUM_TYPES, 3).astype(np.float32)
        self.d_type_colors = gpuarray.to_gpu(colors.flatten())

        print(f"Simulation reset with {NUM_GLOBAL_REPULSIVE} global repulsive, {NUM_GLOBAL_ATTRACTIVE} global attractive, {NUM_TYPE_SPECIFIC} type-specific, and {NUM_DENSITY_REPULSIVE} density-dependent repulsive force layers")

    def update(self):
        block = (256, 1, 1)
        grid = ((NUM_PARTICLES + 255) // 256, 1)

        self.update_kernel(
            self.d_pos_x, self.d_pos_y,
            self.d_vel_x, self.d_vel_y,
            self.d_types,
            self.d_matrices,
            self.d_ranges,
            self.d_strengths,
            self.d_falloff_exponents,
            self.d_force_types,
            self.d_density_scalings,
            np.int32(NUM_PARTICLES),
            np.int32(NUM_TYPES),
            np.int32(NUM_FORCES),
            np.float32(DT),
            np.float32(WORLD_SIZE),
            np.float32(DENSITY_DAMPING_RADIUS),
            np.float32(MIN_NEIGHBORS),
            np.float32(MAX_NEIGHBORS),
            np.float32(MAX_DAMPING),
            np.float32(MIN_DAMPING),
            block=block, grid=grid
        )

    def render(self):
        glClear(GL_COLOR_BUFFER_BIT)
        glLoadIdentity()

        self.render_kernel(
            self.d_pos_x, self.d_pos_y,
            self.d_types,
            self.d_type_colors,
            self.d_rx, self.d_ry,
            self.d_colors,
            np.int32(NUM_PARTICLES),
            np.float32(WORLD_SIZE),
            block=(256,1,1), grid=((NUM_PARTICLES+255)//256, 1)
        )

        x = self.d_rx.get()
        y = self.d_ry.get()
        c = self.d_colors.get().reshape(-1, 3)

        vertices = np.column_stack((x, y)).astype(np.float32)

        glEnableClientState(GL_VERTEX_ARRAY)
        glEnableClientState(GL_COLOR_ARRAY)
        glVertexPointer(2, GL_FLOAT, 0, vertices)
        glColorPointer(3, GL_FLOAT, 0, c)
        glDrawArrays(GL_POINTS, 0, NUM_PARTICLES)
        glDisableClientState(GL_VERTEX_ARRAY)
        glDisableClientState(GL_COLOR_ARRAY)

        pygame.display.flip()

    def run(self):
        clock = pygame.time.Clock()

        running = True

        print("Particle Life with user-defined global repulsive, attractive, type-specific, and density-dependent repulsive force layers")
        print("Now with smooth quadratic density-dependent damping:")
        print(f" - Damping smoothly transitions from {MAX_DAMPING} (low density) to {MIN_DAMPING} (high density)")
        print(f" - Low density threshold: {MIN_NEIGHBORS} neighbors")
        print(f" - High density threshold: {MAX_NEIGHBORS} neighbors")
        print(f" - Density measured within radius {DENSITY_DAMPING_RADIUS}")
        print("Press R to reset with new random matrices")
        print("ESC or close window to quit")

        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_r:
                        self.reset_simulation()

            self.update()
            self.render()
            clock.tick(60)

        pygame.quit()

if __name__ == "__main__":
    sim = MinimalParticleLife()
    sim.run()