using RecordedArrays
using Random
using EvolutionaryModelingTools
using EvolutionaryModelingTools.Scalar

# parameters
const T = 10.0
const β = 0.005
const ν = 0.01
const α = 0.2
const r = 0.5
const d = 0.1
const c = 0.001
const S = 100
const I = 1
const R = 0

## simulation with this package
# epidemic dynamics
@reaction_eq infection β S + I --> I + I
@reaction_eq recovery ν I --> R
@reaction_eq virulence α I --> 0 # death of infection host caused by virus
# demography dynamics
# generate reactions with @eval
for sym in (:S, :I, :R)
    r_name = Symbol(:growth_, sym)
    d_name = Symbol(:death_, sym)
    @eval @reaction_eq $r_name r $sym --> S + $sym # growth
    @eval @reaction_eq $d_name d $sym --> 0 # death
    for sym′ in (:S, :I, :R)
        c_name = Symbol(:competition_, sym, sym′)
        @eval @reaction_eq $c_name c $sym + $sym′ --> $sym′ # competition
    end
end

const REACTIONS = (
    infection, recovery, virulence,
    growth_S, growth_I, growth_R,
    death_S, death_I, death_R,
    competition_SS, competition_SI, competition_SR,
    competition_IS, competition_II, competition_IR,
    competition_RS, competition_RI, competition_RR
)

run_gillespie(rng=Random.GLOBAL_RNG, T=T, β=β, ν=ν, α=α, r=r, d=d, c=c, S=S, I=I, R=R) =
    gillespie(rng, T, (; β, ν, α, r, d, c, S=scalar(S), I=scalar(I), R=scalar(R)), REACTIONS)

function run_gillespie_record(rng=Random.GLOBAL_RNG, T=T, β=β, ν=ν, α=α,
        r=r, d=d, c=c, S=S, I=I, R=R)
    clock = ContinuousClock(T)
    ps = (;β, ν, α, r, d, c,
        S=recorded(DynamicEntry, clock, S),
        I=recorded(DynamicEntry, clock, I),
        R=recorded(DynamicEntry, clock, R),
    )
    return gillespie(rng, clock, ps, REACTIONS)
end

## simulation manually
### most of the code is generated by Github Copilot
### a useful tool for generating code like this
function run_manually(rng=Random.GLOBAL_RNG, T=T, β=β, ν=ν, α=α,
        r=r, d=d, c=c, S=S, I=I, R=R)
    t = 0
    while t <= T
        # calculate rates
        infection_rate = β * S * I
        recovery_rate = ν * I
        virulence_rate = α * I
        growth_S_rate = r * S
        growth_I_rate = r * I
        growth_R_rate = r * R
        death_S_rate = d * S
        death_I_rate = d * I
        death_R_rate = d * R
        competition_SS_rate = c * S * S
        competition_SI_rate = c * S * I
        competition_SR_rate = c * S * R
        competition_IS_rate = c * I * S
        competition_II_rate = c * I * I
        competition_IR_rate = c * I * R
        competition_RS_rate = c * R * S
        competition_RI_rate = c * R * I
        competition_RR_rate = c * R * R
        # summed rate
        summed = infection_rate + recovery_rate + virulence_rate +
            growth_S_rate + growth_I_rate + growth_R_rate +
            death_S_rate + death_I_rate + death_R_rate +
            competition_SS_rate + competition_SI_rate + competition_SR_rate +
            competition_IS_rate + competition_II_rate + competition_IR_rate +
            competition_RS_rate + competition_RI_rate + competition_RR_rate
        # break if summed is zero
        summed == 0 && break
        # update current time
        t += -log(rand(rng)) / summed
        # sample a reaction and adjust population size
        rn = rand(rng) * summed
        if rn < infection_rate
            S -= 1
            I += 1
        elseif (rn -= infection_rate) < recovery_rate
            I -= 1
            R += 1
        elseif (rn -= recovery_rate) < virulence_rate
            I -= 1
        elseif (rn -= virulence_rate) < growth_S_rate
            S += 1
        elseif (rn -= growth_S_rate) < growth_I_rate
            S += 1
        elseif (rn -= growth_I_rate) < growth_R_rate
            S += 1
        elseif (rn -= growth_R_rate) < death_S_rate
            S -= 1
        elseif (rn -= death_S_rate) < death_I_rate
            I -= 1
        elseif (rn -= death_I_rate) < death_R_rate
            R -= 1
        elseif (rn -= death_R_rate) < competition_SS_rate
            S -= 1
        elseif (rn -= competition_SS_rate) < competition_SI_rate
            S -= 1
        elseif (rn -= competition_SI_rate) < competition_SR_rate
            S -= 1
        elseif (rn -= competition_SR_rate) < competition_IS_rate
            I -= 1
        elseif (rn -= competition_IS_rate) < competition_II_rate
            I -= 1
        elseif (rn -= competition_II_rate) < competition_IR_rate
            I -= 1
        elseif (rn -= competition_IR_rate) < competition_RS_rate
            R -= 1
        elseif (rn -= competition_RS_rate) < competition_RI_rate
            R -= 1
        else
            R -= 1
        end
    end
    return (; t, S, I, R)
end

function run_manually_record(rng=Random.GLOBAL_RNG, T=T, β=β, ν=ν, α=α,
        r=r, d=d, c=c, S=S, I=I, R=R)
    clock = ContinuousClock(T)
    S′ = recorded(DynamicEntry, clock, S)
    I′ = recorded(DynamicEntry, clock, I)
    R′ = recorded(DynamicEntry, clock, R)
    for _ in clock
        # calculate rates
        infection_rate = β * S′ * I′
        recovery_rate = ν * I′
        virulence_rate = α * I′
        growth_S_rate = r * S′
        growth_I_rate = r * I′
        growth_R_rate = r * R′
        death_S_rate = d * S′
        death_I_rate = d * I′
        death_R_rate = d * R′
        competition_SS_rate = c * S′ * S′
        competition_SI_rate = c * S′ * I′
        competition_SR_rate = c * S′ * R′
        competition_IS_rate = c * I′ * S′
        competition_II_rate = c * I′ * I′
        competition_IR_rate = c * I′ * R′
        competition_RS_rate = c * R′ * S′
        competition_RI_rate = c * R′ * I′
        competition_RR_rate = c * R′ * R′
        # summed rate
        summed = infection_rate + recovery_rate + virulence_rate +
            growth_S_rate + growth_I_rate + growth_R_rate +
            death_S_rate + death_I_rate + death_R_rate +
            competition_SS_rate + competition_SI_rate + competition_SR_rate +
            competition_IS_rate + competition_II_rate + competition_IR_rate +
            competition_RS_rate + competition_RI_rate + competition_RR_rate
        # break if summed is zero
        summed == 0 && break
        # update current time
        increase!(clock, -log(rand(rng)) / summed)
        # sample a reaction and adjust population size
        rn = rand(rng) * summed
        if rn < infection_rate
            S′[] -= 1
            I′[] += 1
        elseif (rn -= infection_rate) < recovery_rate
            I′[] -= 1
            R′[] += 1
        elseif (rn -= recovery_rate) < virulence_rate
            I′[] -= 1
        elseif (rn -= virulence_rate) < growth_S_rate
            S′[] += 1
        elseif (rn -= growth_S_rate) < growth_I_rate
            S′[] += 1
        elseif (rn -= growth_I_rate) < growth_R_rate
            S′[] += 1
        elseif (rn -= growth_R_rate) < death_S_rate
            S′[] -= 1
        elseif (rn -= death_S_rate) < death_I_rate
            I′[] -= 1
        elseif (rn -= death_I_rate) < death_R_rate
            R′[] -= 1
        elseif (rn -= death_R_rate) < competition_SS_rate
            S′[] -= 1
        elseif (rn -= competition_SS_rate) < competition_SI_rate
            S′[] -= 1
        elseif (rn -= competition_SI_rate) < competition_SR_rate
            S′[] -= 1
        elseif (rn -= competition_SR_rate) < competition_IS_rate
            I′[] -= 1
        elseif (rn -= competition_IS_rate) < competition_II_rate
            I′[] -= 1
        elseif (rn -= competition_II_rate) < competition_IR_rate
            I′[] -= 1
        elseif (rn -= competition_IR_rate) < competition_RS_rate
            R′[] -= 1
        elseif (rn -= competition_RS_rate) < competition_RI_rate
            R′[] -= 1
        else
            R′[] -= 1
        end
    end
    return (; clock, S=S′, I=I′, R=R′)
end
