#lang racket/base

(require math/array)

;; scalar constants
(define years 30)
(define prop-female 0.5)
(define egg-surv 0.6)

;; age-specific fecundity and survival
(define fecundity #(0 0 200 400 800))
(define survival #(0.2 0.4 0.6 0.8 0))
(define capacity #(1e6 1e5 1e4 1e3 1e2))

;; multistage Beverton-Holt model
(define (beverton-holt N p c)
  (/ N (+ (/ 1 p) (/ N c))))

;; initialize empty results matrix and make mutable
(define results (array->mutable-array (make-array (vector years (vector-length fecundity)) 0)))

;; initialize abundances in first year to arbitrary non-zero value
(array-slice-set! results (list '(0) (::)) (make-array (vector 1 (vector-length fecundity)) 10))

;; iterate over results to fill matrix
(for* ([i (in-range (sub1 years))]
       [j (in-range (vector-length fecundity))])
  ;; current abundance vector
  (define N (array-ref results (vector i j)))
  ;; reproduction
  (define fecundity-age-j (vector-ref fecundity j))
  (when (> fecundity-age-j 0)  ;; not all age classes reproduce
    (define N-female (* N prop-female))
    ;; next year age-0
    (define Nt-age-0 (array-ref results (vector (add1 i) 0)))

    (define new-age-0 (beverton-holt
                       N-female
                       (* fecundity-age-j egg-surv)
                       (- (vector-ref capacity 0) Nt-age-0)))
    (array-set! results (vector (add1 i) 0) (+ Nt-age-0 new-age-0))
    )
  ;; survival
  (define survival-age-j (vector-ref survival j))
  (when (> survival-age-j 0)
    (define Nt-age-j (array-ref results (vector (add1 i) (add1 j))))
    (define new-age-j (beverton-holt
                       N
                       survival-age-j
                       (- (vector-ref capacity (add1 j)) Nt-age-j)))
    (array-set! results (vector (add1 i) (add1 j)) (+ Nt-age-j new-age-j))
    )
  )
results