#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ***************************************************************************
# *   Copyright (C) 2020, Paul Lutus                                        *
# *                                                                         *
# *   This program is free software; you can redistribute it and/or modify  *
# *   it under the terms of the GNU General Public License as published by  *
# *   the Free Software Foundation; either version 2 of the License, or     *
# *   (at your option) any later version.                                   *
# *                                                                         *
# *   This program is distributed in the hope that it will be useful,       *
# *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
# *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
# *   GNU General Public License for more details.                          *
# *                                                                         *
# *   You should have received a copy of the GNU General Public License     *
# *   along with this program; if not, write to the                         *
# *   Free Software Foundation, Inc.,                                       *
# *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
# ***************************************************************************

# NOTE: I apologize for all the special flags and states in this program.
# I use it to generate custom graphs and videos for my YouTube video series.

# 2020.07.09 change order of plotted zones to (bottom to top) I,R,S

# user entries:

# frames: 0 = no animation, any > 0 = animate a sequence

frames = 0

# mode: 0 = simple 3-element overlay plot,  1 = vertical array of 3 plots

mode = 0

# fill 0 = lines,1 = filled areas

fill = 1

# vertically stack the filled data zones, only works with fill = 1

stack = 1

# speciala: render only the I group against a
# threshold for healthcare system overload
# for an example video

speciala = 0

# initial, global simulator settings

R_0 = 2.5 # infection time rate
R_1 = .25 # recovery time rate

# sequester values

sequester = True and not speciala
Ss = 10 # start day
Se = 70 # end day
Sp = 90 # percentage (of population or effectiveness)

# days in simulation

maxD = 180

# matplotlib settings

linewidth = 3

# end of user configuration
  
# imports

import os, sys, re, math

import matplotlib

import matplotlib.pyplot as plt

import matplotlib.lines as ml

import matplotlib.ticker as mt

import matplotlib.patches as patch

import matplotlib.gridspec as gridspec

# interpolate two ranges for given x: xa..xb -> ya..yb

def ntrp(x,xa,xb,ya,yb): return (x-xa) * (yb-ya) / (xb-xa) + ya 

# create a quadratic attack-decay envelope

def quadratic_envelope(x,xa,xb):
    ya = ntrp(x,xa,xb,-1,1)
    yb = 1.0-abs(ya)
    s = math.copysign(1,ya)
    yc = -yb*yb * s + s
    return ntrp(yc,-1,1,xa,xb)

alpha = '80'

scolor = '#4060ff'
scolora = scolor #+ alpha
rcolor = '#408040'
rcolora = rcolor #+ alpha
icolor = '#c00000'
icolora = icolor #+ alpha

seq_face = '#00000020'

seq_edge = '#000000'

def calculate():
  global maxD,sa,ia,ra,xref,yref,infx,infy,Ss,Se,Sp,frame
  
  # initial values
  
  I = 0.01 # infected
  S = 1.0 - I # susceptible
  R = 0.0 # recovering
  
  infx = 0
  infy = 0
  sa = []
  ia = []
  ra = []
  xref = []
  yref = []
  
  # create a normalized sequester value
  sp_norm = 1-Sp/100.0
  
  # total count of iterations
  nn = int(maxD * dt/div)
   
  for n in range(nn):
    
    d = n / div
    
    # this is the core numerical DE
    
    # detect sequester zone
    q = (1,sp_norm)[sequester and d >= Ss and d < Se]
    
    # compute di/dt in advance of its application
    # this greatly improves numerical accuracy
    z = (S * I * R_0 * q - R_1 * I) / dt
    
    S -= S * I * R_0 * q / dt
    I += z
    R += R_1 * I / dt
    
    # save a peak value for later plotting
    if I >= infy:
      infy = I
      infx = n
  
    # accumulate intermediate results
    sa += [S]
    ia += [I]
    ra += [R]
    xref += [n]
    # silly, but needed by matplotlib
    yref += [0]
    
  if frames > 0:
    print("%16.2f %16.2f %16.2f" % (framenum,fx,infy * 100))
    
def sequester_title(ax):
  if sequester and Se-Ss > .01:
    a = Ss * 10
    b = Se * 10
    c = b-a
    seqr = patch.Rectangle((a,0),c,1,facecolor=seq_face,edgecolor=seq_edge,label='Sequester',linewidth=2)
    ax.add_patch(seqr)
    s = 'Sequester %.0f%%' % Sp
    ax.text((a+b)/2,1.02,s,ha='center',fontweight=fwb)  
      
def single_plot(ax):
  if fill:
    ax.grid(color='#404040')
    if stack:
      mid = []
      top = []
      for i in range(len(sa)):
        mid += [ra[i]+ia[i]]
        top += [sa[i] + mid[i]]
      ax.fill_between(xref,top,mid, facecolor=scolora)
      ax.fill_between(xref,mid,ia, facecolor=rcolora)
      ax.fill_between(xref,ia,yref, facecolor=icolora)
      plt.legend(('Susceptible','Recovering','Infected'),loc='lower right',fontsize=18)
    else:
      if not speciala:
        ax.fill_between(xref,sa,yref, facecolor=scolora)
      ax.fill_between(xref,ia,yref, facecolor=icolora)
      if not speciala:
        ax.fill_between(xref,ra,yref, facecolor=rcolora)
      if speciala:
        plt.legend(('Infected',),loc='upper right',fontsize=18)
      else:
        plt.legend(('Susceptible','Infected','Recovering'),loc='upper right',fontsize=18)
  else:
    ax.grid()
    if not speciala:
      ax.plot(sa, color=scolor,linewidth=linewidth)
    ax.plot(ia, color=icolor,linewidth=linewidth)
    if not speciala:
      ax.plot(ra, color=rcolor,linewidth=linewidth)
    if speciala:
      plt.legend(('Infected',),loc='upper right',fontsize=18)
    else:
      plt.legend(('Susceptible','Infected','Recovering'),loc='upper right',fontsize=18)
  if fill and stack:
    vl = ml.Line2D([infx,infx],[0,1],color='black')
  else:
    vl = ml.Line2D([infx,infx],[infy-.04,infy+.12],color='black')
    
  ax.add_line(vl)
  
  if speciala:
    spc = '#208020'
    vl = ml.Line2D([0,maxD * div],[.2,.2],color=spc,linewidth=3,linestyle='-.')
    ax.text(maxD*8.4,.22,'Hospital Beds',color=spc,fontsize=22,weight='bold')
    ax.add_line(vl)
  ax.xaxis.set_major_formatter(
     mt.FuncFormatter(lambda x, pos: '%d' % (x/10.0)))
  ax.yaxis.set_major_formatter(
     mt.FuncFormatter(lambda x, pos: '%d%%' % (x*100)))
  ax.set_xticks(xt)
  ax.set_ylabel('Population', fontsize=fsb)
  sequester_title(ax)
  ax.set_ylim(yrange)

def multi_plot(ax,data,color,title):
  ax.grid()
  line = False
  if fill:
    line = ax.fill_between(xref,data,yref, facecolor=color)
  else:
    line, = ax.plot(data, color=color,linewidth=linewidth)
  line.set_label(title)
  ax.legend(fontsize=12)
  if(data == ia):
    vl = ml.Line2D([infx,infx],[infy-.04,infy+.12],color='black')
    ax.add_line(vl)
  ax.xaxis.set_major_formatter(
     mt.FuncFormatter(lambda x, pos: '%d' % (x/10.0)))
  ax.yaxis.set_major_formatter(
     mt.FuncFormatter(lambda x, pos: '%d%%' % (x*100)))
  ax.set_xticks(xt)
  #ax.set_ylabel('Population', fontsize=fsb)
  sequester_title(ax)
  ax.set_ylim(yrange)


def generate():
  global xt
  
  plt.close()
  
  xt = [int(ntrp(n,0,6,0,maxD*10)) for n in range(7)]
  
  calculate()
      
  if mode == 0:
  
    plt.rcParams.update({'font.size': 20})
    fig, ax = plt.subplots(figsize=(19.2,10.8))
   
    plt.title('Pandemic Time Series',fontsize=fsa)
    single_plot(ax)
    
    if speciala:
      
      ax.text(infx+10,infy+.08,'R_0: %.1f' % (R_0),color='black',fontsize=fsc,weight=fwb)
      ax.text(infx+10,infy+.02,'%.0f%% Infected' % (infy*100),color='black',fontsize=fsc,weight=fwb)
    else:
      ax.text(infx+10,infy+.14,'R_0: %.1f' % (R_0),color='black',fontsize=fsc,weight=fwb)
      # plot the peak infection time and rate
      ax.text(infx+10,infy+.08,'Day %d' % (int(infx/10)),color='black',fontsize=fsc,fontweight=fwb)
      ax.text(infx+10,infy+.02,'%.0f%% Infected' % (infy*100),color='black',fontsize=fsc,fontweight=fwb)
    
  elif mode == 1:
    titles = ('Susceptible','Infected','Recovering')
    plt.rcParams.update({'font.size': 8})
    fig, axs = plt.subplots(3,figsize=(10,10))
  
    data = (sa,ia,ra)
    colors = (scolor,icolor,rcolor)
    fig.suptitle('Pandemic Time Series',fontsize=fsa)
    for i in range(3):
      multi_plot(axs[i],data[i],colors[i],titles[i])
    
  plt.xlabel('Days', fontsize=fsb)

# global declarations
  
yrange = (-.1,1.1)

xt = False
  
if mode == 0:
  fsa = 32
  fsb = 26
  fsc = 20
elif mode == 1:
  fsa = 16
  fsb = 12
  fsc = 8
    
dt = 100.0
div = 10.0


infx = 0
infy = 0
sa = []
ia = []
ra = []
xref = []
yref = []

fwb = 'bold'
  
if speciala:
  R_0 = .56
  maxD = 240
    
if frames == 0:
  # generate one plot
  generate()   
  
  plt.savefig('pandemic_time_series.png')
  plt.show()
  
else:
  # generate animation frames
  os.system('mkdir frames > /dev/null 2>&1')
  if os.path.exists('frames/frame0000.png'):
    reply = input('Delete prior frames (y/n):')
    if len(reply) > 0 and reply[0].lower() == 'y':
      os.system('rm frames/*')
      
  print("%16s %16s %16s" % ('Framenum','Frame','Peak infected'))
  for framenum in range(0,frames+1):
    # create a natural-velocity envelope
    fx = quadratic_envelope(framenum,0,frames)
    if speciala:
      R_0 = ntrp(fx,0,frames,5,.57)
    else:
      # create the variable that changes the animation
      #Se = ntrp(fx,0,frames,Ss,70)
      Se = 70
      R_0 = ntrp(fx,0,frames,2.5,.57)
      maxD = 360
    generate()
    plt.savefig('frames/frame%04d.png' % framenum) 
  # create video from stills
  os.system('which ffmpeg && ffmpeg -framerate 30 -y -i frames/frame%04d.png coronavirus_video.mp4')

  # acquire representative frames as stills
  os.system('cp frames/frame0000.png R_minimum.png')
  os.system('cp frames/frame%04d.png R_maximum.png' % frames)
  os.system('cp frames/frame%04d.png R_middle.png' % int(frames/2))