Atmospheric Dispersion model

rays-own [initial-heading]
breed [simulated-rays simulated-ray]
breed [rays ray]
globals [red_em blue_em green_em red_rec blue_rec green_rec test-list]


;;; HELPER PROCEDURES

to-report is-scattered [t_color]
  ;; gives scattering probability as a function of the light ray color
  let probability (ifelse-value
    t_color = blue [base_probability]
    t_color = green [base_probability / 1.8]
    t_color = red [base_probability / 4.6]
                 [0])
  report (random-float 100 <= probability)
end 

to-report scatter-angle
  ;; Rayleigh scattering has a scattering probability proportional to cos(theta)^2
  let which_sign random 2
  let sign (ifelse-value
    which_sign = 0 [ -1 ]
    which_sign = 1 [ 1 ]
  )
  report ((acos (sign * sqrt(random-float 1))) * 2) - 90
end 

to-report is-oob [y]
  ;; check if the y coordinate is out of bounds
  ;; returns 1 if reached the top of the simulation (escaped the atmosphere)
  ;; returns 2 if collided with the ground
  ;; returns 0 if no collisions
  if y > max-pycor - 0.5 [report 1]
  if y < min-pycor + 0.5 [report 2]
  report 0
end 

to-report check-sun-collision [exit-angle]
  ;; Checks if the exit-angle is within the sun's diameter
  if exit-angle > 180 [set exit-angle (exit-angle - 360)]
  report (
    (exit-angle > sun_angle - sun_diameter / 2 ) and
    (exit-angle < sun_angle + sun_diameter / 2 )
  )
end 

to-report get-index [angle]
  ;; For a given angle, gets the index of the cone bin
  if angle > 180 [set angle (angle - 360)]
  set angle (angle + 90)
  report floor(angle / 180 * division-factor)
end 

to-report get-color [red_emitted red_received green_emitted green_received blue_emitted blue_received]
  ;; yields an rgb vector of the color corresponding to the fraction of emitted and received rays
  ifelse relative-gain?
  [
    let multiplier max rgb (red_received / red_emitted) (green_received / green_emitted) (blue_received / blue_emitted)
    ifelse multiplier = 0 [report rgb 0 0 0]
    [
      let new-gain (relative-gain / 100) * (255 / multiplier)
      report rgb (new-gain * red_received / red_emitted) (new-gain * green_received / green_emitted) (new-gain * blue_received / blue_emitted)
    ]
  ]
  [report rgb (gain * red_received / red_emitted) (gain * green_received / green_emitted) (gain * blue_received / blue_emitted)]
end 


;;; SETUP PROCEDURES

to setup
  clear-all
  setup-rays
  setup-atmosphere
  setup-simulated-rays
  setup-counters
  reset-ticks
end 

to setup-rays
  ;; setup the rays with random colour and direction
  set-default-shape rays "line"
  create-rays n_rays [setxy 0 min-pycor + 1]
  ask rays [
    let choice random 3
    set color (ifelse-value
      who mod 3 = 0 [ red ]
      who mod 3 = 1 [ blue ]
      who mod 3 = 2 [ green ])
    set heading (random-float 180) - 90
    set initial-heading heading
  ]
end 

to setup-counters
  ;; setup counters for emitted and received rays of each colours, one counter for each bin of sky
  set red_em (n-values division-factor [1])
  set green_em (n-values division-factor [1])
  set blue_em (n-values division-factor [1])
  set red_rec (n-values division-factor [1])
  set green_rec (n-values division-factor [1])
  set blue_rec (n-values division-factor [1])
end 

to setup-simulated-rays
  ;; setup simulated rays with the calculated colours
  set-default-shape simulated-rays "line"
  let tot 180 * 3
  create-simulated-rays tot [setxy 0 min-pycor + 1]
  ask simulated-rays [
    set hidden? true
    set heading (((who - n_rays) * 180 / tot) - 90)
    set size 2
    set pen-mode "down"
    set pen-size 2
  ]
end 


;;; LOOP PROCEDURES

to increment-color-counter [escaped init-heading escape-heading tcolor]
  ;; increment colour of the correct bin for a ray of color tcolor and initial and escape heading

  let index get-index init-heading
  (
    ifelse
    tcolor = red   [set red_em replace-item index red_em   (item index red_em + 1)]
    tcolor = blue  [set blue_em replace-item index blue_em  (item index green_em + 1)]
    tcolor = green [set green_em replace-item index green_em (item index green_em + 1)]
)
  if escaped and (check-sun-collision escape-heading) [
    (
      ifelse
      tcolor = red   [set red_rec replace-item index red_rec   (item index red_rec + 1)]
      tcolor = blue  [set blue_rec replace-item index blue_rec  (item index green_rec + 1)]
      tcolor = green [set green_rec replace-item index green_rec (item index green_rec + 1)]
    )
]
end 

to setup-atmosphere
  ask patches [
    ifelse pycor + max-pycor < 2 * max-pycor * atmosphere_height / 100 [ set pcolor blue + 3.5 ]
    [ set pcolor black]
  ]
end 

to go
  ;; flow control of what to display and what to compute
  move-rays
  if-else draw-path
  [ask rays [pen-down]]
  [ask rays [pen-up]]
  setup-atmosphere
  ifelse show-result? [
    set draw-path false
    ask rays [set hidden? true]
    ask simulated-rays [
      set hidden? false
      set pen-mode "down"
      get-simulated-color
   ]
    move-simulated-rays
  ]
  [
    ask rays [set hidden? false]
    ask simulated-rays [
      set hidden? true
      set pen-mode "up"
    ]
    clear-drawing
  ]
  tick
end 

to reset-ray
  set pen-mode "up"
  set xcor 0
  set ycor min-pycor + 1
  set heading (random-float 180) - 90
  set initial-heading heading
  if draw-path [set pen-mode "down"]
end 

to move-rays
  ;; move and scatter r
  ask rays [
    if pcolor != black [
      if is-scattered color [
        right (heading + scatter-angle)
      ]
    ]
    forward light_path
    let oob is-oob ycor

    ;; escaped the atmosphere
    if oob = 1 [
      increment-color-counter true initial-heading heading color
      reset-ray
    ]

    ;; collided with ground
    if oob = 2 [
      increment-color-counter true initial-heading heading color
      reset-ray
    ]
  ]
end 

to get-simulated-color
  ask simulated-rays [
    let index get-index heading
    set color get-color
      (item index red_em)
      (item index red_rec)
      (item index green_em)
      (item index green_rec)
      (item index blue_em)
      (item index blue_rec)
  ]
end 

to move-simulated-rays
  ask simulated-rays [
    foreach (n-values 50 [1]) [
      forward 0.1
      if ((xcor < min-pxcor + 0.05) or
        (xcor > max-pxcor - 0.05) or
        (ycor > max-pycor - 0.05)
      ) [setxy 0 min-pycor + 1]
    ]
  ]
end