peloton_complete.jl

using StatsBase
using Plots

mutable struct Agent
    position::Vector{Real}
    cohere_factor::Real
    repel_factor::Real
    energy::Real
    desired_rank::Real
end

struct Model
    peloton::Vector{Agent}
    road_center::Real
end

function Agent()
    position = 0.1 * rand(2) .+ [0.4; 0.4]
    energy = 1
    desired_rank = rand()
    cohere_factor = 0.005 + 0.01 * rand()
    repel_factor = 0.005 + 0.01 * rand()
    Agent(position, cohere_factor, repel_factor, energy, desired_rank)
end

function Model(n_agents::Integer)
    peloton = [ Agent() for i ∈ 1:n_agents ]
    Model(peloton, 0.5)
end

function flocking(agent::Agent, model::Model)
    N = 0
    repel = [0.0; 0.0]
    cohere = [0.0; 0.0]
    
    for other_agent in model.peloton            
        if agent != other_agent
            N += 1
            heading = other_agent.position - agent.position
            distance = sum((heading).^2)^0.5
            
            cohere += heading
            
            if distance < 0.1 && heading[1] > 0
                angle_factor = abs(acot(heading[2] / heading[1])) / π * 2
                repel -= heading ./ distance^2 * angle_factor
            end
        end
    end
    velocity =  agent.cohere_factor * cohere / N + agent.repel_factor * repel / N
    return velocity
end

function energy_usage(agent::Agent, model::Model)
    usage = 0.02
    
    for other_agent in model.peloton
        heading = other_agent.position - agent.position
        distance = sum((heading).^2)^0.5
        
        if agent != other_agent          
            if distance < 0.2 && heading[1] > 0
                angle_factor = abs(acot(heading[2] / heading[1])) / π * 2
                usage -= 0.1 * (0.2 - distance) * angle_factor
            end
        end
    end

    return min(max(-0.01, usage), 0.02)
end

function update_desired_rank!(agent::Agent, position_mean)
    if rand() > agent.energy
            agent.desired_rank = agent.desired_rank - 0.01 * rand()
    else 
        if rand() > 0.99
            agent.desired_rank = agent.desired_rank + max(0, (0.5 - position_mean[1])) * rand()
        end
    end
end

function step!(model::Model)
    positions = [ agent.position for agent in model.peloton ]
    
    position_mean = mean(positions)
    for agent in model.peloton
        # random movement
        random_move =  0.001 *(rand(2) .- 0.5)

        # bias towards center of the road
        bias = [0 ; 0.05 * (model.road_center - agent.position[2])]

        # bias to desired rank
        bias[1] = 0.05 * (agent.desired_rank - agent.position[1])
        
        # energy component
        velocity = flocking(agent, model) + random_move + bias
        agent.position += velocity
        
        agent.energy -= energy_usage(agent, model)
        agent.energy = max(0, min(1, agent.energy))
        
        update_desired_rank!(agent, position_mean)        
    end
end

model = Model(50)

@gif for t ∈ 1:1000            
    step!(model)
    
    # plotting
    postions = [ agent.position for agent in model.peloton ]
    arr = hcat(postions...)
    
    energy = [ agent.energy for agent in model.peloton ]   
    scatter(arr[1,:], arr[2,:], xlims=(-0.2, 1.), ylims=(0,1), marker_z=energy, color=:RdYlGn_9, clims=(0., 1), clabel="Energy")
end