ABM Pandemic

turtles-own [
  state  ;; State of each turtle (S for , I for infectious, R for recovered)
]

globals [
  delay
  infection_rate
  K
  observer
  plot-data
  max-ticks
]

to setup
  clear-all
  reset-ticks
  set max-ticks 80
  set infection_rate initial_infection_rate

  create-turtles population-size [
    set shape "person"
    setxy random-xcor random-ycor
    set state "S"
    set color green
  ]

  ask n-of initial_infections turtles [
    set state "I"
    set color red
  ]

  ask n-of initial_recoveries turtles [
    set state "R"
   set color blue
  ]

  set K population-size / 4 ;; Carrying capacity in population (logistic growth)
end 

to run-simulation
  let tick-count 0

  while [tick-count < max-ticks] [
    ask turtles [
      runge-kutta-step count turtles with [state = "S"] count turtles with [state = "I"] count turtles with [state = "R"]
    ]
    ask turtles [
      move
      if state = "I" [
        ask other turtles in-radius 1 [
          if state = "S" and random-float 1 < infection_rate [
            set state "I"
            set color red
          ]
        ]
        ifelse random-float 1 < recovery_rate [
          set state "R"
          set color blue
        ] [
          ;; Infected turtles that do not recover
        ]
      ]
    ]
    set delay (10 / tick_speed) * 0.1

    wait delay
    tick
    set tick-count tick-count + 1

    ;; Record counts and update plot
    let current-s count turtles with [state = "S"]
    let current-i count turtles with [state = "I"]
    let current-r count turtles with [state = "R"]
    set plot-data (list current-s current-i current-r)

    ;; Check if there are still infectious turtles
    if not any? turtles with [state = "I"] [
      stop
    ]
  ]
end 

to move
  rt random 360
  fd 1
  display
end 

to runge-kutta-step [S I R]
  let dt 0.1  ;; Time step size
  let k1 derivative S I R
  let k2 derivative (S + 0.5 * dt * item 0 k1) (I + 0.5 * dt * item 1 k1) (R + 0.5 * dt * item 2 k1)
  let k3 derivative (S + 0.5 * dt * item 0 k2) (I + 0.5 * dt * item 1 k2) (R + 0.5 * dt * item 2 k2)

  ;;print (list "Before update - S: " S " I: " I " R: " R)

  let k4 derivative (S + dt * item 0 k3) (I + dt * item 1 k3) (R + dt * item 2 k3)

  set S S + (dt / 6) * (item 0 k1 + 2 * item 0 k2 + 2 * item 0 k3 + item 0 k4)
  set I I + (dt / 6) * (item 1 k1 + 2 * item 1 k2 + 2 * item 1 k3 + item 1 k4)
  set R R + (dt / 6) * (item 2 k1 + 2 * item 2 k2 + 2 * item 2 k3 + item 2 k4)

  ;; Calculate logistic growth
  let t ticks
  let logistic-growth-value logistic-growth t

  ;; Update infection rate based on logistic growth
  set infection_rate infection_rate + logistic-growth-value

  ;;print (list "After update - S: " S " I: " I " R: " R)
end 

to-report derivative [S I R]
  let N count turtles
  let dS_dt ((-1) * (infection_rate / N) * I * S)
  let dI_dt ((infection_rate / N) * I * S) - (recovery_rate * I)
  let dR_dt (recovery_rate * I)

  report (list dS_dt dI_dt dR_dt)
end 

to-report logistic-growth [t]
  report K / (1 + initial_infection_rate * exp(- recovery_rate * t))
end