# Моделирование полёта дронов
# с избеганием столкновений

using LinearAlgebra
using Plots
using Printf
using Random
plotlyjs() # Используем plotlyjs для статических графиков
gr()       # Используем gr для анимации

# --- 1. Генерация неравномерного ландшафта ---

"""
    generate_landscape(width_m, length_m; seed=nothing)

Генерирует 3D-ландшафт с неравномерным рельефом.
"""
function generate_landscape(width_m, length_m; seed=nothing)
    if seed !== nothing
        Random.seed!(seed)
    end
    x = 0:1:width_m
    y = 0:1:length_m

    z = zeros(length(x), length(y))
    for (i, xi) in enumerate(x)
        for (j, yi) in enumerate(y)
            # Комбинируем несколько слоев шума для неравномерности
            z1 = 15 * sin(xi / 20) * cos(yi / 25)
            z2 = 10 * sin(xi / 10 + yi / 15)
            z3 = 5 * cos(xi / 30) * sin(yi / 20)
            z4 = 3 * (rand() - 0.5) # Мелкий шум
            z[i, j] = z1 + z2 + z3 + z4 + 10 # Базовая высота
        end
    end
    z = reverse(z, dims=(1, 2))

    plt = surface(x, y, z', c=:terrain, xlabel="X (м)", ylabel="Y (м)", zlabel="Z (м)", title="3D Ландшафт", legend=false)
    display(plt)
    return x, y, z
end

# --- 2. Генерация дронов ---

"""
    is_position_valid(x_pos, y_pos, z_pos, x_grid, y_grid, z_grid, min_altitude=1.0)

Проверяет, является ли позиция над ландшафтом и в допустимых границах.
"""
function is_position_valid(x_pos, y_pos, z_pos, x_grid, y_grid, z_grid, min_altitude=1.0)
    if x_pos < x_grid[1] || x_pos > x_grid[end] || y_pos < y_grid[1] || y_pos > y_grid[end]
        return false
    end
    xi = argmin(abs.(x_grid .- x_pos))
    yi = argmin(abs.(y_grid .- y_pos))
    terrain_height = z_grid[xi, yi]
    return z_pos >= (terrain_height + min_altitude)
end

"""
    generate_drones(N, x_grid, y_grid, z_grid, target_base=nothing)

Генерирует N дронов в случайных уникальных точках и определяет общую базу.
"""
function generate_drones(N, x_grid, y_grid, z_grid, target_base=nothing)
    drones = Vector{Dict{Symbol, Any}}()
    used_positions = Set{Tuple{Float64, Float64, Float64}}()
    min_dist_between_drones = 5.0 # Минимальное расстояние между стартовыми точками
    min_altitude = 2.0
    max_attempts = N * 100 # Ограничение попыток для избежания бесконечного цикла

    # Определяем общую базу (Target)
    if target_base === nothing
        base_x = x_grid[div(length(x_grid), 2)]
        base_y = y_grid[div(length(y_grid), 2)]
        base_xi = argmin(abs.(x_grid .- base_x))
        base_yi = argmin(abs.(y_grid .- base_y))
        base_z = z_grid[base_xi, base_yi] + min_altitude + 5.0 # Немного выше рельефа
        target_pos = [base_x, base_y, base_z]
    else
        target_pos = target_base
    end

    for i in 1:N
        pos = nothing
        attempts = 0
        while pos === nothing && attempts < max_attempts
            attempts += 1
            # Генерируем случайные координаты
            rand_x = x_grid[1] + (x_grid[end] - x_grid[1]) * rand()
            rand_y = y_grid[1] + (y_grid[end] - y_grid[1]) * rand()
            rand_z = 5.0 + (20.0 - 5.0) * rand() # Высота от 5 до 20 метров над "нулем"
            
            # Проверяем валидность позиции
            if is_position_valid(rand_x, rand_y, rand_z, x_grid, y_grid, z_grid, min_altitude)
                # Проверяем минимальное расстояние до уже созданных дронов
                is_too_close = false
                for existing_drone in drones
                    existing_pos = existing_drone[:start_pos]
                    dist = norm([rand_x, rand_y, rand_z] .- existing_pos)
                    if dist < min_dist_between_drones
                        is_too_close = true
                        break
                    end
                end
                if !is_too_close
                    pos = [rand_x, rand_y, rand_z]
                end
            end
        end

        if pos !== nothing
            drone = Dict(
                :id => i,
                :start_pos => pos,
                :target_pos => target_pos, # Общая цель для полета туда
                :home_pos => pos,          # Дом для полета обратно
                :position => copy(pos),
                :velocity => zeros(3),
                :path => [copy(pos)],
                :status => "flying_to_target", # flying_to_target, flying_to_home, landed
                :color => palette(:tab10)[i] # Уникальный цвет
            )
            push!(drones, drone)
            println("Дрон $i создан в точке $(round.(pos, digits=2))")
        else
            println("Не удалось создать дрон $i после $max_attempts попыток.")
        end
    end

    if length(drones) < N
        println("Предупреждение: создано только $(length(drones)) из $N дронов.")
    end

    println("Общая база (цель): $(round.(target_pos, digits=2))")
    return drones, target_pos
end

# --- 3. Динамика дрона ---

"""
    update_drone_physics!(drone, target, x_grid, y_grid, z_grid, dt=1.0;
                          max_speed=1.5, max_acceleration=0.3, max_vertical_speed=1.0,
                          turn_rate_deg=15.0, min_altitude=2.0)
                          
Обновляет физику дрона (позиция, скорость) с учетом динамики.
"""
function update_drone_physics!(drone, target, x_grid, y_grid, z_grid, dt=1.0;
                              max_speed=1.5, max_acceleration=0.3, max_vertical_speed=1.0,
                              turn_rate_deg=15.0, min_altitude=2.0)
    
    current_pos = drone[:position]
    velocity = drone[:velocity]
    turn_rate_rad = deg2rad(turn_rate_deg)

    # --- Желаемое направление и скорость ---
    direction_to_target_xy = normalize(target[1:2] .- current_pos[1:2])
    desired_speed_xy = max_speed
    desired_velocity_xy = direction_to_target_xy * desired_speed_xy

    # --- Ограничение на поворот ---
    current_speed_xy = norm(velocity[1:2])
    if current_speed_xy > 1e-6
        current_direction_xy = velocity[1:2] / current_speed_xy
        dot_product = clamp(dot(current_direction_xy, direction_to_target_xy), -1.0, 1.0)
        angle_to_target = acos(dot_product)

        if angle_to_target > turn_rate_rad * dt # turn_rate_rad уже в рад/сек
            cross_product_z = current_direction_xy[1] * direction_to_target_xy[2] - current_direction_xy[2] * direction_to_target_xy[1]
            sign = cross_product_z > 0 ? 1 : -1
            angle_to_turn = turn_rate_rad * dt
            cos_turn = cos(angle_to_turn)
            sin_turn = sin(angle_to_turn) * sign
            new_vx = current_direction_xy[1] * cos_turn - current_direction_xy[2] * sin_turn
            new_vy = current_direction_xy[1] * sin_turn + current_direction_xy[2] * cos_turn
            desired_velocity_xy = [new_vx, new_vy] * current_speed_xy
        end
    else
        desired_velocity_xy = direction_to_target_xy * min(max_acceleration * dt, desired_speed_xy)
    end

    # --- Применение ускорения к горизонтальной скорости ---
    delta_v_xy = desired_velocity_xy .- velocity[1:2]
    acceleration_magnitude = norm(delta_v_xy) / dt # a = dv/dt
    if acceleration_magnitude > 1e-10
        unit_accel = delta_v_xy / norm(delta_v_xy)
        applied_accel_magnitude = min(acceleration_magnitude, max_acceleration)
        applied_accel = unit_accel * applied_accel_magnitude
        velocity[1:2] .+= applied_accel * dt
    end

    # --- Ограничение максимальной скорости ---
    speed_xy = norm(velocity[1:2])
    if speed_xy > max_speed
        velocity[1:2] .*= (max_speed / speed_xy)
    end

    # --- Вертикальное движение ---
    next_x = current_pos[1] + velocity[1] * dt
    next_y = current_pos[2] + velocity[2] * dt
    next_xi = argmin(abs.(x_grid .- next_x))
    next_yi = argmin(abs.(y_grid .- next_y))
    next_xi = clamp(next_xi, 1, length(x_grid))
    next_yi = clamp(next_yi, 1, length(y_grid))
    terrain_height_next = z_grid[next_xi, next_yi]

    desired_altitude = min_altitude
    desired_z = terrain_height_next + desired_altitude
    desired_vz = (desired_z - current_pos[3]) / dt # Пропорциональный регулятор
    desired_vz = clamp(desired_vz, -max_vertical_speed, max_vertical_speed)
    
    # Простая модель вертикального ускорения
    vz_change = desired_vz - velocity[3]
    max_vz_change = max_acceleration * dt
    vz_change = clamp(vz_change, -max_vz_change, max_vz_change)
    velocity[3] += vz_change

    # Ограничиваем вертикальную скорость
    velocity[3] = clamp(velocity[3], -max_vertical_speed, max_vertical_speed)

    # --- Вычисляем новую позицию ---
    new_pos = current_pos .+ velocity * dt
    new_pos[1] = clamp(new_pos[1], x_grid[1], x_grid[end])
    new_pos[2] = clamp(new_pos[2], y_grid[1], y_grid[end])
    # Вертикальная коррекция уже учтена в vz

    drone[:position] = new_pos
    push!(drone[:path], copy(new_pos))
end

# --- 4. Координация и избегание столкновений ---

"""
    avoid_collisions!(drones, separation_distance=8.0, avoidance_strength=0.5)

Простая логика избегания столкновений на уровне изменения желаемой скорости.
"""
function avoid_collisions!(drones, separation_distance=8.0, avoidance_strength=0.5)
    N = length(drones)
    if N < 2
        return
    end

    # Рассчитываем силы отталкивания для каждого дрона
    repulsion_forces = [zeros(3) for _ in 1:N]

    for i in 1:N
        for j in 1:N
            if i != j
                pos_i = drones[i][:position]
                pos_j = drones[j][:position]
                vec_ij = pos_i .- pos_j
                dist = norm(vec_ij)

                if 0.1 < dist < separation_distance # Избегаем деления на ноль и действуем на расстоянии
                    # Сила отталкивания обратно пропорциональна квадрату расстояния
                    force_magnitude = avoidance_strength / (dist^2 + 1e-6) 
                    force_direction = vec_ij / (dist + 1e-6)
                    repulsion_forces[i] .+= force_direction * force_magnitude
                end
            end
        end
    end

    # Применяем силы как импульсы к скорости
    for i in 1:N
        # Ограничиваем величину изменения скорости
        force_clamped = clamp(norm(repulsion_forces[i]), 0.0, 1.0) * 
                        (repulsion_forces[i] / (norm(repulsion_forces[i]) + 1e-6))
        drones[i][:velocity] .+= force_clamped
    end
end

# --- 5. Моделирование связи ---

"""
    calculate_connection_quality(drone1, drone2, x_grid, y_grid, z_grid;
                                max_range=30.0, path_checks=20)

Рассчитывает качество связи между двумя дронами.
"""
function calculate_connection_quality(drone1, drone2, x_grid, y_grid, z_grid;
                                    max_range=30.0, path_checks=20)
    pos1 = drone1[:position]
    pos2 = drone2[:position]
    
    dx = pos1[1] - pos2[1]
    dy = pos1[2] - pos2[2]
    dz = pos1[3] - pos2[3]
    distance = sqrt(dx^2 + dy^2 + dz^2)

    # 1. Затухание сигнала с расстоянием
    if distance > max_range
        return 0.0
    end
    quality = 1.0 - (distance / max_range)

    # 2. Проверка на препятствия (прямая видимость)
    if quality > 0.01 # Проверяем только если сигнал ещё не нулевой
        for j in 1:path_checks
            t = j / (path_checks + 1) # Не проверяем концы
            check_x = pos1[1] + (pos2[1] - pos1[1]) * t
            check_y = pos1[2] + (pos2[2] - pos1[2]) * t
            xi_check = argmin(abs.(x_grid .- check_x))
            yi_check = argmin(abs.(y_grid .- check_y))
            if 1 ≤ xi_check ≤ size(z_grid, 1) && 1 ≤ yi_check ≤ size(z_grid, 2)
                terrain_z = z_grid[xi_check, yi_check]
                # Линейная интерполяция высоты линии связи
                line_z = pos1[3] + (pos2[3] - pos1[3]) * t
                if terrain_z > line_z - 1.0 # Если рельеф выше линии связи (с запасом)
                    return 0.0 # Полная потеря сигнала
                end
            end
        end
    end
    return max(0.0, quality)
end

"""
    get_network_status(drones, x_grid, y_grid, z_grid)

Возвращает матрицу качества связи между всеми парами дронов.
"""
function get_network_status(drones, x_grid, y_grid, z_grid)
    N = length(drones)
    connection_matrix = zeros(N, N)
    for i in 1:N
        for j in 1:N
            if i != j
                q = calculate_connection_quality(drones[i], drones[j], x_grid, y_grid, z_grid)
                connection_matrix[i, j] = q
            end
        end
    end
    return connection_matrix
end

# --- 6. Основной цикл симуляции ---

"""
    simulate_drones(N, sim_steps=500)
"""
function simulate_drones(N, sim_steps=500)
    println("=== Генерация ландшафта ===")
    # --- Изменено: увеличен размер ландшафта ---
    x_grid, y_grid, z_grid = generate_landscape(150, 150)

    println("\n=== Генерация дронов ===")
    drones, target_base = generate_drones(N, x_grid, y_grid, z_grid)

    if length(drones) == 0
        println("Ошибка: Не создано ни одного дрона. Завершение симуляции.")
        return nothing, nothing, nothing, nothing, nothing
    end

    println("\n=== Начало симуляции ===")
    connection_history = []
    all_paths = []

    # Для визуализации стартов, базы и домов
    start_positions = [d[:start_pos] for d in drones]
    home_positions = [d[:home_pos] for d in drones]

    for step in 1:sim_steps
        if step % 50 == 0
            println("Шаг симуляции: $step")
        end

        all_landed = true
        for drone in drones
            if drone[:status] == "flying_to_target"
                update_drone_physics!(drone, drone[:target_pos], x_grid, y_grid, z_grid)
                # Проверка достижения цели
                dist_to_target = norm(drone[:position][1:2] .- drone[:target_pos][1:2])
                if dist_to_target < 2.0 # Порог достижения
                    println("  Дрон $(drone[:id]) достиг базы. Начинаю возврат домой.")
                    drone[:status] = "flying_to_home"
                end
                all_landed = false
            elseif drone[:status] == "flying_to_home"
                update_drone_physics!(drone, drone[:home_pos], x_grid, y_grid, z_grid)
                # Проверка возвращения домой
                dist_to_home = norm(drone[:position][1:2] .- drone[:home_pos][1:2])
                if dist_to_home < 2.0 # Порог возвращения
                    println("  Дрон $(drone[:id]) вернулся домой и приземлился.")
                    drone[:status] = "landed"
                    # Добавим последнюю точку на дом
                    push!(drone[:path], copy(drone[:home_pos]))
                else
                    all_landed = false
                end
            end # если landed, ничего не делаем
        end

        # Координация: избегание столкновений
        avoid_collisions!(drones, 10.0, 0.8)

        # Моделирование связи
        conn_matrix = get_network_status(drones, x_grid, y_grid, z_grid)
        push!(connection_history, conn_matrix)

        # Проверка завершения
        if all_landed
            println("Все дроны завершили полет на шаге $step.")
            break
        end
    end

    println("\n=== Симуляция завершена ===")
    
    # --- Вывод таблицы связи ---
    if length(drones) > 0
        final_conn_matrix = get_network_status(drones, x_grid, y_grid, z_grid)
        println("\n=== Таблица качества связи между дронами (в конце симуляции) ===")
        print("ID  ")
        for j in 1:length(drones)
            @printf("%6d ", drones[j][:id])
        end
        println()
        for i in 1:length(drones)
            @printf("%3d ", drones[i][:id])
            for j in 1:length(drones)
                if i == j
                    print("   -   ")
                else
                    @printf("%6.2f ", final_conn_matrix[i, j])
                end
            end
            println()
        end
        println("================================================================")
    end
    
    return x_grid, y_grid, z_grid, drones, connection_history
end

# --- 7. Визуализация ---

"""
    visualize_simulation(x_grid, y_grid, z_grid, drones, connection_history, target_base;
                        fps=5, draw_connections=true)
"""
function visualize_simulation(x_grid, y_grid, z_grid, drones, connection_history, target_base;
                            fps=5, draw_connections=true)
    
    println("\n=== Создание визуализации ===")
    
    if drones === nothing
        println("Нет данных для визуализации.")
        return
    end

    N = length(drones)
    # --- Обновлено: создаем два графика ---
    base_surface = surface(x_grid, y_grid, z_grid', c=:terrain, legend=false, alpha=0.7,
                           xlabel="X (м)", ylabel="Y (м)", zlabel="Z (м)", title="3D Вид")

    base_topdown = plot(xlabel="X (м)", ylabel="Y (м)", title="Вид сверху", aspect_ratio=:equal, legend=false)
    contour!(base_topdown, x_grid, y_grid, z_grid', c=:terrain, alpha=0.5, levels=10)
    
    # Отметим общую базу на обоих графиках
    scatter!(base_surface, [target_base[1]], [target_base[2]], [target_base[3]], 
             marker=:circle, markercolor=:black, markersize=8, label="База (B)")
    scatter!(base_topdown, [target_base[1]], [target_base[2]], 
             marker=:circle, markercolor=:black, markersize=8, label="База (B)")

    # Отметим стартовые позиции на обоих графиках
    for drone in drones
        scatter!(base_surface, [drone[:start_pos][1]], [drone[:start_pos][2]], [drone[:start_pos][3]], 
                 marker=:diamond, markercolor=drone[:color], markersize=6, label="Старт $(drone[:id])")
        scatter!(base_topdown, [drone[:start_pos][1]], [drone[:start_pos][2]], 
                 marker=:diamond, markercolor=drone[:color], markersize=6, label=false)
    end

    # Определим общее количество кадров
    max_path_length = maximum(length(drone[:path]) for drone in drones)
    
    # Предварительно вычислим пути для эффективности
    all_paths = [drone[:path] for drone in drones]

    anim = @animate for frame in 1:max_path_length
        plt3d = deepcopy(base_surface)
        plt2d = deepcopy(base_topdown)
        
        current_positions = []
        current_ids = []

        # Рисуем траектории и текущие позиции
        for (i, drone) in enumerate(drones)
            path = all_paths[i]
            path_len = length(path)
            
            if frame <= path_len
                # Текущая позиция
                current_pos = path[frame]
                push!(current_positions, current_pos)
                push!(current_ids, drone[:id])
                
                scatter!(plt3d, [current_pos[1]], [current_pos[2]], [current_pos[3]], 
                         marker=:circle, markercolor=drone[:color], markersize=6, label=false)
                scatter!(plt2d, [current_pos[1]], [current_pos[2]], 
                         marker=:circle, markercolor=drone[:color], markersize=6, label=false)
                
                # Рисуем пройденную часть пути
                if frame > 1
                    path_to_draw = path[1:frame]
                    plot!(plt3d, [p[1] for p in path_to_draw], [p[2] for p in path_to_draw], [p[3] for p in path_to_draw],
                          linecolor=drone[:color], linewidth=2, label=false)
                    plot!(plt2d, [p[1] for p in path_to_draw], [p[2] for p in path_to_draw],
                          linecolor=drone[:color], linewidth=2, label=false)
                end
            else
                # Если дрон уже закончил путь, показываем его финальную позицию
                final_pos = path[end]
                scatter!(plt3d, [final_pos[1]], [final_pos[2]], [final_pos[3]], 
                         marker=:circle, markercolor=drone[:color], markersize=6, label=false)
                scatter!(plt2d, [final_pos[1]], [final_pos[2]], 
                         marker=:circle, markercolor=drone[:color], markersize=6, label=false)
                plot!(plt3d, [p[1] for p in path], [p[2] for p in path], [p[3] for p in path],
                      linecolor=drone[:color], linewidth=2, label=false)
                plot!(plt2d, [p[1] for p in path], [p[2] for p in path],
                      linecolor=drone[:color], linewidth=2, label=false)
            end
        end

        # Рисуем линии связи
        if draw_connections && frame <= length(connection_history)
            conn_matrix = connection_history[frame]
            for i in 1:N
                for j in (i+1):N
                    if frame <= length(all_paths[i]) && frame <= length(all_paths[j])
                        q = conn_matrix[i, j]
                        if q > 0.1 # Рисуем только если есть связь
                            pos1 = all_paths[i][frame]
                            pos2 = all_paths[j][frame]
                            conn_color = :green
                            if q < 0.3
                                conn_color = :red
                            elseif q < 0.7
                                conn_color = :orange
                            end
                            plot!(plt3d, [pos1[1], pos2[1]], [pos1[2], pos2[2]], [pos1[3], pos2[3]],
                                  linecolor=conn_color, linewidth=2*q, alpha=q, label=false)
                            plot!(plt2d, [pos1[1], pos2[1]], [pos1[2], pos2[2]],
                                  linecolor=conn_color, linewidth=2*q, alpha=q, label=false)
                        end
                    end
                end
            end
        end

        # Фокусировка на первом дроне для начала, затем на центре группы
        # Для 3D графика
        #if length(current_positions) > 0
        #     focus_drone_pos = current_positions[1]
        #     xlims!(plt3d, focus_drone_pos[1] - 25, focus_drone_pos[1] + 25)
        #     ylims!(plt3d, focus_drone_pos[2] - 25, focus_drone_pos[2] + 25)
             # Для 2D графика также фокусируемся
        #     xlims!(plt2d, focus_drone_pos[1] - 25, focus_drone_pos[1] + 25)
        #     ylims!(plt2d, focus_drone_pos[2] - 25, focus_drone_pos[2] + 25)
        #end
        
        # --- Обновлено: используем layout для размещения графиков ---
        plot(plt3d, plt2d, layout=(1, 2), size=(1200, 600))
        
    end every 1 # Каждый кадр для плавности

    filename = "multi_drone_flight_simulation.gif"
    gif(anim, filename, fps=fps)
    println("Анимация сохранена в файл: $filename")
    return filename
end

# --- Основной блок выполнения ---
function main()
    NUMBER_OF_DRONES = 10 # <--- ИЗМЕНИТЕ ЭТО ЗНАЧЕНИЕ ДЛЯ ДРУГОГО КОЛИЧЕСТВА ДРОНОВ
    SIMULATION_STEPS = 1000

    x_grid, y_grid, z_grid, drones, connection_history = simulate_drones(NUMBER_OF_DRONES, SIMULATION_STEPS)
    if x_grid !== nothing
        # Найдем target_base из данных дронов для визуализации
        target_base = drones[1][:target_pos] 
        visualize_simulation(x_grid, y_grid, z_grid, drones, connection_history, target_base, fps=10)
    end
    println("\n=== Проект завершен ===")
end

# Запустить симуляцию
main()