mutable struct Bird
    y::Float64
    v::Float64
    s::Int
    a::Bool
end
mutable struct Pipe
    x::Float64
    gy::Int
    p::Bool
end

const GP = [40.0, 15.0, 0.4, -2.0, 4.0, 6.0, 1.2]
function init_game()
    Bird(GP[2] ÷ 2, 0.0, 0, true), Pipe[]
end

function update_bird!(b, jump)
    b.v = jump ? GP[4] : b.v + GP[3]
    b.y += b.v
    if b.y < 1
        b.y = 1; b.v = 0
    elseif b.y > GP[2]
        b.a = false
    end
end

function update_pipes!(pipes, b)
    filter!(p -> p.x > -GP[5], pipes)
    for p in pipes
        p.x -= GP[7]
        if !p.p && p.x < 5
            p.p = true
            b.s += 1
        end
    end
    if isempty(pipes) || pipes[end].x < GP[1] - 20
        push!(pipes, Pipe(GP[1], rand(4:Int(GP[2]) - Int(GP[6]) - 2), false))
    end
end

function check_collision(b, pipes)
    by = round(Int, b.y)
    for p in pipes
        if 5 >= p.x && 5 <= p.x + GP[5] - 1
            if by < p.gy || by > p.gy + GP[6] - 1
                b.a = false
                return
            end
        end
    end
end

function extract_state(b, pipes, state_type=:basic)
    closest = nothing
    for p in pipes
        if !p.p && p.x + GP[5] >= 5
            closest = p
            break
        end
    end
    
    if closest === nothing
        if state_type == :conv
            state = zeros(Float64, Int(GP[2]), Int(GP[1]))
            state[Int(GP[2])÷2, 5] = 1.0
            return state
        else
            return zeros(Float64, 5)
        end
    end
    
    if state_type == :conv
        state = zeros(Float64, Int(GP[2]), Int(GP[1]))
        by = clamp(round(Int, b.y), 1, Int(GP[2]))
        state[by, 5] = 1.0
        
        for p in pipes
            px = round(Int, p.x)
            start_x = max(1, px)
            end_x = min(Int(GP[1]), px + Int(GP[5]) - 1)
            
            if start_x <= end_x
                for x in start_x:end_x
                    for y in 1:Int(GP[2])
                        if y < p.gy || y > p.gy + Int(GP[6]) - 1
                            state[y, x] = 0.5
                        end
                    end
                end
            end
        end
        return state
    else
        return [
            b.y / GP[2],
            clamp(b.v / 5.0, -1.0, 1.0),
            (closest.x - 5) / GP[1],
            closest.gy / GP[2],
            (closest.gy + GP[6]) / GP[2]
        ]
    end
end

function run_episode(network, state_type=:basic; render=false, max_steps=1000)
    bird, pipes = init_game()
    step = 0
    
    while bird.a && step < max_steps
        state = extract_state(bird, pipes, state_type)
        
        if state_type == :conv
            action = conv_forward(network, state) > 0.5
        else
            action = dense_forward(network, state) > 0.5
        end
        
        update_bird!(bird, action)
        update_pipes!(pipes, bird)
        check_collision(bird, pipes)
        
        if render
            draw(bird, pipes)
            sleep(0.1)
        end
        
        step += 1
    end
    
    return bird.s
end

function draw(b, pipes)
    print("\033[2J\033[H")
    w, h = Int(GP[1]), Int(GP[2])
    println("="^w)
    println("FLAPPY BIRD AI | Счет: $(b.s)")
    println("="^w)
    
    for y in 1:h
        print("|")
        print(y == round(Int, b.y) ? "🤖" : " ")
        
        for x in 7:w-2
            is_pipe = false
            for p in pipes
                px = round(Int, p.x)
                if x >= px && x < px + Int(GP[5])
                    if y < p.gy || y > p.gy + Int(GP[6]) - 1
                        print("█")
                        is_pipe = true
                        break
                    end
                end
            end
            !is_pipe && print(" ")
        end
        println("|")
    end
    
    println("="^w)
    println("测试运行")
    !b.a && println("\n💀 Игра окончена! Счет: $(b.s)")
end

mutable struct DenseNet
    w1::Matrix{Float64}
    b1::Vector{Float64}
    w2::Matrix{Float64}
    b2::Vector{Float64}
end

mutable struct ConvNet
    conv_weights::Array{Float64,4}
    conv_bias::Vector{Float64}
    fc_weights::Matrix{Float64}
    fc_bias::Vector{Float64}
end

function init_dense()
    DenseNet(
        randn(8, 5) * 0.1,
        zeros(8),
        randn(1, 8) * 0.1,
        zeros(1)
    )
end
function init_conv()
    height, width = 15, 40
    filters = 4
    conv_size = (height - 2) * (width - 2) * filters
    ConvNet(
        randn(3, 3, 1, filters) * 0.1,
        zeros(filters),
        randn(1, conv_size) * 0.1,
        zeros(1)
    )
end

function dense_forward(net::DenseNet, state)
    h = tanh.(net.w1 * state .+ net.b1)
    logits = net.w2 * h .+ net.b2
    return 1.0 / (1.0 + exp(-logits[1]))
end

function conv_forward(net::ConvNet, state)
    height, width = size(state)
    filters = size(net.conv_weights, 4)
    
    conv_out = zeros(Float64, height-2, width-2, filters)
    
    for f in 1:filters
        for i in 1:height-2
            for j in 1:width-2
                sum_val = 0.0
                for ki in 1:3
                    for kj in 1:3
                        row = i + ki - 1
                        col = j + kj - 1
                        if row <= height && col <= width
                            sum_val += state[row, col] * net.conv_weights[ki, kj, 1, f]
                        end
                    end
                end
                conv_out[i, j, f] = max(sum_val + net.conv_bias[f], 0.0)
            end
        end
    end
    
    flattened = vec(conv_out)
    logits = net.fc_weights * flattened .+ net.fc_bias
    return 1.0 / (1.0 + exp(-logits[1]))
end

function evaluate_network(network, state_type=:basic, episodes=3)
    total = 0.0
    for _ in 1:episodes
        score = run_episode(network, state_type; render=false)
        total += score
    end
    return total / episodes
end

function train_genetic!(population, state_type=:basic; target_score=20, max_generations=100, mutation_rate=0.3)
    results = []
    best_score = 0.0
    best_network = nothing
    
    for generation in 1:max_generations
        scores = [evaluate_network(net, state_type) for net in population]
        gen_best = maximum(scores)
        gen_avg = sum(scores) / length(scores)
        gen_worst = minimum(scores)
        
        push!(results, (generation, gen_best, gen_avg, gen_worst))
        
        println("Поколение $generation | Лучший: ", round(gen_best, digits=2), " /平均: ", round(gen_avg, digits=2), " /最坏的: ", round(gen_worst, digits=2))
        
        if gen_best > best_score
            best_score = gen_best
            best_idx = argmax(scores)
            best_network = deepcopy(population[best_idx])
        end
        
        if gen_best >= target_score
            println("✓ Достигнут целевой счет $target_score")
            break
        end
        
        sorted_idx = sortperm(scores, rev=true)
        elite_size = max(1, length(population) ÷ 4)
        elites = population[sorted_idx[1:elite_size]]
              new_pop = deepcopy(elites)
        
        while length(new_pop) < length(population)
            parent = rand(elites)
            
            if state_type == :conv
                child = ConvNet(
                    copy(parent.conv_weights),
                    copy(parent.conv_bias),
                    copy(parent.fc_weights),
                    copy(parent.fc_bias)
                )
                
                if rand() < mutation_rate
                    child.conv_weights .+= randn(size(child.conv_weights)...) * 0.1
                    child.conv_bias .+= randn(size(child.conv_bias)...) * 0.1
                    child.fc_weights .+= randn(size(child.fc_weights)...) * 0.1
                    child.fc_bias .+= randn(size(child.fc_bias)...) * 0.1
                end
            else
                child = DenseNet(
                    copy(parent.w1),
                    copy(parent.b1),
                    copy(parent.w2),
                    copy(parent.b2)
                )
                
                if rand() < mutation_rate
                    child.w1 .+= randn(size(child.w1)...) * 0.1
                    child.b1 .+= randn(size(child.b1)...) * 0.1
                    child.w2 .+= randn(size(child.w2)...) * 0.1
                    child.b2 .+= randn(size(child.b2)...) * 0.1
                end
            end
            
            push!(new_pop, child)
        end
        
        population = new_pop[1:length(population)]
    end
    
    return population, results, best_network, best_score
end

function argmax(x)
    max_val = -Inf
    max_idx = 1
    for (i, val) in enumerate(x)
        if val > max_val
            max_val = val
            max_idx = i
        end
    end
    return max_idx
end

struct NetworkResult
    name::String
    architecture::String
    network::Any
    state_type::Symbol
    test_score::Float64
    training_stats::Vector{Any}
    params::Int
end

function main_menu()
    results = NetworkResult[]
    
    while true
        print("\033[2J\033[H")
        println("="^60)
        println("          飞扬的鸟-进化算法")
        println("="^60)
        
        if !isempty(results)
            println("\nОБУЧЕННЫЕ СЕТИ:")
            for (i, r) in enumerate(results)
                println("$i. $(r.name) ($(r.architecture)): ", round(r.test_score, digits=2), " | Параметров: $(r.params)")
            end
        end
        
        println("\n" * "-"^60)
        println("学习菜单:")
        println("1. 训练一个完全连接的网络")
        println("2. 训练卷积网络(CNN)")
        println("3. 使用图形进行测试运行")
        println("4. 比较所有训练好的网络")
        println("5. 出口；出口")
        println("-"^60)
        print("\nВыбор: ")
        
        choice = readline()
        
        if choice == "1"
            results = train_dense_menu(results)
        elseif choice == "2"
            results = train_conv_menu(results)
        elseif choice == "3"
            testing_menu(results)
        elseif choice == "4"
            compare_results(results)
        elseif choice == "5"
            println("出口；出口")
            return
        end
    end
end

function train_dense_menu(results)
    print("\033[2J\033[H")
    println("="^60)
    println("全连接网络训练（进化算法）")
    println("="^60)
    println("\nАрхитектура сети:")
    println("输入层：5个神经元（位置，速度，到管道的距离）")
    println("隐藏层：8个具有TANH激活功能的神经元")
    println("输出层：1个带乙状结肠的神经元（跳跃概率）")
    println("总参数: 5×8 + 8 + 8×1 + 1 = 57")
    
    print("\nВведите имя для сети: ")
    name = readline()
    
    print("人口中的网络数量: ")
    pop_size = parse(Int, readline())
    
    print("最大代数: ")
    max_generations = parse(Int, readline())
    
    print("止损目标账户: ")
    target_score = parse(Float64, readline())
    
    print("训练后的测试集: ")
    test_episodes = parse(Int, readline())
    
    println("\nСоздание популяции...")
    population = [init_dense() for _ in 1:pop_size]
       println("进化学习的开始。..")
    println("Цель: достичь счета $target_score")
    println("-"^60)
    
    final_pop, training_stats, best_net, best_score = train_genetic!(
        population, :basic, 
        target_score=target_score, 
        max_generations=max_generations,
        mutation_rate=0.3
    )
    
    println("\nТестовый прогон лучшей сети...")
    test_score = evaluate_network(best_net, :basic, test_episodes)
    
    push!(results, NetworkResult(
        name,
        "Полносвязная (5-8-1)",
        best_net,
        :basic,
        test_score,
        training_stats,
        57
    ))
    
    println("\nОбучение завершено!")
    println("最佳学习成绩: ", round(best_score, digits=2))
    println("测试帐户: ", round(test_score, digits=2))
    println("按Enter继续。..")
    readline()
    
    return results
end

function testing_menu(results)
    if isempty(results)
        println("没有训练有素的网络！")
        sleep(2)
        return
    end
    
    while true
        print("\033[2J\033[H")
        println("="^60)
        println("使用图形进行测试运行")
        println("="^60)
        println("\nВыберите сеть для тестирования (0 для выхода):")
        
        for (i, r) in enumerate(results)
            println("$i. $(r.name) ($(r.architecture)): ", round(r.test_score, digits=2))
        end
        
        print("\nНомер сети: ")
        input = readline()
        
        if input == "0"
            return
        end
        
        idx = tryparse(Int, input)
        if idx === nothing || idx < 1 || idx > length(results)
            println("打错了！")
            sleep(2)
            continue
        end
        
        chosen = results[idx]
        
        println("\nЗапуск демонстрации сети '$(chosen.name)'...")
        println("按Ctrl+C退出演示")
        println("3秒后开始。..")
        sleep(3)
        
        try
            score = run_episode(chosen.network, chosen.state_type; render=true, max_steps=1000)
            println("\nДемонстрация завершена! Счет: $score")
        catch e
            if !(e isa InterruptException)
                println("\nОшибка при выполнении: $e")
            end
        end
        
        println("\nНажмите Enter для продолжения...")
        readline()
    end
end

function train_conv_menu(results)
    print("\033[2J\033[H")
    println("="^60)
    println("卷积网络训练（CNN）是一种进化算法")
    println("="^60)
    println("\nАрхитектура сети:")
    println("入口：15×40（比赛场地的高度×宽度）")
    println("卷积层：4个3×3滤波器，ReLU激活")
    println("全连接层：1个具有乙状结肠的神经元")
    println("总参数:~500")
    
    print("\nВведите имя для сети: ")
    name = readline()
    
    print("人口中的网络数量: ")
    pop_size = parse(Int, readline())
    
    print("最大代数: ")
    max_generations = parse(Int, readline())
    
    print("止损目标账户: ")
    target_score = parse(Float64, readline())
    
    print("训练后的测试集: ")
    test_episodes = parse(Int, readline())
    
    println("\nСоздание популяции...")
    population = [init_conv() for _ in 1:pop_size]
    
    println("进化学习的开始。..")
    println("Цель: достичь счета $target_score")
    println("-"^60)
        final_pop, training_stats, best_net, best_score = train_genetic!(
        population, :conv, 
        target_score=target_score, 
        max_generations=max_generations,
        mutation_rate=0.3
    )
    
    println("\nТестовый прогон лучшей сети...")
    test_score = evaluate_network(best_net, :conv, test_episodes)
    
    height, width = 15, 40
    filters = 4
    conv_params = 3 * 3 * 1 * filters + filters
    fc_params = (height-2) * (width-2) * filters + 1
    total_params = conv_params + fc_params
    
    push!(results, NetworkResult(
        name,
        "Сверточная (Conv3x3×4 → FC)",
        best_net,
        :conv,
        test_score,
        training_stats,
        total_params
    ))
    
    println("\nОбучение завершено!")
    println("最佳学习成绩: ", round(best_score, digits=2))
    println("测试帐户: ", round(test_score, digits=2))
    println("按Enter继续。..")
    readline()
    
    return results
end

function compare_results(results)
    if isempty(results)
        println("没有结果可以比较！")
        sleep(2)
        return
    end
    
    print("\033[2J\033[H")
    println("="^60)
    println("训练网络的比较")
    println("="^60)
    
    sorted_results = sort(results, by=r->r.test_score, rev=true)
    
    println("\nРейтинг сетей (по тестовому счету):")
    println("-"^60)
    println("地点/名称/架构|帐户/参数")
    println("-"^60)
    
    for (i, r) in enumerate(sorted_results)
        place = i == 1 ? "🥇" : i == 2 ? "🥈" : i == 3 ? "🥉" : "  $i."
        println("$place $(r.name) | $(r.architecture) | ", round(r.test_score, digits=2), " | $(r.params)")
    end
    
    if any(r -> !isempty(r.training_stats), sorted_results)
        println("\nСтатистика обучения лучшей сети:")
        best = sorted_results[1]
        if !isempty(best.training_stats)
            println("Количество поколений: $(length(best.training_stats))")
            last_stat = best.training_stats[end]
            first_stat = best.training_stats[1]
            println("初始平均分: ", round(first_stat[3], digits=2))
            println("最后最好成绩: ", round(last_stat[2], digits=2))
            println("改善工程: ", round(last_stat[2] - first_stat[2], digits=2))
        end
    end
    
    println("\nНажмите Enter для продолжения...")
    readline()
end