# Ecosystem Models

using AlgebraicDynamics
using AlgebraicDynamics.DWDDynam
using AlgebraicDynamics.UWDDynam

using Catlab.CategoricalAlgebra
using Catlab.WiringDiagrams
using Catlab.Graphics
using Catlab.Programs

using OrdinaryDiffEq
using Plots, Plots.PlotMeasures

## Land Ecosystem

### Rabbits and foxes

A standard Lotka Volterra predator-prey model is the composition of three primitive resource sharers:

1. a model of rabbit growth: this resource sharer has dynamics $\dot r(t) = \alpha r(t)$ and one port which exposes the rabbit population.
2. a model of rabbit/fox predation: this resource sharer has dynamics $\dot r(t) = -\beta r(t) f(t), \dot f(t) = \gamma r(t)f(t)$ and two ports which expose the rabbit and fox populations respectively.
3. a model of fox population decline: this resource sharer has dynamics $\dot f(t) = -\delta f(t)$ and one port which exposes the fox population.

However, there are not two independent rabbit populations – one that grows and one that gets eaten by foxes. Likewise, there are not two independent fox populations – one that declines and one that feasts on rabbits. To capture these interactions between the trio of resource sharers, we compose them by identifying the exposed rabbit populations and identifying the exposed fox populations. The syntax for this undirected composition is defined by an undirected wiring diagram.

# Define the primitive systems
dotr(u,p,t) = p[1]*u
dotrf(u,p,t) = [-p[2]*u[1]*u[2], p[3]*u[1]*u[2]]
dotf(u,p,t) = -p[4]*u

rabbit_growth = ContinuousResourceSharer{Float64}(1, dotr)
rabbitfox_predation = ContinuousResourceSharer{Float64}(2, dotrf)
fox_decline = ContinuousResourceSharer{Float64}(1, dotf)

# Define the composition pattern
rabbitfox_pattern = @relation (rabbits, foxes) begin
rabbit_growth(rabbits)
rabbitfox_predation(rabbits,foxes)
fox_decline(foxes)
end

# Compose
rabbitfox_system = oapply(rabbitfox_pattern, [rabbit_growth, rabbitfox_predation, fox_decline])

Previously, when we derived the Lotka-Volterra model via undirected composition, we by-hand defined the undirected wiring diagram that implements the composition pattern. In contrast, here we implement the same composition pattern as before but this time using the @relation macro. This strategy simplifies the definition and explicitly names the boxes and variables. We visualize the composition pattern below.

to_graphviz(rabbitfox_pattern, box_labels = :name, junction_labels = :variable, edge_attrs=Dict(:len => ".75"))

We can now construct an ODEProblem from the resource sharer rabbitfox_system and plot the solution.

α, β, γ, δ = 0.3, 0.015, 0.015, 0.7
params = [α, β, γ, δ]

u0 = [10.0, 100.0]
tspan = (0.0, 100.0)

prob = ODEProblem(rabbitfox_system, u0, tspan, params)
sol = solve(prob, Tsit5())

plot(sol, lw=2, title = "Lotka-Volterra Predator-Prey Model", bottom_margin=10mm, left_margin=10mm, label=["rabbits" "foxes"])
xlabel!("Time")
ylabel!("Population size")

### Rabbits, foxes, and hawks

Suppose we now have a three species ecosystem containing rabbits, foxes, and hawks. Foxes and hawks both prey upon rabbits but do not interact with each other. This ecosystem consists of five primitive systems which share variables.

1. rabbit growth: $\dot r(t) = \alpha r(t)$
2. rabbit/fox predation: $\dot r(t) = -\beta r(t) f(t), \dot f(t) = \delta r(t)f(t)$
3. fox decline: $\dot f(t) = -\gamma f(t)$
4. rabbit/hawk predation: $\dot r(t) = -\beta' r(t)h(t), \dot h(t) = \delta' r(t)h(t)$
5. hawk decline: $\dot h(t) = -\gamma' h(t)$

This means the desired composition pattern has five boxes and many ports and wires to keep track of. Instead of implementing this composition pattern by hand, we construct it as a pushout.

# Define the composition pattern for rabbit growth
rabbit_pattern = @relation (rabbits,) -> rabbit_growth(rabbits)

# Define the composition pattern for the rabbit/hawk Lotka Volterra model
rabbithawk_pattern = @relation (rabbits, hawks) begin
rabbit_growth(rabbits)
rabbit_hawk_predation(rabbits,hawks)
hawk_decline(hawks)
end

# Define transformations between the composition patterns
rabbitfox_transform  = ACSetTransformation((Box=[1], Junction=[1], Port=[1], OuterPort=[1]), rabbit_pattern, rabbitfox_pattern)
rabbithawk_transform = ACSetTransformation((Box=[1], Junction=[1], Port=[1], OuterPort=[1]), rabbit_pattern, rabbithawk_pattern)

# Take the pushout to define the composition pattern for the rabbit, fox, hawk system
rabbitfoxhawk_pattern = ob(pushout(rabbitfox_transform, rabbithawk_transform))

# Visualize the compsition pattern
to_graphviz(rabbitfoxhawk_pattern, box_labels = :name, junction_labels = :variable, edge_attrs=Dict(:len => ".9"))
# Define the additional primitive systems
dotrh(x, p, t) = [-p[5]*x[1]*x[2], p[6]*x[1]*x[2]]
doth(x, p, t)  = -p[7]*x

rabbithawk_predation = ContinuousResourceSharer{Float64}(2, dotrh)
hawk_decline         = ContinuousResourceSharer{Float64}(1, doth)

# Compose
land_system = oapply(rabbitfoxhawk_pattern,
[rabbit_growth, rabbitfox_predation, fox_decline,
rabbithawk_predation, hawk_decline])

# Solve and plot
β′, γ′, δ′ = .01, .01, .5
params = vcat(params, [β′, γ′, δ′])

u0 = [10.0, 100.0, 50.0]
tspan = (0.0, 100.0)

prob = ODEProblem(land_system, u0, tspan, params)
sol = solve(prob, Tsit5())

plot(sol, lw=2, title = "Land Ecosystem", bottom_margin=10mm, left_margin=10mm, label=["rabbits" "foxes" "hawks"])
xlabel!("Time")
ylabel!("Population size")

Unfortunately, the hawks are going extinct in this model. We'll have to give hawks something else to eat!

## Ocean Ecosystem

Consider a ocean ecosystem containing three species —- little fish, big fish, and sharks -— with two predation interactions —- sharks eat big fish and big fish eat little fish.

This ecosystem can be modeled as the composition of 3 machines:

1. Evolution of the little fish population: this machine has one exogenous variable which represents a population of predators $h(t)$ that hunt little fish. This machine has one output which emits the little fish population. The dynamics of this machine is the driven ODE $\dot f(t) = \alpha f(t) - \beta f(t)h(t)$
2. Evolution of the big fish population: this machine has two exogenous variables which represent a population of prey $e(t)$ that are eaten by big fish and a population of predators $h(t)$ which hunt big fish. This machine has one output which emits the big fish population. The dynamics of this machine is the drive ODE $\dot F(t) = \gamma F(t)e(t) - \delta F(t) - \beta'F(t)h(t)$
3. Evolution of the shark population: this machine has one exogenous variable which represents a population of prey $e(t)$ that are eaten by sharks. This machine has one output which emits the shark population. The dynamics of this machine is the driven ODE $\dot s(t) = \gamma's(t)e(t) - \delta's(t)$
# Define the primitive systems
dotfish(f, x, p, t) = [p[1]*f[1] - p[2]*x[1]*f[1]]
dotFISH(F, x, p, t) = [p[3]*x[1]*F[1] - p[4]*F[1] - p[5]*x[2]*F[1]]
dotsharks(s, x, p, t) = [p[6]*s[1]*x[1]-p[7]*s[1]]

fish   = ContinuousMachine{Float64}(1,1,1, dotfish,   f->f)
FISH   = ContinuousMachine{Float64}(2,1,1, dotFISH,   F->F)
sharks = ContinuousMachine{Float64}(1,1,1, dotsharks, s->s)

We compose these machines by (1) sending the output of the big fish machine as the input to both the little fish and shark machines and (2) sending the output of the little fish and shark machines as the inputs to the big fish machine. The syntax for this directed composition is given by a directed wiring diagram.

# Define the composition pattern
ocean_pattern = WiringDiagram([], [])
fish_box = add_box!(ocean_pattern, Box(:fish, [:pop], [:pop]))
Fish_box = add_box!(ocean_pattern, Box(:Fish, [:pop, :pop], [:pop]))
shark_box = add_box!(ocean_pattern, Box(:shark, [:pop], [:pop]))

(fish_box, 1)  => (Fish_box, 1),
(shark_box, 1) => (Fish_box, 2),
(Fish_box, 1)  => (fish_box, 1),
(Fish_box, 1)  => (shark_box, 1)
])

# Visualize the composition pattern
to_graphviz(ocean_pattern, orientation=TopToBottom)
# Compose
ocean_system = oapply(ocean_pattern, [fish, FISH, sharks])

# Solve and plot
α, β, γ, δ, β′, γ′, δ′ = 0.35, 0.015, 0.015, 0.7, 0.017, 0.017, 0.35
params = [α, β, γ, δ, β′, γ′, δ′]

u0 = [100.0, 10, 2.0]
tspan = (0.0, 100.0)

prob = ODEProblem(ocean_system, u0, tspan, params)
sol = solve(prob, FRK65(0))

plot(sol, lw=2, title = "Ocean Ecosystem", bottom_margin=10mm, left_margin=10mm, label=["little fish" "big fish" "sharks"])
xlabel!("Time")
ylabel!("Population size")