import numpy as np
from numba import cuda
import pygame
import random
import time

# ====================== META DECAY RULE ======================
RULE_TRIGGERING_NEIGHBORS = [1, 1, 1, 1, 1, 1, 1, 1, 1]

# ====================== TWEAKABLE PARAMETERS ======================
GRID_WIDTH = 1920
GRID_HEIGHT = 1080
MAX_VALUE = 20
CIRCLE_RADIUS = 180
CIRCLE_INTENSITY = MAX_VALUE
FALLOFF_EXPONENT = 1.6
NOISE_DENSITY = 0.002

MAX_POSSIBLE_SUM = MAX_VALUE * 8

# ── Classic rule counts ──────────────────────────────────────────
RANDOM_RULE_COUNT_MIN_IGNITE  = [0, 2, 2, 2, 0, 0, 0, 0]
RANDOM_RULE_COUNT_MAX_IGNITE  = [0, 2, 2, 2, 0, 0, 0, 0]
RANDOM_RULE_COUNT_MIN_KILL    = [0, 6, 6, 6, 0, 0, 0, 0]
RANDOM_RULE_COUNT_MAX_KILL    = [0, 6, 6, 6, 0, 0, 0, 0]
RANDOM_RULE_COUNT_MIN_SURVIVE = [0, 0, 0, 0, 0, 0, 0, 0]
RANDOM_RULE_COUNT_MAX_SURVIVE = [0, 0, 0, 0, 0, 0, 0, 0]

DEFAULT_NUM_IGNITES  = [2, 3, 4, 5, 6, 6, 7, 7]
DEFAULT_NUM_KILLS    = [3, 4, 5, 6, 8, 10, 12, 16]
DEFAULT_NUM_SURVIVES = [0, 0, 1, 2, 2, 3, 3, 4]

# ── Neighbor-based rules ─────────────────────────────────────────
DEFAULT_NUM_NEIGHBOR_IGNITES = [1, 1, 1, 1, 1, 1, 1, 1]
DEFAULT_NUM_NEIGHBOR_KILLS   = [0, 1, 0, 1, 0, 1, 0, 1]   # example: sparse kills

RANDOM_RULE_COUNT_MIN_NEIGHBOR_IGNITE = [2] * 8
RANDOM_RULE_COUNT_MAX_NEIGHBOR_IGNITE = [2] * 8
RANDOM_RULE_COUNT_MIN_NEIGHBOR_KILL   = [6] * 8
RANDOM_RULE_COUNT_MAX_NEIGHBOR_KILL   = [6] * 8

# ── Initial rule generation ──────────────────────────────────────
random.seed()
initial_random_state = random.getstate()

current_num_ignites          = DEFAULT_NUM_IGNITES.copy()
current_num_kills            = DEFAULT_NUM_KILLS.copy()
current_num_survives         = DEFAULT_NUM_SURVIVES.copy()
current_num_neighbor_ignites = DEFAULT_NUM_NEIGHBOR_IGNITES.copy()
current_num_neighbor_kills   = DEFAULT_NUM_NEIGHBOR_KILLS.copy()

def generate_random_rules(num_ignites, num_kills, num_survives):
    ignition_sums = [[] for _ in range(9)]
    kill_sums     = [[] for _ in range(9)]
    survive_sums  = [[] for _ in range(9)]

    for n in range(1, 9):
        min_sum = n
        max_sum = n * MAX_VALUE
        all_possible = list(range(min_sum, max_sum + 1))
        possible_count = len(all_possible)

        num_i = num_ignites[n-1]
        num_k = num_kills[n-1]
        num_s = num_survives[n-1]
        total = num_i + num_k + num_s

        if total == 0:
            continue
        if total > possible_count:
            ratio = possible_count / total
            num_i = max(0, int(num_i * ratio))
            num_k = max(0, int(num_k * ratio))
            num_s = max(0, int(num_s * ratio))
            total = num_i + num_k + num_s

        if total == 0:
            continue

        chosen = random.sample(all_possible, total)
        ignition_sums[n] = sorted(chosen[:num_i])
        kill_sums[n]     = sorted(chosen[num_i:num_i + num_k])
        survive_sums[n]  = sorted(chosen[num_i + num_k:])

    return ignition_sums, kill_sums, survive_sums


def generate_neighbor_conditions(num_per_dir):
    conditions = [[] for _ in range(8)]
    for dir_idx in range(8):
        num = num_per_dir[dir_idx]
        if num == 0:
            continue
        possible = [(cv, ss) for cv in range(MAX_VALUE + 1) for ss in range(MAX_POSSIBLE_SUM + 1)]
        selected = random.sample(possible, min(num, len(possible)))
        conditions[dir_idx] = selected
    return conditions


IGNITION_SUMS, KILL_SUMS, SURVIVE_SUMS = generate_random_rules(
    current_num_ignites, current_num_kills, current_num_survives
)
NEIGHBOR_IGNITE_CONDITIONS = generate_neighbor_conditions(current_num_neighbor_ignites)
NEIGHBOR_KILL_CONDITIONS   = generate_neighbor_conditions(current_num_neighbor_kills)

# Lookup tables
ignition_lookup        = np.zeros((9,   MAX_POSSIBLE_SUM + 1), dtype=bool)
kill_lookup            = np.zeros((9,   MAX_POSSIBLE_SUM + 1), dtype=bool)
survive_lookup         = np.zeros((9,   MAX_POSSIBLE_SUM + 1), dtype=bool)
neighbor_ignite_lookup = np.zeros((8, MAX_VALUE + 1, MAX_POSSIBLE_SUM + 1), dtype=bool)
neighbor_kill_lookup   = np.zeros((8, MAX_VALUE + 1, MAX_POSSIBLE_SUM + 1), dtype=bool)

d_meta_rule = cuda.to_device(np.array(RULE_TRIGGERING_NEIGHBORS, dtype=np.int32))

def update_lookups():
    global d_ignition_lookup, d_kill_lookup, d_survive_lookup
    global d_neighbor_ignite_lookup, d_neighbor_kill_lookup

    ignition_lookup.fill(False)
    kill_lookup.fill(False)
    survive_lookup.fill(False)
    neighbor_ignite_lookup.fill(False)
    neighbor_kill_lookup.fill(False)

    for n in range(1, 9):
        for s in IGNITION_SUMS[n]:   ignition_lookup[n, s] = True
        for s in KILL_SUMS[n]:       kill_lookup[n, s] = True
        for s in SURVIVE_SUMS[n]:    survive_lookup[n, s] = True

    for dir_idx in range(8):
        for cv, ss in NEIGHBOR_IGNITE_CONDITIONS[dir_idx]:
            neighbor_ignite_lookup[dir_idx, cv, ss] = True
        for cv, ss in NEIGHBOR_KILL_CONDITIONS[dir_idx]:
            neighbor_kill_lookup[dir_idx, cv, ss] = True

    d_ignition_lookup        = cuda.to_device(ignition_lookup)
    d_kill_lookup            = cuda.to_device(kill_lookup)
    d_survive_lookup         = cuda.to_device(survive_lookup)
    d_neighbor_ignite_lookup = cuda.to_device(neighbor_ignite_lookup)
    d_neighbor_kill_lookup   = cuda.to_device(neighbor_kill_lookup)

update_lookups()

# ====================== CUDA SETUP ======================
ZOOM_LEVELS = [1, 4, 9, 16, 25]
current_zoom_idx = 0
pixels_per_cell = ZOOM_LEVELS[current_zoom_idx]

viewport_x_float = 0.0
viewport_y_float = 0.0
viewport_x = 0
viewport_y = 0

dragging = False
last_mouse_pos = (0, 0)

d_current = cuda.device_array((GRID_HEIGHT, GRID_WIDTH), dtype=np.uint16)
d_next    = cuda.device_array((GRID_HEIGHT, GRID_WIDTH), dtype=np.uint16)
d_visible = cuda.device_array((GRID_HEIGHT, GRID_WIDTH, 3), dtype=np.uint8)

SCALE_FACTOR = 255.0 / MAX_VALUE
d_scale_factor = cuda.to_device(np.array([SCALE_FACTOR], dtype=np.float32))

threads_per_block = (16, 16)
blocks = ((GRID_HEIGHT + 15) // 16, (GRID_WIDTH + 15) // 16)

meta_enabled = True
neighbor_ignite_enabled = True
neighbor_kill_enabled   = True

@cuda.jit
def update_kernel(current, next_state,
                  ignition_lookup, kill_lookup, survive_lookup,
                  meta_rule, max_val, meta_enabled,
                  neighbor_ignite_lookup, neighbor_ignite_enabled,
                  neighbor_kill_lookup,   neighbor_kill_enabled):
    i, j = cuda.grid(2)
    if i >= current.shape[0] or j >= current.shape[1]:
        return

    neighbor_sum = 0
    living_count = 0
    for di in range(-1, 2):
        for dj in range(-1, 2):
            if di == 0 and dj == 0:
                continue
            ni = (i + di) % current.shape[0]
            nj = (j + dj) % current.shape[1]
            val = current[ni, nj]
            neighbor_sum += val
            if val > 0:
                living_count += 1

    # Meta decay calculation
    rule_triggering_neighbors = 0
    if meta_enabled:
        for di in range(-1, 2):
            for dj in range(-1, 2):
                if di == 0 and dj == 0:
                    continue
                ni = (i + di) % current.shape[0]
                nj = (j + dj) % current.shape[1]

                nn_sum = 0
                nn_living = 0
                for ddi in range(-1, 2):
                    for ddj in range(-1, 2):
                        if ddi == 0 and ddj == 0:
                            continue
                        nni = (ni + ddi) % current.shape[0]
                        nnj = (nj + ddj) % current.shape[1]
                        v = current[nni, nnj]
                        nn_sum += v
                        if v > 0:
                            nn_living += 1

                if ignition_lookup[nn_living, nn_sum] or kill_lookup[nn_living, nn_sum]:
                    rule_triggering_neighbors += 1

    this_cell_decay = 1
    if meta_enabled:
        idx = min(rule_triggering_neighbors, meta_rule.shape[0] - 1)
        this_cell_decay = meta_rule[idx]

    val = current[i, j]

    # Neighbor ignition check
    ignited_by_neighbor = False
    if neighbor_ignite_enabled:
        dir_idx = 0
        for di in [-1,0,1]:
            for dj in [-1,0,1]:
                if di == 0 and dj == 0:
                    continue
                ni = (i + di) % current.shape[0]
                nj = (j + dj) % current.shape[1]

                second_sum = 0
                for ddi in [-1,0,1]:
                    for ddj in [-1,0,1]:
                        if ddi == 0 and ddj == 0:
                            continue
                        nni = (ni + ddi) % current.shape[0]
                        nnj = (nj + ddj) % current.shape[1]
                        second_sum += current[nni, nnj]

                if second_sum < neighbor_ignite_lookup.shape[2]:
                    if neighbor_ignite_lookup[dir_idx, val, second_sum]:
                        ignited_by_neighbor = True
                dir_idx += 1

    # Neighbor kill check
    killed_by_neighbor = False
    if neighbor_kill_enabled:
        dir_idx = 0
        for di in [-1,0,1]:
            for dj in [-1,0,1]:
                if di == 0 and dj == 0:
                    continue
                ni = (i + di) % current.shape[0]
                nj = (j + dj) % current.shape[1]

                second_sum = 0
                for ddi in [-1,0,1]:
                    for ddj in [-1,0,1]:
                        if ddi == 0 and ddj == 0:
                            continue
                        nni = (ni + ddi) % current.shape[0]
                        nnj = (nj + ddj) % current.shape[1]
                        second_sum += current[nni, nnj]

                if second_sum < neighbor_kill_lookup.shape[2]:
                    if neighbor_kill_lookup[dir_idx, val, second_sum]:
                        killed_by_neighbor = True
                dir_idx += 1

    # Final decision (order: neighbor ignite > neighbor kill > classic rules)
    if ignited_by_neighbor:
        next_state[i, j] = max_val
    elif killed_by_neighbor:
        next_state[i, j] = 0
    elif ignition_lookup[living_count, neighbor_sum]:
        next_state[i, j] = max_val
    elif kill_lookup[living_count, neighbor_sum]:
        next_state[i, j] = 0
    elif survive_lookup[living_count, neighbor_sum]:
        next_state[i, j] = val
    else:
        next_state[i, j] = max(0, val - this_cell_decay)


@cuda.jit
def render_kernel(green, rgb, scale_factor_array):
    i, j = cuda.grid(2)
    if i < green.shape[0] and j < green.shape[1]:
        v = green[i, j]
        scaled = int(v * scale_factor_array[0] + 0.5)
        scaled = min(255, max(0, scaled))
        rgb[i, j, 0] = 0
        rgb[i, j, 1] = scaled
        rgb[i, j, 2] = 0


# ── Grid seeding ─────────────────────────────────────────────────
def create_centered_circle():
    np.random.seed(42)
    y, x = np.ogrid[:GRID_HEIGHT, :GRID_WIDTH]
    dist = np.sqrt((x - GRID_WIDTH//2)**2 + (y - GRID_HEIGHT//2)**2)
    falloff = np.clip(1 - dist / CIRCLE_RADIUS, 0, 1)
    return (falloff ** FALLOFF_EXPONENT * CIRCLE_INTENSITY).astype(np.uint16)


def create_noise_seed(density=NOISE_DENSITY):
    grid = np.zeros((GRID_HEIGHT, GRID_WIDTH), dtype=np.uint16)
    avg_cluster_size = 5.0
    num_clusters = int(GRID_HEIGHT * GRID_WIDTH * density / avg_cluster_size)
    possible_rows = np.arange(1, GRID_HEIGHT - 1)
    possible_cols = np.arange(1, GRID_WIDTH - 1)
    num_possible = len(possible_rows) * len(possible_cols)
    indices = np.random.choice(num_possible, num_clusters, replace=False)
    centers_rows = possible_rows[indices // len(possible_cols)]
    centers_cols = possible_cols[indices % len(possible_cols)]
    deltas = np.array([[-1,-1],[-1,0],[-1,1],[0,-1],[0,1],[1,-1],[1,0],[1,1]])
    for clus in range(num_clusters):
        cr = centers_rows[clus]
        cc = centers_cols[clus]
        k = np.random.randint(0, 9)
        grid[cr, cc] = np.random.randint(1, MAX_VALUE + 1)
        if k > 0:
            delta_idx = np.random.choice(8, k, replace=False)
            for dr, dc in deltas[delta_idx]:
                grid[cr + dr, cc + dc] = np.random.randint(1, MAX_VALUE + 1)
    return grid


# Initial seed
d_current.copy_to_device(create_centered_circle())

noise_mode = False

# ── Main loop ────────────────────────────────────────────────────
pygame.init()
screen = pygame.display.set_mode((GRID_WIDTH, GRID_HEIGHT), pygame.RESIZABLE)
pygame.display.set_caption("Ignition Sim – Neighbor Ignite & Kill")
clock = pygame.time.Clock()
small_font = pygame.font.SysFont("consolas", 24)

rule_flash_frames = 0
MAX_FLASH_FRAMES = 90

running = True

while running:
    mouse_x, mouse_y = pygame.mouse.get_pos()

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

        elif event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 1 and pixels_per_cell > 1:
                dragging = True
                last_mouse_pos = event.pos

        elif event.type == pygame.MOUSEBUTTONUP:
            if event.button == 1:
                dragging = False

        elif event.type == pygame.MOUSEMOTION:
            if dragging:
                dx = event.pos[0] - last_mouse_pos[0]
                dy = event.pos[1] - last_mouse_pos[1]
                viewport_x_float -= dx / pixels_per_cell
                viewport_y_float -= dy / pixels_per_cell
                view_w = GRID_WIDTH // pixels_per_cell
                view_h = GRID_HEIGHT // pixels_per_cell
                viewport_x_float = max(0, min(viewport_x_float, GRID_WIDTH - view_w))
                viewport_y_float = max(0, min(viewport_y_float, GRID_HEIGHT - view_h))
                viewport_x = int(viewport_x_float)
                viewport_y = int(viewport_y_float)
                last_mouse_pos = event.pos

        elif event.type == pygame.MOUSEWHEEL:
            if event.y != 0:
                old_ppc = pixels_per_cell
                if event.y > 0:
                    current_zoom_idx = min(current_zoom_idx + 1, len(ZOOM_LEVELS) - 1)
                else:
                    current_zoom_idx = max(current_zoom_idx - 1, 0)
                pixels_per_cell = ZOOM_LEVELS[current_zoom_idx]
                if pixels_per_cell != old_ppc:
                    world_x = viewport_x_float + mouse_x / old_ppc
                    world_y = viewport_y_float + mouse_y / old_ppc
                    viewport_x_float = world_x - mouse_x / pixels_per_cell
                    viewport_y_float = world_y - mouse_y / pixels_per_cell
                    view_w = GRID_WIDTH // pixels_per_cell
                    view_h = GRID_HEIGHT // pixels_per_cell
                    viewport_x_float = max(0, min(viewport_x_float, GRID_WIDTH - view_w))
                    viewport_y_float = max(0, min(viewport_y_float, GRID_HEIGHT - view_h))
                    viewport_x = int(viewport_x_float)
                    viewport_y = int(viewport_y_float)

        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_r:
                d_current.copy_to_device(create_noise_seed() if noise_mode else create_centered_circle())
            elif event.key == pygame.K_g:
                seed = time.time_ns()
                random.seed(seed)
                IGNITION_SUMS, KILL_SUMS, SURVIVE_SUMS = generate_random_rules(
                    current_num_ignites, current_num_kills, current_num_survives)
                NEIGHBOR_IGNITE_CONDITIONS = generate_neighbor_conditions(current_num_neighbor_ignites)
                NEIGHBOR_KILL_CONDITIONS   = generate_neighbor_conditions(current_num_neighbor_kills)
                update_lookups()
                rule_flash_frames = MAX_FLASH_FRAMES
            elif event.key == pygame.K_n:
                noise_mode = not noise_mode
                d_current.copy_to_device(create_noise_seed() if noise_mode else create_centered_circle())
            elif event.key == pygame.K_q:
                random.seed()
                current_num_ignites  = [random.randint(RANDOM_RULE_COUNT_MIN_IGNITE[i],  RANDOM_RULE_COUNT_MAX_IGNITE[i])  for i in range(8)]
                current_num_kills    = [random.randint(RANDOM_RULE_COUNT_MIN_KILL[i],    RANDOM_RULE_COUNT_MAX_KILL[i])    for i in range(8)]
                current_num_survives = [random.randint(RANDOM_RULE_COUNT_MIN_SURVIVE[i], RANDOM_RULE_COUNT_MAX_SURVIVE[i]) for i in range(8)]
                current_num_neighbor_ignites = [random.randint(RANDOM_RULE_COUNT_MIN_NEIGHBOR_IGNITE[i], RANDOM_RULE_COUNT_MAX_NEIGHBOR_IGNITE[i]) for i in range(8)]
                current_num_neighbor_kills   = [random.randint(RANDOM_RULE_COUNT_MIN_NEIGHBOR_KILL[i],   RANDOM_RULE_COUNT_MAX_NEIGHBOR_KILL[i])   for i in range(8)]
                IGNITION_SUMS, KILL_SUMS, SURVIVE_SUMS = generate_random_rules(
                    current_num_ignites, current_num_kills, current_num_survives)
                NEIGHBOR_IGNITE_CONDITIONS = generate_neighbor_conditions(current_num_neighbor_ignites)
                NEIGHBOR_KILL_CONDITIONS   = generate_neighbor_conditions(current_num_neighbor_kills)
                update_lookups()
                rule_flash_frames = MAX_FLASH_FRAMES
            elif event.key == pygame.K_d:
                meta_enabled = not meta_enabled
                print("Meta decay enabled:", meta_enabled)
            elif event.key == pygame.K_h:
                neighbor_ignite_enabled = not neighbor_ignite_enabled
                print("Neighbor ignite enabled:", neighbor_ignite_enabled)
            elif event.key == pygame.K_j:
                neighbor_kill_enabled = not neighbor_kill_enabled
                print("Neighbor kill enabled:", neighbor_kill_enabled)
            elif event.key == pygame.K_p:
                print("Meta decay rule:         ", RULE_TRIGGERING_NEIGHBORS)
                print("Ignite counts:            ", current_num_ignites)
                print("Kill counts:              ", current_num_kills)
                print("Survive counts:           ", current_num_survives)
                print("Neighbor ignite counts:   ", current_num_neighbor_ignites)
                print("Neighbor kill counts:     ", current_num_neighbor_kills)
            elif event.key == pygame.K_ESCAPE:
                running = False

    # Simulation step
    update_kernel[blocks, threads_per_block](
        d_current, d_next,
        d_ignition_lookup, d_kill_lookup, d_survive_lookup,
        d_meta_rule, MAX_VALUE, meta_enabled,
        d_neighbor_ignite_lookup, neighbor_ignite_enabled,
        d_neighbor_kill_lookup,   neighbor_kill_enabled
    )
    d_current, d_next = d_next, d_current

    render_kernel[blocks, threads_per_block](d_current, d_visible, d_scale_factor)
    host_rgb = d_visible.copy_to_host()

    # Extract visible portion
    step = pixels_per_cell
    view_h = GRID_HEIGHT // step
    view_w = GRID_WIDTH // step

    y_start = viewport_y
    y_end = min(viewport_y + view_h, GRID_HEIGHT)
    x_start = viewport_x
    x_end = min(viewport_x + view_w, GRID_WIDTH)

    view = host_rgb[y_start:y_end, x_start:x_end, :]

    if view.shape[0] < view_h or view.shape[1] < view_w:
        padded = np.zeros((view_h, view_w, 3), dtype=np.uint8)
        padded[:view.shape[0], :view.shape[1]] = view
        view = padded

    upscaled = np.repeat(np.repeat(view, step, axis=0), step, axis=1)
    upscaled = upscaled[:GRID_HEIGHT, :GRID_WIDTH, :]

    surf = pygame.surfarray.make_surface(np.swapaxes(upscaled, 0, 1))
    screen.blit(surf, (0, 0))

    # UI
    y_pos = 20

    if rule_flash_frames > 0:
        screen.blit(small_font.render("RULES CHANGED!", True, (255,120,255)), (20, y_pos))
        rule_flash_frames -= 1

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

pygame.quit()