### ============== ### ============== ###
##    Behavioural Rules Model          ##
##    Martin Zumaya Hernandez          ##
##    EXAMPLE SIMULATION SCRIPT        ##
### ============== ### ============== ###

### ============ INCLUDE PACKAGES ============ ###

@everywhere using CollectiveDynamics.BehaviouralRules

### ============ SYSTEM'S PARAMETERS ============ ###

@everywhere dt = 1.0 # time step
@everywhere ρ = 0.3 # initial density

@everywhere v0 = 1.0 # speed

@everywhere η = 0.15 # noise intensity

@everywhere θ = 40.0 # maximum turn
@everywhere δ = 0.05 # deviation from aligned velocity

### ============ METRIC BEHAVIORAL THRESHOLDS ============ ###

n = parse(Int64, ARGS[1]) # number of particles

o = parse(Float64, ARGS[2]) # size of the orientation zone relative to system size
a = parse(Float64, ARGS[3]) # size of the attraction zone relative to system size

T = parse(Int64, ARGS[4]) # 10 ^ T iterations

rep = parse(Int64, ARGS[5]) # ensemble index

init = ARGS[6] # random or aligned initial velocities

@eval @everywhere init_e = $init

@eval @everywhere N = $n
@eval @everywhere Δo = $o
@eval @everywhere Δa = $a

@everywhere L  = cbrt(N / ρ) # size of box

@everywhere zor = 0.1 # zone of repulsion
# @everywhere zor = 1.0 # zone of repulsion
@everywhere zoo = zor + Δo*L # zone of orientation
@everywhere zoa = zoo + Δa*L # zone of attraction

println("rep:\t", zor, "\torient:\t", zoo, "\tattr:\t", zoa )

### ============ SHARED ARRAY INITIALIZATION ============ ###

pos = SharedArray{Float64}(3N) # particles positions
vel = SharedArray{Float64}(3N) # array of particles' velocities

v_int = SharedArray{Float64}(3N) # total signal
v_r   = SharedArray{Float64}(3N) # repulsion metric interactions
v_o   = SharedArray{Float64}(3N) # orientation metric interactions
v_a   = SharedArray{Float64}(3N) # attraction metric interactions

Rij = SharedArray{Float64}(N,N)

output_path = ""

if init_e == "R" # random initial conditions

    output_path = set_output_data_structure("BEHAV_R_01_N_015", N, ARGS[2], ARGS[3])
    println(output_path)

    ### ============ RANDOM INITIAL CONDITIONS ============ ###
    for i in 1:length(pos)
        pos[i] = 2*rand()*L - L
        vel[i] = 2*rand() - 1
    end

elseif init_e == "A" # ordered state initial condition

    output_path = set_output_data_structure("BEHAV_R_01_N_015_AL", N, ARGS[2], ARGS[3])
    println(output_path)

    ### ============ RANDOM POSITIONS BUT ALIGNED VELOCITIES ============ ###

    vel_0 = [2*rand() - 1, 2*rand() - 1, 2*rand() - 1]

    for i in 1:3:length(pos)
        pos[i]   = 2*rand()*L - L
        pos[i+1] = 2*rand()*L - L
        pos[i+2] = 2*rand()*L - L
        vel[i]   = vel_0[1] + δ*rand() - δ
        vel[i+1] = vel_0[2] + δ*rand() - δ
        vel[i+2] = vel_0[3] + δ*rand() - δ
    end

end

### ============ VELOCITY NORMALIZATION ============ ###

for i in 1:3:length(vel)
    norm = sqrt(vel[i]^2 + vel[i+1]^2 + vel[i+2]^2)
    vel[i] /= norm
    vel[i+1] /= norm
    vel[i+2] /= norm
end

### ============ SET OUTPUT ============ ###

pos_file = open(joinpath(output_path,"pos_$(rep).dat"), "w+")
vel_file = open(joinpath(output_path,"vel_$(rep).dat"), "w+")

# write initial conditions 
write(pos_file, pos)
write(vel_file, vel)
println("//////// ", 1)

### ============ TIME EVOLUTION ============ ###

times = [convert(Int, exp10(i)) for i in 0:T]

for i in 1:(length(times) - 1)

    for t in (times[i]+1):times[i+1]

        evolve_system(pos, vel, v_int, v_r, v_o, v_a, Rij, N, zor, zoo, zoa, η, θ, v0)

        if t % times[i] == 0 || t % div(times[i], exp10(1)) == 0
            println("//////// ", t)
            write(pos_file, pos)
            write(vel_file, vel)
        end
    end

end

close(pos_file)
close(vel_file)

rmprocs(workers())

println("Done all")

### ============== ### ============== ### ============== ###
