########## 
##########
# Functions to perform perturbations on the food webs 
##########
##########
# AUTHOR: Cole B. Brookson
# DATE OF CREATION: 2021-08-09
##########
##########

# set up =======================================================================
using DrWatson
@quickactivate "trait-based-rewiring"
# logger set up 
using Logging: global_logger
using TerminalLoggers: TerminalLogger
global_logger(TerminalLogger())
using BenchmarkTools

include(srcdir("01_eqn_component_functions.jl"))
include(srcdir("00_species_list_creation.jl"))
include(srcdir("03_initial_sims_functions.jl"))
include(srcdir("04_perturbation_sims_functions.jl"))

# set all positional values 
organism_id_pos = 1
organism_type_pos = 2
body_size_pos = 3
lower_prey_limit_pos = 4
upper_prey_limit_pos = 5
habitat_midpoint_pos = 6
habitat_lower_limit_pos = 7
habitat_upper_limit_pos = 8
nutritional_value_pos = 9
biomass_pos = 10
network_id_pos = 11
alive_pos = 12 

# pulse perturbations ==========================================================

# little line to set the network id (this happens from the shell line)
network_id = parse(Int64, ARGS[1]) 

species_list = readdlm(datadir("./intermediate-files/species-lists-pre-perturbation/species_list_0001.csv"), ',', Any)
network = create_interaction_matrix(species_list)

# set pulse amounts 
c = 210
x = 100
f = 1
z = c/(x*(1/f))
# set vector callback
extinction_callback = 
    VectorContinuousCallback(extinction_threshold,
                            extinction_affect!,
                            size(species_list, 1),
                            save_positions = ntuple(x->true,
                                                    size(species_list, 1)))
steady_callback = TerminateSteadyState(1e-8, 
                                        1e-6,
                                        DiffEqCallbacks.allDerivPass,
                                        min_t = 1205)

# set callback set
"""
    affect_inc!(integrator)

Affect for the condition to work on. Increase the mortality by making the amount
(the fourth element in p) equal to z.  
"""
function affect_inc!(integrator, z)
  integrator.p[4] = z
end

"""
    affect_dec!(integrator)

Affect for the condition to work on. Decrease the fourth element in p, 
returning it to 1.0.  
"""
function affect_dec!(integrator)
  integrator.p[4] = nothing 
end

save_positions1 = (true, true)
cb_inc = DiscreteCallback(inc_mort_cond, affect_inc!,
                            save_positions=save_positions1)
save_positions2 = (false, true)
cb_dec = DiscreteCallback(dec_mort_cond, affect_dec!, 
                            save_positions=save_positions2)
callback_set = CallbackSet(extinction_callback, 
                            #steady_callback, #keep this out until fixed 
                            cb_inc,
                            cb_dec)

#  allocate x0
u0 = Float64[]
for i in 1:size(species_list, 1)
    append!(u0, species_list[i, biomass_pos])
end

# set perturbation tspan
tspan = (1000.0, 1100.0)

# set constant stop values
const stop1 = [1010.0]
const stop2 = [1080.0]
const tstop = [1010.;1080.]

# loop through random selections of species to perturb
#@info{"starting loop..."}

# set times to save from the simulation 
save_times = [collect(range(1000.0, 1300.0, step = 0.1));
                collect(range(1250, 4800, step = sample(200:400, 1)[1]));
                collect(range(4800.0, 5000.0, step = 0.1))]

for j in 1:100

    # pick set of random species to perturb
    all = Array(1:size(species_list, 1))
    perturb_sp = rand(Array(1:size(species_list, 1)),
                        Int(round(size(species_list, 1)*(1/2))))
    non_perturb = setdiff(Set(all), Set(perturb_sp))

    # set p 
    amount = nothing 
    p = [species_list, network, 0.3, amount, perturb_sp, non_perturb]

    prob_pulse = ODEProblem(biomass_perturb_random,
                            u0, 
                            tspan, 
                            p, 
                            tstops=tstop) 
    @benchmark sol_prob_pulse = solve(prob_pulse,
                                        OwrenZen5(),
                                        callback = callback_set,
                                        reltol = 1e-4,
                                        # save every part for the perturbation and 
                                        # the first 150 time steps after the 
                                        # perturbation, but then just grab a few 
                                        # from there to the end 
                                        saveat = save_times, 
                                        progress = true,
                                        progress_steps = 1)
    # save output of the solver
    sol_name = string("sol-pulse-net-", lpad(network_id,4,"0"), "-sim-", lpad(j, 4, "0"), ".csv")
    writedlm(datadir("./simulation-outpus/pulse-sims/output/$name"), sol_prob_pulse, ",")
    #plot(sol_prob_pulse)
    u_matrix = zeros(length(sol_prob_pulse.u), length(sol_prob_pulse.u[1]))
    for i in 1:size(u_matrix, 1)
        u_matrix[i,:] = sol_prob_pulse.u[i]
    end 
    output = hcat(sol_prob_pulse.t, u_matrix)
    # write out the biomass outputs
    output_name = string("pulse-net-", lpad(network_id,4,"0"), "-sim-", lpad(j, 4, "0"), ".csv") 
    writedlm(datadir("./simulation-outpus/pulse-sims/output/$name"), output_name, ",")
end

