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

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

include(srcdir("01_eqn_component_functions.jl"))
include(srcdir("00_species_list_creation.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 

# periodic 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 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-7, 
                                        1e-5, 
                                        DiffEqCallbacks.allDerivPass,
                                        min_t = 1205)

# set callback set
save_positions = (true, true)
cb_inc = DiscreteCallback(inc_mort_cond, affect_inc!, save_positions=save_positions)
save_positions = (false, true)
cb_dec = DiscreteCallback(dec_mort_cond, affect_dec!, save_positions=save_positions)
callback_set = CallbackSet(extinction_callback, 
                            #steady_callback, #keep this out in the period ones 
                            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, 5000.0)

# set constant stop values
const stop1 = [1000.0, 2000.0, 3000.0, 4000.0]
const stop2 = [1100.0, 2100.0, 3100.0, 4100.0]
const tstop = [1000.;1100.;2000.;2100;3000.;3100.;4000.;4100.]

c = 210
x = 100
f = 
z = c/(x*(1/f))

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

for j in 1:10

    # 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/3))))
    non_perturb = setdiff(Set(all), Set(perturb_sp))

    # set times to save from the simulation 
    save_times = [collect(range(1000.0, 1250.0, step = 0.1));
                collect(range(1250, 1999, step = sample(200:400, 1)[1]));
                collect(range(2000.0, 2250.0, step = 0.1));
                collect(range(2250, 2999, step = sample(200:400, 1)[1]));
                collect(range(3000.0, 3250.0, step = 0.1));
                collect(range(3250, 3999, step = sample(200:400, 1)[1]));
                collect(range(4000.0, 4250.0, step = 0.1));
                collect(range(4250, 5000, step = sample(200:400, 1)[1]));]

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

    prob_period = ODEProblem(biomass_perturb_random,
                                        u0, 
                                        tspan, 
                                        p, 
                                        tstops=tstop) 
    @time sol_prob_period = solve(prob_period, 
                                        Tsit5(),
                                        callback = callback_set,
                                        # 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)
    # save output of the solver
    output = hcat(sol_prob_period.t, sol_prob_period.u)

    # write out the biomass outputs
    name = string("periodic-", network_id, "-", j, "-", "sims.csv") 
    path = "./data/simulation-outputs/period-sims/"
    writedlm(string(path,name), output, ",")
end