#!/usr/bin/env python3

from math import pi,cos

def paperstrategy(day,distancing):
  if current['I'] >= 37.5/10000: distancing = 1
  if current['I'] <= lowwater/10000: distancing = 0
  return distancing

def chinastrategy(day,distancing):
  return 1 if day >= 22 else 0
  
for R0,distancingimprovement,seasonal,startday,pulsewidth,stopday,lowwater,strategy,outputfile in (
  (2.5,0.60,52,10,2.0,121,10.0,chinastrategy,'gigo-25-60-china.pdf'),
  (2.0,0.60,52,4,2.0,121,10.0,chinastrategy,'gigo-20-60-china.pdf'),
  (2.0,0.99,52,0,2.0,121,10.0,chinastrategy,'gigo-20-99-china.pdf'),
  (3.5,0.99,52,9,2.0,121,10.0,chinastrategy,'gigo-35-99-china.pdf'),
  (2.5,0.60,55,70,3.5,1096,10.0,paperstrategy,'gigo-25-60-55.pdf'),
  (2.5,0.60,52,70,3.5,1096,10.0,paperstrategy,'gigo-25-60-52.pdf'),
  (2.5,0.60,0,70,3.5,1096,10.0,paperstrategy,'gigo-25-60-0.pdf'),
  (2.0,0.60,55,70,3.5,1096,10.0,paperstrategy,'gigo-20-60-55.pdf'),
  (2.0,0.60,52,70,3.5,1096,10.0,paperstrategy,'gigo-20-60-52.pdf'),
  (2.0,0.60,0,70,3.5,1096,10.0,paperstrategy,'gigo-20-60-0.pdf'),
  (2.0,0.80,52,70,3.5,1096,10.0,paperstrategy,'gigo-20-80-52.pdf'),
  (2.0,0.99,52,70,3.5,1096,2.0,paperstrategy,'gigo-20-99-52.pdf'),
  (1.5,0.99,52,70,3.5,1096,2.0,paperstrategy,'gigo-15-99-52.pdf'),
):

  def seasonalimprovement(day):
    summerimprovement = 0.3
    wintershift = 3.8*7 # peak of winter before end of December
    angle = 2*pi*(day+wintershift)/(7.0*seasonal)
    return summerimprovement*0.5*(1-cos(angle))

  def paperbeta(day,distancing):
    R = R0
    if seasonal: R *= 1-seasonalimprovement(day)
    if distancing: R *= 1-distancingimprovement
    return R/5.0
  
  def paperinitialpulse(day,distancing):
    return 0.01/7 if day < startday+pulsewidth else 0

  transition = (
    ('S','E',paperbeta,'I'),
    ('S','E',paperinitialpulse),
    ('E','I',1/5.0), # E->I averaging 5 days
    ('I','HH',0.0308/5.0), # 3.08% chance I->HH, averaging 5 days
    ('I','HC',0.0132/5.0), # 1.32% chance I->HC, averaging 5 days
    ('I','R',0.956/5.0), # 95.6% chance I->R, averaging 5 days
    ('HC','C',1/8.0), # HC->C averaging 8 days
    ('HH','R',1/6.0), # HH->R averaging 6 days
    ('C','R',1/10.0), # C->R averaging 10 days
  )
  
  current = {}
  for t in transition:
    current[t[0]] = 0
    current[t[1]] = 0
  current['S'] = 1
  
  historyday = []
  historydistancing = []
  history = {c:[] for c in current}
  
  distancing = 0 # will always be 0 or 1
  day = startday
  while day < stopday:
  
    historyday += [day]
    historydistancing += [distancing]
    for c in current: history[c] += [current[c]]

    distancing = strategy(day,distancing)
  
    change = {c:0 for c in current}
  
    for t in transition:
      assert len(t) in [3,4]
      speed = t[2]
      if callable(speed): speed = speed(day,distancing)
      if len(t) == 4: speed *= current[t[3]]
      speed *= current[t[0]]
      change[t[0]] -= speed
      change[t[1]] += speed
  
    daydelta = 1/72.0
    if current['I'] < 37.5/10000:
      if current['I']+change['I']*daydelta > 37.5/10000:
        daydelta = (37.50001/10000-current['I'])/change['I']
    if current['I'] > lowwater/10000:
      if current['I']+change['I']*daydelta < lowwater/10000:
        daydelta = (0.999999*lowwater/10000-current['I'])/change['I']
  
    assert daydelta > 0
    day += daydelta
    for c in current:
      current[c] += change[c]*daydelta
  
  
  # ----- plotting
  
  import datetime
  xmin = datetime.datetime(2020,1,1)
  xmax = datetime.datetime(2020,1,1)+datetime.timedelta(days=stopday)
  days = [xmin+datetime.timedelta(days=x) for x in historyday]
  
  import matplotlib.pyplot as plt
  import matplotlib.dates as mdates
  from matplotlib.backends.backend_pdf import PdfPages
  
  fig,ax1 = plt.subplots()
  ax2 = ax1.twinx()
  ax3 = ax2.twinx()
  fig.set_size_inches(8,2.5)
  
  ax1.set_xlim([xmin,xmax])
  ax2.set_xlim([xmin,xmax])
  ax3.set_xlim([xmin,xmax])

  ax1.set_ylim([0,3.0])
  ax1.fill_between(days,[3.0*y for y in historydistancing],color='lightblue')
  ax1.tick_params(axis='y',labelcolor='orange',labelleft=False,labelright=True,pad=30)
  ax1.plot(days,[10000*(y1+y2) for y1,y2 in zip(history['HH'],history['HC'])],color='orange')
  
  ax2.set_ylim([0,1.2])
  ax2.tick_params(axis='y',labelcolor='red',labelleft=False,labelright=True)
  if strategy == paperstrategy:
    ax2.axhline(y=0.89,xmin=0,xmax=1,color='tomato')
    ax2.axhline(y=37.5/50.0,xmin=0,xmax=1,color='gray',linestyle='--')
    ax2.axhline(y=lowwater/50.0,xmin=0,xmax=1,color='gray',linestyle='--')
  ax2.plot(days,[1-y for y in history['S']],color='green')
  ax2.plot(days,[10000*y for y in history['C']],color='red')

  ax3.set_ylim([0,60.0])
  ax3.tick_params(axis='y',labelcolor='black',labelleft=True,labelright=False)
  ax3.plot(days,[10000*y for y in history['I']],color='black')
  
  ax1.xaxis.set_major_locator(mdates.MonthLocator())
  ax1.xaxis.set_major_formatter(mdates.DateFormatter('%b'))
  plt.setp(ax1.get_xticklabels(),rotation=90)

  fig.tight_layout()
  with PdfPages(outputfile) as pdf:
    pdf.savefig()
    plt.close()


  # ----- statistics

  maxcritical = max(history['C'])
  maxhospital = max(y1+y2 for y1,y2 in zip(history['HH'],history['HC']))

  previousday = startday
  distancingdays = 0
  halfcriticaldays = 0
  lastmaxcritical = 0
  halfhospitaldays = 0
  lastmaxhospital = 0
  for day,distancing,critical,hh,hc in zip(historyday,historydistancing,history['C'],history['HH'],history['HC']):
    if distancing:
      distancingdays += day-previousday
    if critical >= 0.5*maxcritical:
      halfcriticaldays += day-previousday
    if critical == maxcritical:
      lastmaxcritical = day
    hospital = hh+hc
    if hospital >= 0.5*maxhospital:
      halfhospitaldays += day-previousday
    if hospital == maxhospital:
      lastmaxhospital = day
    previousday = day

  print('%s: maxcritical on day %.6f, %.6f days >=0.5*maxcritical, %.6f days >=0.5*maxhospital, %.6f days distancing'
    % (outputfile,lastmaxcritical,halfcriticaldays,halfhospitaldays,distancingdays))