#!/usr/bin/env python3

from collections import OrderedDict
import csv
from datetime import datetime
from io import StringIO

from matplotlib import pyplot as plt
import requests


URL = 'https://data.cdc.gov/api/views/3rge-nu2a/rows.csv?accessType=DOWNLOAD'
COLUMNS = {
  'outcome': ('outcome', lambda s: s == 'case'),
  'Week date': ('date', lambda s: datetime.strptime(s.split('-')[0], '%b %d').replace(year=datetime.now().year)),
  'Age group': ('age', lambda s: s.startswith('all')),
  'Vaccine product': ('vaccine', lambda s: s.startswith('all')),
  'Vaccinated with outcome': ('breakthrough', int),
  'Fully vaccinated population': ('vaccinated', int),
  'Unvaccinated with outcome': ('infected', int),
  'Unvaccinated population': ('unvaccinated', int),
}
CALC = OrderedDict([
  ('breakthrough_rate', lambda row: row['breakthrough'] / row['vaccinated']),
  ('infected_rate', lambda row: row['infected'] / row['unvaccinated']),
  ('vaccine_factor', lambda row: row['infected_rate'] / row['breakthrough_rate']),
])


def main():

  (header, data) = fetch(URL)
  data = parse(header, data)
  graph(data)


def fetch(url):

  result = []

  data = StringIO(requests.get(url).text)
  for row in csv.reader(data):
    result.append(row)

  return (result[0], result[1:])


def parse(header, data):

  cols = OrderedDict()
  for (i, col) in enumerate(header):
    if col in COLUMNS:
      cols[i] = COLUMNS[col]

  result = []
  for row in data:
    new  = {}
    for (i, (name, func)) in cols.items():
      val = func(row[i])
      if isinstance(val, bool):
        if not val:
          break
      else:
        new[name] = val
    else:
      for (name, func) in CALC.items():
        new[name] = func(new)
      result.append(new)

  return result


def graph(data):

  x = [row['date'] for row in data]

  (fig, ax1) = plt.subplots()
  ax2 = ax1.twinx()

  ax1.plot(x, [100*row['infected_rate'] for row in data], c='green', label='Unvaccinated')
  ax1.plot(x, [100*row['breakthrough_rate'] for row in data], c='red', label='Vaccinated')
  ax1.set_ylabel('New Cases (% of Population)')
  ax1.grid(axis='x')

  ax2.plot(x, [row['vaccine_factor'] for row in data], c='blue', label='Factor')
  ax2.set_ylabel('Vaccinated Factor')
  ax2.set_ylim([0, None])
  ax2.grid(axis='y')

  fig.autofmt_xdate()

  ax1.legend()
  ax2.legend()
  plt.title('Breakthrough Cases (CDC)')
  plt.show()


if __name__ == '__main__':
  main()