neat snake test map

snake_neat.odin

package main

import "core:encoding/json"
import "core:fmt"
import "core:math"
import "core:math/rand"
import "core:os"
import "core:slice"
import "core:strings"
import "vendor:raylib"

// ============================================================================
// PART 1: NEAT LIBRARY
// ============================================================================

NodeType :: enum {
    Input,
    Hidden,
    Output,
    Bias,
}

NodeGene :: struct {
    id: int,
    type: NodeType,
    layer: f32,
    value: f64,
}

ConnectionGene :: struct {
    in_node: int,
    out_node: int,
    weight: f64,
    enabled: bool,
}

Genome :: struct {
    id: int,
    fitness: f64,
    inputs_count: int,
    outputs_count: int,

    nodes: map[int]NodeGene,
    connections: [dynamic]ConnectionGene,

    sorted_nodes: [dynamic]^NodeGene,
}

free_genome :: proc(g: ^Genome) {
    delete(g.nodes)
    delete(g.connections)
    delete(g.sorted_nodes)
}

genome_compile :: proc(g: ^Genome) {
    clear(&g.sorted_nodes)

    for _, &node in g.nodes {
        append(&g.sorted_nodes, &node)
    }

    slice.sort_by_cmp(g.sorted_nodes[:], proc(a, b: ^NodeGene) -> slice.Ordering {
        if a.layer < b.layer do return .Less
        if a.layer > b.layer do return .Greater
        return .Equal
    })
}

sigmoid :: proc(x: f64) -> f64 {
    return 1.0 / (1.0 + math.exp(-4.9 * x))
}

genome_feedforward :: proc(g: ^Genome, inputs: []f64) -> []f64 {
    for node_ptr in g.sorted_nodes {
        node_ptr.value = 0.0
    }

    for i in 0 ..< len(inputs) {
        if i in g.nodes {
            n := &g.nodes[i]
            n.value = inputs[i]
        }
    }

    for node in g.sorted_nodes {
        if node.type == .Input do continue

        total : f64 = 0
        for conn in g.connections {
            if conn.out_node == node.id && conn.enabled {
                if in_n, ok := g.nodes[conn.in_node]; ok {
                    total += in_n.value * conn.weight
                }
            }
        }

        node.value = sigmoid(total)
    }

    result := make([]f64, g.outputs_count, context.temp_allocator)
    out_idx := 0
    for n in g.sorted_nodes {
        if n.type == .Output {
            if out_idx < len(result) {
                result[out_idx] = n.value
                out_idx += 1
            }
        }
    }

    return result
}

// --- RANDOM GENERATION (NEW) ---

create_random_genome :: proc(id, inputs, outputs: int) -> ^Genome {
    g := new(Genome)
    g.id = id
    g.inputs_count = inputs
    g.outputs_count = outputs
    g.nodes = make(map[int]NodeGene)

    // 1. Create Inputs
    for i in 0 ..< inputs {
        g.nodes[i] = NodeGene{
            id = i,
            type = .Input,
            layer = 0.0,
            value = 0.0,
        }
    }

    // 2. Create Outputs
    // В NEAT-Python ID выходов обычно идут сразу после входов
    for i in 0 ..< outputs {
        node_id := inputs + i
        g.nodes[node_id] = NodeGene{
            id = node_id,
            type = .Output,
            layer = 1.0, // Output layer
            value = 0.0,
        }
    }

    // 3. Dense Connections (Full mesh)
    // Соединяем каждый вход с каждым выходом случайным весом
    for i in 0 ..< inputs {
        for j in 0 ..< outputs {
            out_id := inputs + j
            weight := rand.float64_range(-2.0, 2.0)

            append(&g.connections, ConnectionGene{
                in_node = i,
                out_node = out_id,
                weight = weight,
                enabled = true,
            })
        }
    }

    genome_compile(g)
    return g
}

create_random_population :: proc(size, inputs, outputs: int) -> [dynamic]^Genome {
    pop := make([dynamic]^Genome)
    for i in 0 ..< size {
        append(&pop, create_random_genome(i, inputs, outputs))
    }
    return pop
}

// --- SAVE SYSTEM (NEW) ---

save_population_to_file :: proc(pop: [dynamic]^Genome, filename: string) {
    sb := strings.builder_make()
    defer strings.builder_destroy(&sb)

    strings.write_string(&sb, "{\"population\": [\n")

    for g, i in pop {
        fmt.sbprintf(&sb, "  {{\n    \"key\": %d,\n    \"fitness\": %f,\n", g.id, g.fitness)

        // Nodes
        strings.write_string(&sb, "    \"nodes\": [\n")

        // Need to iterate map cleanly. Convert to slice for index control or just use counter
        node_count := len(g.nodes)
        k := 0
        for _, node in g.nodes {
            type_str := "hidden"
            if node.type == .Input do type_str = "input"
            if node.type == .Output do type_str = "output"

            fmt.sbprintf(&sb, "      {{\"id\": %d, \"type\": \"%s\", \"layer\": %f}}", node.id, type_str, node.layer)
            if k < node_count - 1 do strings.write_string(&sb, ",")
            strings.write_string(&sb, "\n")
            k += 1
        }
        strings.write_string(&sb, "    ],\n")

        // Connections
        strings.write_string(&sb, "    \"connections\": [\n")
        conn_count := len(g.connections)
        for c, j in g.connections {
            enabled_str := c.enabled ? "true" : "false"
            fmt.sbprintf(&sb, "      {{\"in\": %d, \"out\": %d, \"weight\": %f, \"enabled\": %s}}", c.in_node, c.out_node, c.weight, enabled_str)
            if j < conn_count - 1 do strings.write_string(&sb, ",")
            strings.write_string(&sb, "\n")
        }
        strings.write_string(&sb, "    ]\n")

        strings.write_string(&sb, "  }")
        if i < len(pop) - 1 do strings.write_string(&sb, ",")
        strings.write_string(&sb, "\n")
    }

    strings.write_string(&sb, "]}")

    os.write_entire_file(filename, transmute([]u8)strings.to_string(sb))
    fmt.println("Population SAVED to", filename)
}

// --- JSON Helpers ---

json_to_f64 :: proc(v: json.Value) -> f64 {
    switch t in v {
    case json.Float: return t
    case json.Integer: return f64(t)
    case json.Null, json.Boolean, json.String, json.Array, json.Object: return 0.0
    }
    return 0.0
}

json_to_int :: proc(v: json.Value) -> int {
    switch t in v {
    case json.Integer: return int(t)
    case json.Float: return int(t)
    case json.Null, json.Boolean, json.String, json.Array, json.Object: return 0
    }
    return 0
}

json_to_string :: proc(v: json.Value) -> string {
    switch t in v {
    case json.String: return t
    case json.Null, json.Integer, json.Float, json.Boolean, json.Array, json.Object: return ""
    }
    return ""
}

json_to_bool :: proc(v: json.Value) -> bool {
    switch t in v {
    case json.Boolean: return t
    case json.Null, json.Integer, json.Float, json.String, json.Array, json.Object: return false
    }
    return false
}

load_population_from_file :: proc(filename: string, inputs, outputs: int) -> [dynamic]^Genome {
    if !os.exists(filename) {
        return nil
    }

    data, ok := os.read_entire_file(filename)
    if !ok {
        return nil
    }
    defer delete(data)

    json_data, err := json.parse(data)
    if err != .None {
        return nil
    }
    defer json.destroy_value(json_data)

    root := json_data.(json.Object)

    pop_array: json.Array
    if "population" in root {
        pop_array = root["population"].(json.Array)
    } else if "pop" in root {
        pop_array = root["pop"].(json.Array)
    } else {
        return nil
    }

    genomes := make([dynamic]^Genome)

    for val in pop_array {
        g_obj := val.(json.Object)

        new_g := new(Genome)
        new_g.inputs_count = inputs
        new_g.outputs_count = outputs

        if "key" in g_obj {
            new_g.id = json_to_int(g_obj["key"])
        } else if "id" in g_obj {
            new_g.id = json_to_int(g_obj["id"])
        }

        if "fitness" in g_obj {
            new_g.fitness = json_to_f64(g_obj["fitness"])
        }

        new_g.nodes = make(map[int]NodeGene)
        nodes_arr := g_obj["nodes"].(json.Array)
        for n_val in nodes_arr {
            n_obj := n_val.(json.Object)

            id := json_to_int(n_obj["id"])
            type_str := json_to_string(n_obj["type"])
            layer := f32(json_to_f64(n_obj["layer"]))

            nt: NodeType
            switch type_str {
            case "input": nt = .Input
            case "output": nt = .Output
            case: nt = .Hidden
            }

            new_g.nodes[id] = NodeGene{
                id = id,
                type = nt,
                layer = layer,
                value = 0,
            }
        }

        conns_arr := g_obj["connections"].(json.Array)
        new_g.connections = make([dynamic]ConnectionGene, 0, len(conns_arr))
        for c_val in conns_arr {
            c_obj := c_val.(json.Object)

            conn := ConnectionGene{
                in_node  = json_to_int(c_obj["in"]),
                out_node = json_to_int(c_obj["out"]),
                weight   = json_to_f64(c_obj["weight"]),
                enabled  = json_to_bool(c_obj["enabled"]),
            }
            append(&new_g.connections, conn)
        }

        genome_compile(new_g)
        append(&genomes, new_g)
    }

    return genomes
}

// ============================================================================
// PART 2: SNAKE DEMO
// ============================================================================

CELL_SIZE :: 12
GRID_WIDTH :: 100
GRID_HEIGHT :: 100
SCREEN_W :: 1600
SCREEN_H :: 900

DEATH_SPEED :: 0.05
BASE_FOOD :: 50
SNAKE_COUNT :: 160

Direction :: enum {
    Up, Right, Down, Left
}

SnakeState :: enum {
    Alive,
    DyingShrink,
    DyingRot,
}

Point :: struct {
    x, y: int
}

DemoSnake :: struct {
    genome: ^Genome,
    body: [dynamic]Point,
    dir: Direction,
    stamina: int,
    color: raylib.Color,
    scale: f32,
    state: SnakeState,
    score: int,
}

Game :: struct {
    population: [dynamic]^Genome,
    snakes: [dynamic]DemoSnake,
    foods: [dynamic]Point,

    camera_pos: [2]f32,
    mode: int,
    msg_timer: f32,
}

init_snake :: proc(g: ^Genome) -> DemoSnake {
    pad :: 5
    sx := rand.int_max(GRID_WIDTH - 2 * pad) + pad
    sy := rand.int_max(GRID_HEIGHT - 2 * pad) + pad

    s := DemoSnake{
        genome  = g,
        dir     = .Up,
        stamina = 400,
        color   = raylib.Color{
            u8(rand.int_max(155) + 100),
            u8(rand.int_max(155) + 100),
            u8(rand.int_max(155) + 100),
            255,
        },
        scale   = 1.0,
        state   = .Alive,
        body    = make([dynamic]Point),
    }

    append(&s.body, Point{ sx, sy })
    append(&s.body, Point{ sx, sy + 1 })
    append(&s.body, Point{ sx, sy + 2 })
    return s
}

destroy_snake :: proc(s: ^DemoSnake) {
    delete(s.body)
}

die_snake :: proc(s: ^DemoSnake) {
    if s.state != .Alive do return

    if len(s.body) > 2 && rand.float32() < 0.3 {
        s.state = .DyingRot
        s.scale = 1.2
    } else {
        s.state = .DyingShrink
    }
}

update_visuals :: proc(s: ^DemoSnake, foods: ^[dynamic]Point) -> bool {
    if s.state == .Alive do return false

    if s.state == .DyingShrink {
        s.scale -= DEATH_SPEED
        if s.scale <= 0 do return true
    } else if s.state == .DyingRot {
        s.scale -= DEATH_SPEED * 2.0
        if s.scale <= 0.8 {
            for segment in s.body {
                exists := false
                for f in foods {
                    if f == segment {
                        exists = true; break
                    }
                }
                if !exists {
                    append(foods, segment)
                }
            }
            return true
        }
    }
    return false
}

think :: proc(s: ^DemoSnake, occupied: map[Point]bool, foods: [dynamic]Point) {
    if s.state != .Alive do return

    head := s.body[0]

    closest : Point = { -1, -1 }
    min_dist_sq := f32(999999.0)

    for f in foods {
        dx := f32(f.x - head.x)
        dy := f32(f.y - head.y)
        d_sq := dx * dx + dy * dy
        if d_sq < min_dist_sq {
            min_dist_sq = d_sq
            closest = f
        }
    }

    inputs := make([dynamic]f64, 0, 24, context.temp_allocator)

    // A. Danger
    dirs := [4]Point{ { 0, -1 }, { 1, 0 }, { 0, 1 }, { -1, 0 } }
    for d in dirs {
        nx, ny := head.x + d.x, head.y + d.y
        danger := 0.0
        if nx < 0 || nx >= GRID_WIDTH || ny < 0 || ny >= GRID_HEIGHT || (Point{ nx, ny } in occupied) {
            danger = 1.0
        }
        append(&inputs, danger)
    }

    // B. Food Direction
    if closest.x != -1 {
        dx := closest.x - head.x
        dy := closest.y - head.y
        append(&inputs, dy < 0 ? 1.0 : 0.0)
        append(&inputs, dx > 0 ? 1.0 : 0.0)
        append(&inputs, dy > 0 ? 1.0 : 0.0)
        append(&inputs, dx < 0 ? 1.0 : 0.0)
    } else {
        append(&inputs, 0, 0, 0, 0)
    }

    // C. Radar
    radars := [8]Point{ { 0, -1 }, { 1, -1 }, { 1, 0 }, { 1, 1 }, { 0, 1 }, { -1, 1 }, { -1, 0 }, { -1, -1 } }
    for r in radars {
        cx, cy := head.x, head.y
        dist := 0.0
        found := false
        for _ in 0 ..< 15 {
            cx += r.x
            cy += r.y
            dist += 1.0
            if cx < 0 || cx >= GRID_WIDTH || cy < 0 || cy >= GRID_HEIGHT || (Point{ cx, cy } in occupied) {
                found = true
                break
            }
        }
        append(&inputs, found ? 1.0 / dist : 0.0)
    }

    // D. Direction
    append(&inputs, s.dir == .Up    ? 1.0 : 0.0)
    append(&inputs, s.dir == .Right ? 1.0 : 0.0)
    append(&inputs, s.dir == .Down  ? 1.0 : 0.0)
    append(&inputs, s.dir == .Left  ? 1.0 : 0.0)

    // E. Coords
    append(&inputs, f64(head.y) / f64(GRID_HEIGHT))
    append(&inputs, f64(GRID_WIDTH - head.x) / f64(GRID_WIDTH))
    append(&inputs, f64(GRID_HEIGHT - head.y) / f64(GRID_HEIGHT))
    append(&inputs, f64(head.x) / f64(GRID_WIDTH))

    outputs := genome_feedforward(s.genome, inputs[:])

    best_idx := -1
    max_val := -999.0

    current_idx := int(s.dir)
    opposite := (current_idx + 2) % 4

    limit := min(4, len(outputs))
    for i in 0 ..< limit {
        if i == opposite do continue
        if outputs[i] > max_val {
            max_val = outputs[i]
            best_idx = i
        }
    }

    if best_idx != -1 {
        s.dir = Direction(best_idx)
    }
}

move_snake :: proc(s: ^DemoSnake, occupied: map[Point]bool, foods: ^[dynamic]Point) {
    if s.state != .Alive do return

    head := s.body[0]
    next_pos := head

    switch s.dir {
    case .Up: next_pos.y -= 1
    case .Right: next_pos.x += 1
    case .Down: next_pos.y += 1
    case .Left: next_pos.x -= 1
    }

    if next_pos.x < 0 || next_pos.x >= GRID_WIDTH ||
    next_pos.y < 0 || next_pos.y >= GRID_HEIGHT ||
    (next_pos in occupied) {
        die_snake(s)
        return
    }

    inject_at(&s.body, 0, next_pos)
    s.stamina -= 1

    ate := false
    for i in 0 ..< len(foods) {
        if foods[i] == next_pos {
            unordered_remove(foods, i)
            ate = true
            s.score += 1
            s.stamina += 150
            if s.stamina > 500 do s.stamina = 500
            break
        }
    }

    if !ate {
        pop(&s.body)
    }

    if s.stamina <= 0 {
        die_snake(s)
    }
}

restart_game :: proc(g: ^Game) {
    for &s in g.snakes {
        destroy_snake(&s)
    }
    clear(&g.snakes)
    clear(&g.foods)

    for _ in 0 ..< BASE_FOOD {
        append(&g.foods, Point{ rand.int_max(GRID_WIDTH), rand.int_max(GRID_HEIGHT) })
    }

    if g.mode == 1 {
        top_k := 3
        if len(g.population) < 3 do top_k = len(g.population)

        fmt.println("Restarting in ELITE mode")
        for i in 0 ..< SNAKE_COUNT {
            if len(g.population) > 0 {
                gen := g.population[i % top_k]
                append(&g.snakes, init_snake(gen))
            }
        }
    } else {
        fmt.println("Restarting in CHAOS mode")
        for i in 0 ..< SNAKE_COUNT {
            if len(g.population) > 0 {
                idx := rand.int_max(len(g.population))
                gen := g.population[idx]
                append(&g.snakes, init_snake(gen))
            }
        }
    }
}

// ============================================================================
// MAIN
// ============================================================================

main :: proc() {
    raylib.InitWindow(SCREEN_W, SCREEN_H, "NEAT Snake: Odin Implementation")
    defer raylib.CloseWindow()
    raylib.SetTargetFPS(60)

    game := new(Game)
    game.mode = 1

    // 1. Попытка загрузить файл
    fmt.println("Checking for neat_snake_battle.json...")
    loaded_pop := load_population_from_file("neat_snake_battle.json", 24, 4)

    if loaded_pop != nil && len(loaded_pop) > 0 {
        fmt.println("Loaded population from file.")
        game.population = loaded_pop
    } else {
        fmt.println("File not found or invalid. Creating RANDOM population.")
        // Создаем случайную популяцию (24 входа, 4 выхода)
        game.population = create_random_population(100, 24, 4)
    }

    // Сортировка по фитнесу (если он есть)
    slice.sort_by_cmp(game.population[:], proc(a, b: ^Genome) -> slice.Ordering {
        if a.fitness > b.fitness do return .Less
        if a.fitness < b.fitness do return .Greater
        return .Equal
    })

    restart_game(game)

    for !raylib.WindowShouldClose() {
        dt := raylib.GetFrameTime()
        if game.msg_timer > 0 do game.msg_timer -= dt

        // Inputs
        if raylib.IsMouseButtonDown(.LEFT) {
            delta := raylib.GetMouseDelta()
            game.camera_pos.x += delta.x
            game.camera_pos.y += delta.y
        }
        if raylib.IsKeyPressed(.ONE) {
            game.mode = 1; restart_game(game)
        }
        if raylib.IsKeyPressed(.TWO) {
            game.mode = 2; restart_game(game)
        }
        if raylib.IsKeyPressed(.R)   {
            restart_game(game)
        }

        // SAVE BUTTON
        if raylib.IsKeyPressed(.S) {
            save_population_to_file(game.population, "neat_snake_battle.json")
            game.msg_timer = 2.0
        }

        occupied := make(map[Point]bool, 1000, context.temp_allocator)
        alive_count := 0
        for &s in game.snakes {
            if s.state == .Alive {
                alive_count += 1
                for p in s.body {
                    occupied[p] = true
                }
            }
        }

        for i := len(game.snakes) - 1; i >= 0; i -= 1 {
            s := &game.snakes[i]

            if update_visuals(s, &game.foods) {
                destroy_snake(s)
                unordered_remove(&game.snakes, i)
                continue
            }

            think(s, occupied, game.foods)
            move_snake(s, occupied, &game.foods)
        }

        if len(game.foods) < BASE_FOOD / 2 {
            append(&game.foods, Point{ rand.int_max(GRID_WIDTH), rand.int_max(GRID_HEIGHT) })
        }

        if alive_count == 0 && len(game.snakes) == 0 {
            restart_game(game)
        }

        raylib.BeginDrawing()
        raylib.ClearBackground(raylib.Color{ 20, 20, 25, 255 })

        ox := i32(game.camera_pos.x) + 50
        oy := i32(game.camera_pos.y) + 50

        fw := i32(GRID_WIDTH * CELL_SIZE)
        fh := i32(GRID_HEIGHT * CELL_SIZE)
        raylib.DrawRectangle(ox, oy, fw, fh, raylib.Color{ 30, 30, 35, 255 })
        raylib.DrawRectangleLines(ox, oy, fw, fh, raylib.GRAY)

        time := raylib.GetTime()
        pulse := 1.0 + math.sin(time * 5.0) * 0.1

        for f in game.foods {
            size := i32(f32(CELL_SIZE) * f32(pulse))
            offset := (i32(CELL_SIZE) - size) / 2
            raylib.DrawRectangle(
            ox + i32(f.x) * CELL_SIZE + offset,
            oy + i32(f.y) * CELL_SIZE + offset,
            size, size, raylib.RED)
        }

        for &s in game.snakes {
            draw_size := i32(f32(CELL_SIZE) * s.scale)
            offset := (i32(CELL_SIZE) - draw_size) / 2

            col := s.color
            if s.state == .DyingShrink {
                col = raylib.Fade(col, s.scale)
            } else if s.state == .DyingRot {
                col = raylib.WHITE
            }

            for p, idx in s.body {
                px := ox + i32(p.x) * CELL_SIZE + offset
                py := oy + i32(p.y) * CELL_SIZE + offset

                final_col := col
                if idx == 0 && s.state == .Alive {
                    final_col = raylib.WHITE
                }

                raylib.DrawRectangle(px, py, draw_size, draw_size, final_col)
            }
        }

        mode_str := game.mode == 1 ? "ELITE (Top 3)" : "CHAOS (Random)"
        raylib.DrawText(fmt.ctprintf("MODE: %s", mode_str), 20, 20, 30, raylib.GOLD)
        raylib.DrawText("Press '1', '2' or 'R'", 20, 60, 20, raylib.GRAY)
        raylib.DrawText("Press 'S' to Save", 20, 85, 20, raylib.GRAY)
        raylib.DrawText(fmt.ctprintf("Alive: %d", alive_count), 20, 110, 20, raylib.WHITE)
        raylib.DrawFPS(SCREEN_W - 100, 10)

        if game.msg_timer > 0 {
            raylib.DrawText("SAVED!", SCREEN_W / 2 - 50, SCREEN_H / 2, 40, raylib.GREEN)
        }

        raylib.EndDrawing()

        free_all(context.temp_allocator)
    }

    for &s in game.snakes {
        destroy_snake(&s)
    }
    delete(game.snakes)
    delete(game.foods)
    for g in game.population {
        free_genome(g)
        free(g)
    }
    delete(game.population)
    free(game)
}