using Plots
using LinearAlgebra
using Random
Random.seed!(2018)
gr(size=(300,300),titlefont=font("sans-serif",9))


module Boids
    struct Param
        d::Int64 #dimention
        N::Int64 #Number of boid
        distThresh::Float64 #threshold distance of rule1
        maxSpeed::Float64 #maximum speed
        minSpeed::Float64 #minimum speed
        maxWorld::Float64 #world upper limit
        minWorld::Float64 #world lower limit

        #R::Float64
        ω1::Float64 #rule1 weight
        ω2::Float64 #rule2 weight
        ω3::Float64 #rule3 weight
    end
    mutable struct Variables
        p::Array{Float64,2} #Boid positions
        v::Array{Float64,2} #Boid velocities
        v1::Array{Float64,2} #rule1 output
        v2::Array{Float64,2} #rule2 output
        v3::Array{Float64,2} #rule3 output
    end
end

using .Boids
d=2
N=50

distThresh=10.0

#R=10.0 #radius
ω1=0.1
ω2=0.4
ω3=1/300

maxSpeed=4.0
minSpeed=1.0

minWorld=0
maxWorld=150

worldSize=abs(maxWorld-minWorld)

param1=Boids.Param(d,N,distThresh,maxSpeed,minSpeed,maxWorld,minWorld,ω1,ω2,ω3)

p=rand(N,d).*worldSize
v=(rand(N,d).-0.5)
#a=rand(N,d) #acceleration

v1=zeros(N,d)
v2=zeros(N,d)
v3=zeros(N,d)

var1=Boids.Variables(p,v,v1,v2,v3)

#clear output vector
function clear_vec(var)
    var.v1.=0.0
    var.v2.=0.0
    var.v3.=0.0
end

#rule1 Collision avoidance
function rule1(param,var)
    for i in 1:param.N
        for j in 1:param.N
            if i==j;continue;end
            dist=norm(var.p[j,:].-var.p[i,:])
            #If the distance to a boid j is shorter than constant distance
            if dist < param.distThresh
                var.v1[i,:].-=(var.p[j,:].-var.p[i,:])
            end
        end
    end
end

#rule2 Velocity Matching
function rule2(param,var)
    for i in 1:param.N
        for j in 1:param.N
            if i==j;continue;end
            var.v2[i,:].+=var.v[j,:]
        end
    end
    var.v2.=var.v2./(param.N-1)
    var.v2.=var.v2.-var.v
end

#rule3 Flock Centering
function rule3(param,var)
    for i in 1:param.N
        for j in 1:param.N
            if i==j;continue;end
            var.v3[i,:].+=var.p[j,:]
        end
    end
    var.v3.=var.v3./(param.N-1)
    var.v3.=var.v3.-var.p
end

#limit velocity
function limit_vel(param,var)
    for i in 1:param.N
        speed=norm(var.v[i,:])
        clamped=clamp(speed,param.minSpeed,param.maxSpeed)
        var.v[i,:].*=(clamped/speed)
    end
end

#check boundary
function check_boundary(param,var)
    for i in 1:param.N
        for n in 1:param.d
            if ((var.p[i,n]<param.minWorld && var.v[i,n]<0) || (var.p[i,n]>param.maxWorld && var.v[i,n]>0))
                var.v[i,n]*=-1.0
            end
        end
    end
end

function update(param,var)
    clear_vec(var)
    rule1(param,var)
    rule2(param,var)
    rule3(param,var)
    var.v.+=param.ω1*var.v1+param.ω2*var.v2+param.ω3*var.v3
    limit_vel(param,var)
    var.p.+=var.v
    check_boundary(param,var)
end
const dt=1.0

@time anim = @animate for m in 0:300
    t=round(m*dt,digits=2)
    quiver(var1.p[:,1],var1.p[:,2],
        quiver=(var1.v[:,1].*5.0,var1.v[:,2].*5.0),
        xlab="x",xlims=(minWorld-10,maxWorld+10),
        ylab="y",ylims=(minWorld-10,maxWorld+10),
        #zlab="z",zlims=(minZ,maxZ),
        leg=false,
        title="Boids,t = $t")
    plot!(var1.p[:,1],var1.p[:,2],seriestype=:scatter,markersize=2)
    update(param1,var1)

end
@time gif(anim,"tmp/Boids4.gif",fps=10)