########## 
##########
# Definition of functions used to look at species rewiring in relevant paper
##########
##########
# AUTHOR: Cole B. Brookson
# DATE OF CREATION: 2021-04-29
##########
##########

# set up =======================================================================

# set packages
using DataFrames, Distributions, Random, StatsBase, Documenter, Colors
using LightGraphs, GraphPlot, GraphRecipes, Plots

# trait-related functions ======================================================

"""
    assign_trait_values(organism_type::String, organism_id::Int)

Takes in the type of organism and it's unique id and uses the percentage of its
own body size as defined, to return a list of the various traits for that 
organism to be passed to the next list.

"""
function assign_trait_values(
    organism_type::String, 
    organism_id::Int,
    lower_prey_limit::Float64 = 0.001,
    upper_prey_limit::Float64 = 0.65
    )
    # determine producer or consumer
    if organism_type == "producer"
        # scale for producers should be 0.34 (Brose 2019 paper on body size)
        body_size = rand(Exponential(0.34), 1)
        # set prey limits - empty for producers
        prey_limits = Array([NaN, NaN])
    else # everything here down is for consumers
        # scale for consumers should be 800 (Brose 2019 paper on body size)
        body_size = rand(Exponential(20), 1)
        # set prey limits - empty for producers
        lower_limit = body_size*lower_prey_limit
        upper_limit = body_size*upper_prey_limit
        prey_limits = [lower_limit[1], upper_limit[1]]

    end
    # habitat values get set by same process for both consumers and producers
    
    # randomly select habitat midpoint 
    habitat_midpoint = round(rand(Uniform(0,1), 1)[1], digits = 3)
    # define vector of the midpoint and upper/lower habitat bounds
    habitat = [habitat_midpoint, 0, 0]
    if habitat_midpoint >= 0.4 # for lower bound
        habitat[2] = round((habitat_midpoint - 0.4), digits = 3)
    else
        habitat[2] = 0
    end
    if habitat_midpoint <= 0.6 # for upper bound 
        habitat[3] = round((habitat_midpoint + 0.4), digits = 2)
    else
        habitat[3] = 1
    end
    #nutritional value 
    nutritional_value = rand(Beta(5, 2))

    #habitat = [habitat_midpoint, 0, 1]
    trait_list = [organism_id, organism_type, round(body_size[1], digits = 2),
                    round(prey_limits[1], digits = 2), 
                    round(prey_limits[2], digits = 2), habitat[1], 
                    round(habitat[2], digits = 2), 
                    round(habitat[3], digits = 2), 
                    round(nutritional_value, digits = 2)]

    return trait_list
end

"""
    define_species_list(n::Int, shape::String)

Takes in the number of species and a shape for the web and returns an array 
with each element being the trait list for each species in the web


# examples & checks
```jldoctest
# make sure ArgumentError works
julia> define_species_list(50, "topheavy")
ERROR: ArgumentError: shape is not recognized: can only be 'top-heavy', 
'bottom-heavy', or 'uniform'

```
"""
function define_species_list(n::Int, richness_shape::String)
    # make sure shape is of the right kind 
    if richness_shape ∉ ["top-heavy", "bottom-heavy", "uniform"]
        ArgumentError("richness_shape is not recognized: can only be 'top-heavy', 
        'bottom-heavy', or 'uniform'")
    end
    # define number of species in the array with predefined types of the sublist
    species_list = Array{Any}(undef, 50, 9)
    # set items in organism type list to be selected from 
    types = ["producer", "consumer"]
    # if statements to determine the weights of these
    if richness_shape == "top-heavy"
        weights = [1, 5]
    elseif richness_shape == "bottom-heavy"
        weights = [5,1]
    else richness_shape == "uniform"
        weights = [1,1]
    end
    # use while loop to populate list
    i = 1
    while i <= n
        organism_type = sample(types, Weights(weights))
        species_list[i, :] = assign_trait_values(organism_type, i)
        

        i += 1 #increment
    end
    return species_list
end 

"""
    get_species(species_list, i)

Finds the species list entry according to species id, NOT simply location in the 
array
"""
function get_species(species_list::Matrix, id::Int)
    
    # find focal species in array
    matrix_sp = species_list[species_list[:, organism_id_pos] .== id, :]
    focal_sp = []
    focal_sp = append!(focal_sp, matrix_sp)

    return focal_sp
end 

"""
    determine_interaction(predator_id::Int, prey_id::Int, species_list::Array)

Assigns whether or not a given predator & prey have an interaction according to 
the stated rules of the algorithm

"""
function determine_interaction(
    predator_id::Int, 
    prey_id::Int, 
    species_list::Array
    )
    # extract the values for habitat, size, and predation limits for ease 
    predator = get_species(species_list, predator_id)
    prey = get_species(species_list, prey_id)

    pred_hab = predator[habitat_midpoint_pos:habitat_upper_limit_pos]
    prey_hab = prey[habitat_midpoint_pos:habitat_upper_limit_pos]
    prey_size = prey[body_size_pos]
    pred_min = predator[lower_prey_limit_pos]
    pred_max = predator[upper_prey_limit_pos]

    if (predator_id == prey_id) || (predator[organism_type_pos] == "Producer")
        value = 0 # assuming no canibalism and that producers don't predate
    elseif (((prey_hab[2] >= pred_hab[2] && prey_hab[2] <= pred_hab[3]) ||
        (prey_hab[3] >= pred_hab[2] && prey_hab[3] <= pred_hab[3])) && 
        (prey_size[1] <= pred_max && prey_size[1] >= pred_min))
        value = 1 # assign 1 if within body size and habitat range
    else
        value = 0
    end 

    return value 
end

"""
    biomass_interaction(predator_id::Int, prey_id::Int, species_list::Array)

Assigns whether or not a given predator & prey have an interaction according the
same stated rules as determine_interaction(), but now including whether either 
the predator or prey are extinct (i.e. biomass = 0.0)

"""
function biomass_interaction(
    predator_id::Int, 
    prey_id::Int, 
    species_list::Array
    )
    # extract the values for habitat, size, and predation limits for ease 
    predator = get_species(species_list, predator_id)
    prey = get_species(species_list, prey_id)

    pred_hab = predator[habitat_midpoint_pos:habitat_upper_limit_pos]
    prey_hab = prey[habitat_midpoint_pos:habitat_upper_limit_pos]
    prey_size = prey[body_size_pos]
    pred_min = predator[lower_prey_limit_pos]
    pred_max = predator[upper_prey_limit_pos]

    if (predator_id == prey_id) || (predator[organism_type_pos] == "Producer")
        value = 0 # assuming no canibalism and that producers don't predate
    elseif predator[biomass_pos] == 0.0
        value = 0
    elseif prey[biomass_pos] == 0.0
        value = 0
    elseif (((prey_hab[2] >= pred_hab[2] && prey_hab[2] <= pred_hab[3]) ||
        (prey_hab[3] >= pred_hab[2] && prey_hab[3] <= pred_hab[3])) && 
        (prey_size[1] <= pred_max && prey_size[1] >= pred_min))
        value = 1 # assign 1 if within body size and habitat range
    else
        value = 0
    end 

    return value 
end

"""
    create_interaction_matrix(species_list::Array)

Takes in the species list and returns an interaction matrix that will form the 
basis of the networks to then be created.

"""
function create_interaction_matrix(species_list::Array)
    interaction_matrix = fill(0, (size(species_list, 1), size(species_list, 1)))
    for i in 1:size(interaction_matrix)[1] #rows
        for j in 1:size(interaction_matrix)[2] #cols
            interaction_matrix[i,j] = determine_interaction(j, i, species_list)
        end
    end
    return interaction_matrix
end

"""
    biomass_interaction_matrix(species_list::Array)

Takes in the species list and returns an interaction matrix that will allow for
inspection of networks after biomass simulations have occured.

"""
function biomass_interaction_matrix(species_list::Array)
    interaction_matrix = fill(0, (size(species_list, 1), size(species_list, 1)))
    for i in 1:size(interaction_matrix)[1] #rows
        for j in 1:size(interaction_matrix)[2] #cols
            interaction_matrix[i,j] = biomass_interaction(j, i, species_list)
        end
    end
    return interaction_matrix
end

"""
    check_connectedness(species_list)

Check the connectedness of the interaction_matrix formed by a specific 
species list draw.
"""
function check_connectedness(
    species_list::Array,
    )

    interaction_matrix = create_interaction_matrix(species_list)
    # get indices of the consumers to check if they are consuming something
    consumers = species_list[species_list[:,organism_type_pos] .== "consumer", 
                                organism_id_pos]
    producers = species_list[species_list[:,organism_type_pos] .== "producer", 
                                organism_id_pos]
    producers = species_list[species_list[:,organism_type_pos] .== "g", 
                                organism_id_pos]
    # make sure all consumers are eating at least one other organism (i.e. the 
    # sum of their column is > 0) 
    consumer_matrix = interaction_matrix[:,consumers]
    if size(consumer_matrix, 1) > 0
        consumer_sums = minimum(sum(consumer_matrix, dims = 1))
    else
        return "fail" # this makes sure there are no all-producer networks
    end
    # make sure all producers are being consumeed by least one other organism 
    # (i.e. the sum of their column is > 0) 
    producer_matrix = interaction_matrix[producers,:]
    if size(producer_matrix, 1) > 0
        producer_sums =  minimum(sum(producer_matrix, dims = 2))
    else
        producer_sums = 1 # just make it non-zero so it will pass the test
    end 

    if consumer_sums != 0 && producer_sums != 0
        return "pass"
    else 
        return "fail"
    end 
end 

"""
    network_draw(n::Int64, richness_shape::String)

This function takes in the number of species and the richness shape and draws
an initial network and makes the checks to see if it's valid - if it is 
connected, and then resamples the connectance to prepare it to have the network 
metrics sampled.
"""
function network_draw(
    n::Int64, 
    richness::String
    )

    # draw species list 
    species_list = define_species_list(n, richness)
    interaction_matrix = create_interaction_matrix(species_list)
    # check connectedness, and go through again if need be
    if check_connectedness(species_list) == "pass"
        #print("first time success ")
        return species_list, interaction_matrix 
    else
        #print("first time fail ")
        # go through while loop to re-draw until connected
        connect_acomplish = false
        connected_counter = 0
        while connect_acomplish == false
            connected_counter += 1
            # draw species list 
            species_list = define_species_list(n, richness)
            interaction_matrix = create_interaction_matrix(species_list)
            if check_connectedness(species_list) == "fail"
                continue 
            elseif check_connectedness(species_list) == "pass"
                connect_acomplish = true 
                #print("success after ", connected_counter, " additional draws.")
                return species_list, interaction_matrix
            end 
        end
    end 
end 

"""
    get_thirds(sequence::Array, num::Int64)
    
Takes in sequence and an integer and returns the sequence split into `num` 
groups as an array of arrays.
"""
function get_thirds(sequence::Array, num::Int64)

    k, m = divrem(length(sequence), num)
    thirds = []
    for i in 0:(num-1)
        append!(thirds, [sequence[(i*k + min(i, m) +1):((i+1)*k + min(i+1, m))]])
    end
    thirds 
    return thirds
end

"""
    initialize_biomass(species_list::Array, biomass_shape::String="pyramid")

Initializes the amount of biomass for each species in the species list according
to the shape of the web (i.e. pyramid, uniform, inverted)

"""
function initialize_biomass(
    species_list::Array, 
    biomass_shape::String
    )

    
    # get all body body sizes 
    body_sizes = zeros(0)
    for i in 1:size(species_list,1)
        append!(body_sizes, species_list[i,body_size_pos])
    end
    # make three equal groups of the body sizes
    sorted_body_sizes = sort(body_sizes)
    thirds = get_thirds(sorted_body_sizes, 3)
    small = [thirds[1][1],thirds[1][end]] 
    medium = [thirds[2][1],thirds[2][end]]
    large = [thirds[3][1],thirds[3][end]]

    if biomass_shape == "pyramid"
        small_biomass = 12.5
        medium_biomass = 5
        large_biomass = 1.25
    elseif biomass_shape == "inverted"
        small_biomass = 1.25
        medium_biomass = 5
        large_biomass = 12.5
    else biomass_shape == "uniform"
        small_biomass = 6.2333333
        medium_biomass = 6.2333333
        large_biomass = 6.2333333
    end

    # get mean biomass for each size bin 
    num_small_sp = length(thirds[1])
    num_medium_sp = length(thirds[2])
    num_large_sp = length(thirds[3])
    
    large_mean = large_biomass/num_large_sp
    medium_mean = medium_biomass/num_medium_sp
    small_mean = small_biomass/num_small_sp
    
    # get random biomass sample and append it to the end of each species trait info 
    temp = zeros(50)
    for i in 1:size(species_list,1)
        if species_list[i,3] ≤ small[2]
            temp[i] = rand(TruncatedNormal(small_mean, 0.09, 0.05, 1))
        elseif species_list[i,3] >= large[1]
            temp[i] = rand(TruncatedNormal(large_mean, 0.09, 0.05, 1))
        else 
            temp[i] = rand(TruncatedNormal(medium_mean, 0.09, 0.05, 1))
        end
    end
    species_list = hcat(species_list, temp)

    return species_list
end 


  