#!/usr/bin/perl

# model the spread of a pandemic
# Andrew Daviel, Nov 2009

# generate 2 output files: "pan.gif", an animated GIF, and
# "pan0.dat", a summary file suitable for gnuplot

# Initially populate a grid with uninfected people seeded
# with a few infected ones
# Each loop (one "day") people move around according to their mobility
# index. Except the dead ones, or if they stay home
# If they come close to an uninfected person, they may infect them

$steps = shift or $steps = 50 ;

$type = shift or $type = 0 ;
$type[0] = 'none' ; $type[1] = 'vulnerable' ;
$type[2] = 'mobile' ; $type[3] = 'random' ;

$size = shift or $size = 200 ;
$ndose = shift  or $ndose = 500 ;

use GD;
$| = 1 ;

$spont = 0.05 ; # chance of spontaneous development
$morb = 0.003 ; # chance of dying if infected
$maxd = 4 ; # max NY distance to catch
$trec = 10 ; # recovery time
$tinf = 7 ; # till realize sick
$pinf = 0.5 ; # probability of infection if meet
#$ndose = 500 ; # doses of vaccine available
$maxi = 15 ; # max initial infections

$maxx = $maxy = $size ;
$im = new GD::Image($maxx,$maxy);
$white = $im->colorAllocate(255,255,255); # background
$black = $im->colorAllocate(0,0,0);
$red = $im->colorAllocate(255,0,0);
$blue = $im->colorAllocate(0,0,255);
$blue2 = $im->colorAllocate(0,0,155);
$green = $im->colorAllocate(0,255,0) ;
$yellow = $im->colorAllocate(255,255,0) ;

# setup
$i = 0 ;
for ($x=1;$x<$maxx;$x+=5) {
  for ($y=1;$y<$maxy;$y+=5) {
    $x[$i] = $x ; $y[$i] = $y ; # initial position
    $health[$i] = 0 ; $tinf[$i] = 0 ;
    $mov[$i] = 1 ;
    $mov[$i] = 2 + 150* rand(1) * rand(1) * rand(1) ; # mobility
    $morb[$i] = 1 * rand(1) * rand(1) * rand(1) ; # vulnerability
    if (rand(1) < $spont and $ni<$maxi ) { $health[$i] = 1 ; $ni++ ; }
    $i++ ;
  }
}
open (OUT,">pan0.dat") ;
print OUT "# k x y morb health mov  (strategy=$type[$type])\n" ;
print OUT "# tinf $tinf pinf $pinf trec $trec maxi $maxi\n" ;
$N = $i ;
$i = 0 ;
for ($k=0;$k<$N;$k++) {
  if ($morb[$k] > 0.15 and $type == 1) { $health[$k] = 4 ; $i++ ; } # vaccinate vulnerable
  if ($mov[$k] > 20 and $type == 2) { $health[$k] = 4 ; $i++ ; } # vaccinate mobile
  if (rand(1)<0.3 and $type == 3) { $health[$k] = 4 ; $i++ ; } # vaccinate randomly
  if ($i>=$ndose) { last ; } # out of vaccine
  print OUT "$k $x[$k] $y[$k] $morb[$k] $health[$k] $mov[$k] \n" ;
}
close(OUT) ;
$nv = $i ;

#$health[15] = $health[139] = 1 ; # 1 infection
#$mov[15] = $mov[39] =  $mov[79] = $mov[139] = 16 ; #mobile

print "# $N people $ni infected, strategy $type ($type[$type])\n" ;
print "# T sick immune dead vaccinated\n" ;
print "# tinf $tinf pinf $pinf trec $trec maxi $maxi\n" ;

my $gifdata = $im->gifanimbegin;
$gifdata   .= $im->gifanimadd;    # first frame
for ($t=1;$t<$steps;$t++) {
  my $frame  = GD::Image->new($im->getBounds);
  $white = $frame->colorAllocate(255,255,255); # background
  $black = $frame->colorAllocate(0,0,0);
  $red = $frame->colorAllocate(255,0,0);
  $blue = $frame->colorAllocate(0,0,255);
  $blue2 = $frame->colorAllocate(0,0,155);
  $green = $frame->colorAllocate(0,255,0) ;
  $yellow = $frame->colorAllocate(255,255,0) ;

  $nsick = $nimmune = $nhealthy = $ndead = 0 ;

  # $health 0 healthy 1 sick 2 dead 3 immune 4 = vaccinated
  $color[0] = $green ; $color[1] = $red ; $color[2] = $black ;
  $color[3] = $blue ; $color[4] = $blue2 ;

  for ($i=0;$i<$N;$i++) {
    if ($health[$i] == 0) { $nhealthy++ ; }
    if ($health[$i] == 1) { $nsick ++ ; }
    if ($health[$i] == 2) { $ndead++ ; }
    if ($health[$i] == 3) { $nimmune++ ; }
    $frame->filledEllipse($x[$i],$y[$i],2,2,$color[$health[$i]]) ;
  }
  $gifdata .= $frame->gifanimadd(0,0,0,5);     # add frame, 100ms

  print "$t $nsick $nimmune $ndead $nv\n" ;
  if ($nsick == 0) { last ; }
  undef (@XY) ;
# speed up neighbour search by populating array of "geographic" points
# so to find possible new infectees, do not need to compute distance
# to everyone else but just to close-by people
  for ($k=0;$k<$N;$k++) {
    $XY[$x[$k]][$y[$k]] .= "$k," ;
  }

  for ($k=0;$k<$N;$k++) {
    if ($health[$k] == 1 and ($t - $tinf[$k]) > $trec) {
      if (rand(1) < $morb[$k]) {
        $health[$k] = 2 ; # die
      } else {
        $health[$k] = 3 ; # recover
      }
    }
    for ($nx=$x[$k]-$maxd;$nx<=$x[$k]+$maxd;$nx++) {
      if ($nx<0 or $nx>$maxx) { next ; }
      for ($ny=$y[$k]-$maxd;$ny<=$y[$k]+$maxd;$ny++) {
        if ($ny<0 or $ny>$maxy) { next ; }
	foreach $k2 (split(',',$XY[$nx][$ny])) {
	  if ($k == $k2) { next ; }
          if ($k2 eq '') { next ; }
	  $nyd = abs($x[$k]-$x[$k2]) + abs($y[$k]-$y[$k2]) ;
	  if ($nyd > $maxd ) { next ; }
	  if ($health[$k] == 0 and $health[$k2] == 1 and rand(1) < $pinf
# stay in bed
           and ($t - $tinf[$k2]) < $tinf ) {
	    $health[$k] = 1 ; $tinf[$k] = $t ;
	    next ;
	  }
	}
      }
    }
 #   if ($health[$k] == 1 and rand(1) < $morb[$k]) { $health[$k] = 2 ; }
    
    unless ($health[$k] == 2 or ( $health[$k] ==1 and ($t-$tinf[$k] > $tinf)) ) {
      $mov2 = $mov[$k]*2 + 1 ;
      $x[$k] += int(rand($mov2))-$mov[$k] ;
      $y[$k] += int(rand($mov2))-$mov[$k] ;
      if ($x[$k] < 0) { $x[$k] += $maxx  ; }
      if ($y[$k] < 0) { $y[$k] += $maxy  ; }
      if ($x[$k] > $maxx) { $x[$k] -= $maxx ; }
      if ($y[$k] > $maxy) { $y[$k] -= $maxy ; }
    }
  }

}
$gifdata .= $im->gifanimend;   # finish the animated GIF
open (OUT,">pan.gif") ;
print OUT $gifdata;

#print OUT $im->png;