#! /usr/bin/env python3
#
def sir_solve_ivp ( ):

#*****************************************************************************80
#
## sir_solve_ivp() solves the SIR ODE using solve_ivp().
#
  from scipy.integrate import solve_ivp
  from sir_conserved import sir_conserved
  from sir_dydt import sir_dydt
  import matplotlib.pyplot as plt
  import numpy as np

  print ( '' )
  print ( 'sir_euler:' )
  print ( '  Solve the SIR ODE using solve_ivp()' )

  t0 = 0.0
  tstop = 7500.0
  tspan = np.array ( [ t0, tstop ] )
  y0 = np.array ( [ 99.0, 1.0, 0.0 ] )

  sol = solve_ivp ( sir_dydt, tspan, y0 )

  y = np.transpose ( sol.y )
#
#  Time plot.
#
  plt.clf ( )
  plt.plot ( sol.t, y, linewidth = 3 )
  plt.grid ( True );
  plt.xlabel ( '<-- Time -->' )
  plt.ylabel ( '<-- Cases -->' )
  plt.legend ( [ 'Susceptible', 'Infected', 'Recovered' ] )
  plt.title ( 'SIR disease model, time plot' )
  filename = 'sir_solve_ivp_time.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )
  plt.show ( )
  plt.close ( )
#
#  Conservation plot.
#
  P = sir_conserved ( sol.t, y )

  plt.clf ( )
  plt.plot ( sol.t, P, 'c-', lineWidth = 3 )
  plt.plot ( [ t0, tstop ], [ 0.0, 0.0 ], 'k-', linewidth = 3 )
  plt.grid ( 'on' )
  plt.xlabel ( '<-- Time -->' )
  plt.ylabel ( '<-- P(t) -->' )
  plt.title ( 'SIR disease model, conservation plot' )
  filename = 'sir_solve_ivp_conservation.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )
  plt.show ( )
  plt.close ( )
#
#  Report final values.
#
  dydt = sir_dydt ( tstop, y[-1,:] )

  print ( '' )
  print ( '  At final time t = ', tstop )
  print ( '  S = ', y[-1,0], '  dSdt = ', dydt[0] )
  print ( '  I = ', y[-1,1], '  dIdt = ', dydt[1] )
  print ( '  R = ', y[-1,2], '  dRdt = ', dydt[2] )

  return

if ( __name__ == "__main__" ):
  sir_solve_ivp ( )