<gh_stars>1-10
#!/usr/bin/env python
# coding: utf-8

# # COVID-19 model for policy makers in the United Kingdom

#  <p align="center">
#  <img src="images/dynamicalModel.png" width="70%">
#  </p>

# We use an extended [SEIR model](https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#The_SEIR_model) to capture available information both about the disease progression, as well as how accessible various disease states are by testing. Being tested might cause a transition in the <span style="display: inline-block;background-color:#FFFCCC">Testing State</span>; the rate of such a transition depends both on the <span style="display: inline-block;background-color:#D1E2FF">Health State</span> as well the parameters of the test used.
#
# Due to age being both an extra risk factor as well as a potential way for decision makers to clearly apply policies to only parts of the total population, we include it directly in our model, and most transition rates are strongly influenced by <span style="display: inline-block;background-color:#FFD4D1">Age State</span>.
#
# Finally, the main policy making tool as well as conundrum nowadays is the implementation of quarantining and social distancing in order to keep hospitals and medical staff under tolerable levels of pressure. We represent <span style="display: inline-block;background-color:#C2EDC0">Isolation States</span> to investigate the effects of various interventions on policy targets relating to hospitalisation rates and economic freedom, while describing the different health outcomes via the <span style="display: inline-block;background-color:#D1E2FF">Health States</span>.
#

# ### Health states and Disease progression
#
# Susceptible people may be Exposed to the virus by mixing with other infected people (E,A,I<sub>1</sub> and I<sub>2</sub>).
#
# They may get through the infection Asymptomatic and eventually Recover, or become symptomatic and more Infectious, spreading the disease, and potentially Dying or Recovering.
#
# Recovered people develop more effective antibodies against the virus, and are considered immune<sup>*</sup>.
#
#
# | State | Description | Testing |
# | ----- | ----------- | -------- |
# | S | Susceptible | Negative |
# | E | Exposed | Very weakly virus positive
# | A | Asymptomatic | Weakly virus positive
# | I<sub>1</sub> | Symptomatic early | Strongly virus positive
# | I<sub>2</sub> | Symptomatic late | Medium virus positive <br>Weakly IgM antibody positive
# | R<sub>1</sub> | Recovered early | IgM antibody positive
# | R<sub>2</sub> | Recovered late | IgM/IgG antibody positive
# | D | COVID-related death | May be virus or antibody positive
#
#
# <sub><sup>*</sup>We plan to consider partial / short-term immunity, see further discussion in [Continued research](#header_contd).</sub>

# # Code
#
# The below consists of the implementation of the model described above, with extended comments including assumptions, sources of data and areas to improve upon

# In[1]:


# Use available data up until this day; cutoff is important due to more recent data being less complete.
CONST_DATA_CUTOFF_DATE = "20200414"


# # Packages and Helper functions
#
# To preserve the single-notebook formulation, we include all packages used as well as subfunctions here
#
# To skip to the start of model implementation, <a href="#modelImplementation">click here</a>!

# In[2]:


# Packages

# Basic packages
import numpy as np

from scipy import integrate, stats, spatial
from scipy.special import expit, binom

import pandas as pd
import xlrd # help read excel files directly from source into pandas

import copy
import warnings

# Building parameter/computation graph
import inspect
from collections import OrderedDict

# OS/filesystem tools
import time
from datetime import datetime
import random
import string
import os
import shutil
import sys
import cloudpickle

# Distributed computing tools
import dask
import distributed
from dask.distributed import Client
from dask.distributed import as_completed
import itertools


# In[3]:


# Regroup various age-group representations into our internal one, and vice versa
def regroup_by_age(
    inp, # first dimension is ages, others dont matter.
    fromAgeSplits, toAgeSplits, maxAge=100., maxAgeWeight = 5.):
    fromAgeSplits = np.concatenate([np.array([0]), fromAgeSplits, np.array([maxAge])]) # Add a zero at beginning for calculations
    toAgeSplits = np.concatenate([np.array([0]), toAgeSplits, np.array([maxAge])]) # Add inf at end for calculations
    def getOverlap(a, b):
        return max(0, min(a[1], b[1]) - max(a[0], b[0]))
    out = np.zeros((len(toAgeSplits)-1,)+inp.shape[1:])
    for from_ind in range(1, len(fromAgeSplits)):
        # Redistribute to the new bins by calculating how many years in from_ind-1:from_ind falls into each output bin
        cur_out_distribution = (
        [getOverlap(toAgeSplits[cur_to_ind-1:cur_to_ind+1],fromAgeSplits[from_ind-1:from_ind+1])  for cur_to_ind in range(1, len(toAgeSplits))]
        )

        if cur_out_distribution[-1] > 0:
            cur_out_distribution[-1] = maxAgeWeight # Define the relative number of ages if we have to distribute between second to last and last age groups

        cur_out_distribution = cur_out_distribution/np.sum(cur_out_distribution)

        for to_ind in range(len(out)):
            out[to_ind] += cur_out_distribution[to_ind] * inp[from_ind-1]

    return out


# PARAMETER DICTIONARIES AND TABLES
# -----------------------------------------------------------------------------------------

# Build the nested parameter/computation graph of a single function.
def build_paramDict(cur_func):
    """
    This function iterates through all inputs of a function,
    and saves the default argument names and values into a dictionary.

    If any of the default arguments are functions themselves, then recursively (depth-first) adds an extra field to
    the dictionary, named <funcName + "_params">, that contains its inputs and arguments.

    The output of this function can then be passed as a "kwargs" object to the highest level function,
    which will then pass the parameter values to the lower dictionary levels appropriately
    """

    paramDict = OrderedDict()

    allArgs = inspect.getfullargspec(cur_func)

    # Check if there are any default parameters, if no, just return empty dict
    if allArgs.defaults is None:
        return paramDict


    for argname, argval in zip(allArgs.args[-len(allArgs.defaults):], allArgs.defaults):
        # Save the default argument
        paramDict[argname] = argval
        # If the default argument is a function, inspect it for further

        if callable(argval):
            # print(argname)
            paramDict[argname+"_params"] = build_paramDict(argval)

    return paramDict


# Do a mapping between dictionary and parameter table row and vice versa (for convenient use)

# Flatten the dictionary into a table with a single row (but many column):
def paramDict_toTable(paramDict):
    paramTable = pd.DataFrame()
    def paramDictRecurseIter(cur_table, cur_dict, preString):
        # Iterate through the dictionary to find all keys not ending in "_params",
        # and add them to the table with name <preString + key>
        #
        # If the key doesn end in "_params", then append the key to preString, in call this function on the value (that is a dict)
        for key, value in cur_dict.items():
            if key.endswith("_params"):
                paramDictRecurseIter(cur_table, value, preString+key+"_")
            else:
                paramTable[preString+key] = [value]

        # For the rare case where we want to keep an empty dictionary, the above for cycle doesn't keep it
        if len(cur_dict)==0:
            paramTable[preString] = [OrderedDict()]

        return cur_table

    return paramDictRecurseIter(paramTable, paramDict, preString="")


def paramTable_toDict(paramTable, defaultDict=None):
    # enable to pass a default dict (if paramTable is incomplete), in which we'll just add / overwrite the values
    paramDict = defaultDict if defaultDict is not None else OrderedDict()
    def placeArgInDictRecurse(argName, argVal, cur_dict):
        # Find all "_params_" in the argName, and for each step more and more down in the dictionary
        strloc = argName.find("_params_")
        if strloc == -1:
            # We're at the correct level of dictionary
            cur_dict[argName] = argVal
            return cur_dict
        else:
            # step to the next level of dictionary
            nextKey = argName[:strloc+len("_params_")-1]
            nextArgName = argName[strloc+len("_params_"):]
            if not nextKey in cur_dict:
                cur_dict[nextKey] = OrderedDict()
            placeArgInDictRecurse(nextArgName, argVal, cur_dict[nextKey])

        return cur_dict

    for key in paramTable.columns:
        paramDict = placeArgInDictRecurse(key, paramTable.at[0,key], paramDict)

    return paramDict


# <a name="modelImplementation"></a>
# # Model implementation
#
# The model is governed by these two main tensors:
# - State tensor: a 4th order tensor containing axes:
#     - Age groups
#     - Health states
#     - Isolation states
#     - Testing states
#
#     In our extended state tensor formulation, we also keep track of not just people currently in each state, but also people newly arriving to each state, as a large number of available datasets refer to "new cases" or "new hospitalisations" each day, rather than current state occupancies normally represented by ODE solutions.
#
#
# - Dynamic state transition rate tensor
#     - Rates that govern all allowed transitions between states
#     - Dynamically recomputed every iteration, based on number of infected, their current social mixing and number and types of tests available, amongst other variables.
#     - Briefly:
#         - No transition between age groups
#         - No transitions between testing states without performing a test
#         - No transitions into S or out of D and R_IgG health states
#     - Allowed transitions are as showcased in the model image above
#     - Represented by a 7th order, sparse tensor, containing all transitions except age (unimportant due to relatively short time scales compared to coarse age grouping)
#

# ```
# NOTICE
# ------
# THE "MODEL IMPLEMENTATION" SECTION CONTAINS A LARGE NUMBER OF PARAMETER VALUES SET TO A DEFAULT VALUE.
# THESE ARE LARGELY INFORMED BY DATA, BUT NOT FIXED!
# THEY ARE VARIED DURING THE FITTING OF THE MODEL, ACCORDING TO HOW UNCERTAIN WE ARE IN THE PARAMETER
# ```
#
# The priors are defined <a href="#defineEnsemblePriors">below the model</a>. Although many of our uncertain / weak assumptions are signalled by "TODO" comments, we feel that the overall conclusions would not be affected by finding better parameter values, especially given our fitting and exploration approach.
#

# ### The state tensor

# In[4]:


# State Dimensions

# Health states (S, E and D are fixed to 1 dimension)
nI_symp = 2 # number of sympyomatic infected states
nI = 2+nI_symp # number of total infected states (disease stages), the +2 are Exposed and I_nonsymptomatic
nR = 2 # number of recovery states (antibody development post-disease, IgM and IgG are two stages)
nHS = 2+nI+nR # number of total health states, the +2: S, D are suspectible and dead

# Age groups (risk groups)
nAge = 9 # In accordance w Imperial #13 report (0-9, 10-19, ... 70-79, 80+)

# Isolation states
nIso = 4 # None/distancing, Case isolation, Hospitalised, Hospital staff

# Testing states
nTest = 4 # untested/negative, Virus positive, Antibody positive, Both positive


stateTensor = np.zeros((nAge, nHS, nIso, nTest))


# ### Transition rates (1 / day)
#
# The full transition rate structure is an 8th order tensor,
# mapping from any 4D state in the state tensor, to any other 4D state in the state tensor
#
# However, many of these transitions are non-sensical (ie a 3 year old cannot suddenly become 72, or a dead person be infected again), therefore during the construction of the full model below, we fill in the rates on all "allowed" transitions.
#
# We attempt to do so based on existing data either describing particular rates (like COVID-related hospitalisation),
# or data that helps constrain the product or ratios of multiple rates (such as the R0, or the case fatality ratio [noting this latter depends on testing policy and test availability]).
#
# Further, to massively reduce the number of weakly constrained parameters, we will approximate many of the weakly correlated transition rates as rank 1 (uncorrelated) matrices. For example the rate of hospitalisation for a patient at a given age and stage of infection will be computed as a product of two indepent rates, one based purely on the age (older people are generally more at risk of hospitalisation), and the other purely on how far the patient has progressed into the disease. This allows us to estimate more of required parameters from available published data.
#
# There of course still is a lot of uncertainty about how the virus behaves, and all of the data that we use is likely uncomplete and noisy. In order to better represent the things we do not know, we use advanced machine learning techniques, and investigate many possible scenarios (settings of parameters) and for all analyses we retain all plausible scenarios (various parameter settings that explain the available data well enough).
#
# Any policies we suggest for the near future are investigated for all plausible scenarios, such that decision makers know how likely each one will work as expected in these uncertain times. We further note that as we progress further into the pandemic, the number of plausible scenarios reduces more and more, enabling us to see the way out clearer and clearer.
#
#

# In[5]:


# Population (data from Imperial #13 ages.csv/UK)
agePopulationTotal = 1000.*np.array([8044.056,7642.473,8558.707,9295.024,8604.251,9173.465,7286.777,5830.635,3450.616])
#agePopulationTotal = 1000.*pd.read_csv("https://raw.githubusercontent.com/ImperialCollegeLondon/covid19model/master/data/ages.csv").iloc[3].values[2:]
# Currently: let's work with england population only instead of full UK, as NHS England + CHESS data is much clearer than other regions
agePopulationTotal *= 55.98/66.27 # (google england/uk population 2018, assuming age dist is similar)
agePopulationRatio = agePopulationTotal/np.sum(agePopulationTotal)

# Helper function to adjust average rates to age-aware rates
def adjustRatesByAge_KeepAverageRate(rate, ageRelativeAdjustment, agePopulationRatio=agePopulationRatio, maxOutRate=10):
    """This is a helper function and wont be picked up as a model parameter!"""
    if rate == 0:
        return np.zeros_like(ageRelativeAdjustment)
    if rate >= maxOutRate:
        warnings.warn("covidTesting::adjustRatesByAge_KeepAverageRate Input rate {} > maxOutRate {}, returning input rates".format(rate, maxOutRate))
        return rate*np.ones_like(ageRelativeAdjustment)
    out = np.zeros_like(ageRelativeAdjustment)
    out[0] = maxOutRate+1 # just to start the while loop below
    while np.sum(out>=maxOutRate)>0:
        corrFactor = np.sum(agePopulationRatio/(1+ageRelativeAdjustment))
        out =  rate * (1+ageRelativeAdjustment) * corrFactor
        if np.sum(out>=maxOutRate)>0:
            warnings.warn("covidTesting::adjustRatesByAge_KeepAverageRate Adjusted rate larger than {} encountered, reducing ageAdjustment variance by 10%".format(maxOutRate))
            #print(out)
            tmp_mean = np.mean(ageRelativeAdjustment)
            ageRelativeAdjustment = tmp_mean + np.sqrt(0.9)*(ageRelativeAdjustment-tmp_mean)
    return out


# ## Getting infected

# In[7]:


# Getting infected
# ------------------
# We wish to calibrate overall infection rates to match
# - previous R0 estimates,
# - available age-attack-ratios,



# Age-dependent mixing affects state transition S -> I1 (data available eg Imperial #13 report)
# The mixing-related data is nowhere to be found!
# This is an Age x Age symmetric matrix, showing which groups mix with which other ones.
# Data from DOI: 10.1097/EDE.0000000000001047 via http://www.socialcontactdata.org/tools/ interactive tool in data folder
# This is assumed to be contacts per day (but may need to be time-rescaled)
ageSocialMixingBaseline = pd.read_csv('data/socialcontactdata_UK_Mossong2008_social_contact_matrix.csv', sep=',').iloc[:,1:].values
ageSocialMixingDistancing = pd.read_csv('data/socialcontactdata_UK_Mossong2008_social_contact_matrix_with_distancing.csv', sep=',').iloc[:,1:].values

# Symmetrise these matrices (not sure why they aren't symmetric)
ageSocialMixingBaseline = (ageSocialMixingBaseline+ageSocialMixingBaseline.T)/2.
ageSocialMixingDistancing = (ageSocialMixingDistancing+ageSocialMixingDistancing.T)/2.

# For simplicity, let's assume scenario of perfect isolation in state-issued home quarantine, see commented below for alternatives
ageSocialMixingIsolation = np.zeros_like(ageSocialMixingBaseline)
#isolationEffectComparedToDistancing = 3. # TODO - find better numbers for proper isolation mixing estimation!
#ageSocialMixingIsolation = ageSocialMixingBaseline/(isolationEffectComparedToDistancing * np.mean(ageSocialMixingBaseline/ageSocialMixingDistancing))



# For the S->I1 transition we also need a product mapping,
# as the AS->AI1 rate is variable and depend on all AI via social mixing (ages) and transmission rates (I stages)
# this vector is nI long only, calibrated together with other variables to reproduce overall R0
# These numbers should represent rate of transmission given contact [will be multiplied by social mixing matrices]
transmissionInfectionStage = np.array([0.001, 0.1, 0.6, 0.5]) # We vary this during model fitting



# In[8]:


# Getting Infected in the Hospital
# ---------------------------------------

# The general experience is that infections spread faster in a hospital environment,
# we capture this intuition with an age-independent but increased "social Mixing" amongst hospital patients and staff

# TODO - This requires further data-driven calibration!

# Capture S->I1 within hospital, given the number of total infected inside hospitals
elevatedMixingRatioInHospital = 3. # TODO - fact check this number, atm just set based on intuition
# Called "Nosocomial viral infection", some data: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5414085/
# HAP: hospital acquired pneumonia, apparently quite common
# more data: https://cmr.asm.org/content/14/3/528
# on covid-19: https://www.thelancet.com/journals/lanpub/article/PIIS2468-2667(20)30073-6/fulltext "Nosocomial infection risk among health-care workers and patients has been identified as a research gap to be prioritised in the next few months by WHO."
withinHospitalSocialMixing = elevatedMixingRatioInHospital * np.sum(np.dot(agePopulationRatio, ageSocialMixingBaseline))


# In[9]:


# Also add new infected from travelling based on time-within-simulation

# TODO - get real travel data to make these numbers more realistic. For now based on the following assumptions:
# - people's age distribution in travel is square of the usual age distribution
# - travel rates declined from a base rate as a sigmoid due to border closures, with given mean and slope
# - infection rates due to travel are modelled as a gamma pdf over time, with given peak value, loc, and scale parameter
def trFunc_travelInfectionRate_ageAdjusted(
    t, # Time (int, in days) within simulation

    travelMaxTime = 200,
    travelBaseRate = 5e-4, # How many people normally travel back to the country per day # TODO - get data
    travelDecline_mean = 15.,
    travelDecline_slope = 1.,

    travelInfection_peak = 1e-1,
    travelInfection_maxloc = 10.,
    travelInfection_shape = 2.,

    **kwargs
):

    tmpTime = np.arange(travelMaxTime)
    # nAge x T TODO get some realistic data on this
    travelAgeRateByTime = travelBaseRate * np.outer(agePopulationRatio, 1-expit((tmpTime-travelDecline_mean)/travelDecline_slope))


    # 1 x T TODO get some realistic data on this, maybe make it age weighted
    travelContractionRateByTime = stats.gamma.pdf(tmpTime, a=travelInfection_shape, loc=0., scale=travelInfection_maxloc/(travelInfection_shape-1))
    travelContractionRateByTime = travelInfection_peak*travelContractionRateByTime/np.max(travelContractionRateByTime)


    if t >= travelAgeRateByTime.shape[-1]:
        return np.zeros(travelAgeRateByTime.shape[0])
    else:
        return travelAgeRateByTime[:,int(t)] * travelContractionRateByTime[int(t)]


# In[10]:


# Overall new infections include within quarantine and hospital infections
# ------------------------------------------------------------------------

def trFunc_newInfections_Complete(
    stateTensor,
    policySocialDistancing, # True / False, no default because it's important to know which one we use at any moment!
    policyImmunityPassports, # True / False, no default because it's important to know which one we use at any moment!
    ageSocialMixingBaseline = ageSocialMixingBaseline,
    ageSocialMixingDistancing = ageSocialMixingDistancing,
    ageSocialMixingIsolation = ageSocialMixingIsolation,
    withinHospitalSocialMixing = withinHospitalSocialMixing,
    transmissionInfectionStage = transmissionInfectionStage,

    **kwargs
):
    """
    All new infections, given infected people on all different isolation states (normal, home, hospital)
    We use the following assumptions:

    - Infectiousness only depends on infection stage, not age or location

    - Hospitalised people are assumed to only mix with other hospitalised people (this includes staff!),
    in a non-age-dependent manner: withinHospitalSocialMixing

    If policySocialDistancing is True
    - Non- and home-isolated people mix with non- and home isolated via ageSocialMixingDistancing (averaging interactions)

    If policySocialDistancing is False, we assume home-isolation is taken more seriously, but with little effect on non-isolated people
    - Non-isolated people mix with each other via ageSocialMixingBaseline, and with home-isolated people via ageSocialMixingIsolation
    - Home-isolated people do not mix with each other

    This separation will help disentangle the effects of simply a blanket lessening of social distancing
    (keeping the policy True but with less effective ageSocialMixingDistancing matrix),
    vs case isolation (policy = False, but with serious ageSocialMixingIsolation)
    """

    ageIsoContractionRate = np.zeros((nAge, nIso, nTest))


    # Add non-hospital infections
    #--------------------------------

    curNonIsolatedSocialMixing = ageSocialMixingDistancing if policySocialDistancing else ageSocialMixingBaseline

    # Add baseline interactions only between non-isolated people
    for k1 in [0,3]:
        for k2 in [0,3]:
            ageIsoContractionRate[:,k1,:] += np.expand_dims(
                np.matmul(
                    curNonIsolatedSocialMixing,
                    np.einsum('ijl,j->i',
                        stateTensor[:,1:(nI+1),k2,:], transmissionInfectionStage) # all infected in non-isolation
                ),
                axis=1
            )

    if policyImmunityPassports:
        # If the immunity passports policy is on, everyone who tested antibody positive, can roam freely
        # Therefore replace the interactions between people with testingState = 2 with ageSocialMixingBaseline
        # we do this by using the distributive property of matrix multiplication, and adding extra interactions
        # "ageSocialMixingBaseline"-"curNonIsolatedSocialMixing" with each other (this is zero if no social distancing!)
        # TODO - this is a bit hacky?, but probably correct - double check though!
        for k1 in [0,3]:
            for k2 in [0,3]:
                ageIsoContractionRate[:,k1,2:] += np.matmul(
                        ageSocialMixingBaseline-curNonIsolatedSocialMixing,
                        np.einsum('ijk,j->ik',
                            stateTensor[:,1:(nI+1),k2,2:], transmissionInfectionStage) # all infected in non-isolation
                    )

    # Add isolation interactions only between isolated and non-isolated people
    # non-isolated contracting it from isolated
    for k1 in [0,3]:
        ageIsoContractionRate[:,k1,:] += np.expand_dims(
            np.matmul(
                ageSocialMixingIsolation,
                np.einsum('ijl,j->i',
                    stateTensor[:,1:(nI+1),1,:], transmissionInfectionStage) # all infected in isolation
            ),
            axis=1
        )

    # isolated contracting it from non-isolated
    for k1 in [0,3]:
        ageIsoContractionRate[:,1,:] += np.expand_dims(
            np.matmul(
                ageSocialMixingIsolation,
                np.einsum('ijl,j->i',
                    stateTensor[:,1:(nI+1),k1,:], transmissionInfectionStage) # all infected in non-hospital, non-isolation
            ),
            axis = 1
        )

        # isolated cannot contracting it from another isolated


    # Add in-hospital infections (of hospitalised patients, and staff)
    #--------------------------------
    # (TODO - within hospitals we probably want to take into effect the testing state;
    #      tested people are better isolated and there's less mixing)

    ageIsoContractionRate[:,2:,:] += np.expand_dims(
            withinHospitalSocialMixing *
            np.einsum('ijkl,j->i',
                stateTensor[:,1:(nI+1),2:,:], transmissionInfectionStage), # all infected in hospital (sick or working)
        axis = (1,2))


    return ageIsoContractionRate/np.sum(stateTensor) # Normalise the rate by total population



# ## Hospitalisation and hospital staff
#
# Disease progression in severe cases naturally leads to hospitalisation before death.
# One of the important policy questions we wish to estimate is how many people at any one time would require a hospital bed during the treatment of their disease.
#
# Hospitalisation is generally a simple situation modeling wise. People with either symptomatic infection (I3-...In states), or for other sicknesses (baseline hospitalisation) end up in hospital. People in S health state may return to non-hospitalised S state, however people in (informed, see later) I state generally remain in hospital until they recovered or dead.
#
# Home quarantine / social distancing is a different situation. Unlike other reports, here we do not (yet) wish to disentagle the effects of individual quarantine operations (school closures, working from home, social distancing), but rather investigate the effects of current full lockdown (coming into effect on 24 March in the UK), versus testing-based informed individual quarantining.
#
# Numerically:
#
# - People in home isolation change their social mixing patterns. The overall social mixing matrix between people in no isolation and home has been estimated via the http://www.socialcontactdata.org/tools/ software, see details in the data_cleaning notebook, this will determine the S->I transition overall.
#
# - People in hospitals (sick) dramatically reduce their contacts outside the hospital, but increase the chance of transmission within the hospitalised community. For the purpose of this simulation, hospital staff will also in effect be suspecitble to higher risk of infection due to "hospitalised" patients and they also keep their normal interaction.
#
# - Reported numbers regarding pressure on the health system will report both COVID-19 and non-COVID-19 patients
#

# In[11]:


# Getting Hospitalised
# ---------------------------------------

# Describe the transitions to-from hospitals
# Note that this implementation will assume that hospitalisation takes an extra day,
# due to the discrete nature of the simulation, might need to be re-thought.
# -> if simulation of a single day is done in multiple steps (first disease progression, then potential hospitalisation),
#.    then this problem is avoided. Can do the same with testing.

# Further we assume that hospitalisation does not change health state,
# but if happens in a non-S state, then it persists until R1 or D
# (this may need to be relaxed for early untested I states, where the hospitalisation is not COVID-related).

# Hospitalisation mainly depends on disease severity
# Baseline hospitalisation rate (Data from Scotland: https://www.isdscotland.org/Health-Topics/Hospital-Care/Publications/Acute-Hospital-Publication/data-summary/)
#hospitalisationRateBaseline = 261278./(91.*(5.425*10**6)) # hospitalisation / (period in days * population) -> frac of pop hospitalised per day
#hospitalisationRecoveryRateBaseline = 1./4.2 # inverse of mean length of stay in days

# Larger data driver approaches, with age distribution, see data_cleaning_R.ipynb for details
ageHospitalisationRateBaseline = pd.read_csv('data/clean_hosp-epis-stat-admi-summ-rep-2015-16-rep_table_6.csv', sep=',').iloc[:,-1].values
ageHospitalisationRecoveryRateBaseline = 1./pd.read_csv('data/clean_10641_LoS_age_provider_suppressed.csv', sep=',').iloc[:,-1].values

# Calculate initial hospitalisation (occupancy), that will be used to initialise the model
initBaselineHospitalOccupancyEquilibriumAgeRatio = ageHospitalisationRateBaseline/(ageHospitalisationRateBaseline+ageHospitalisationRecoveryRateBaseline)


# Take into account the NHS work-force in hospitals that for our purposes count as "hospitalised S" population,
# also unaffected by quarantine measures
ageNhsClinicalStaffPopulationRatio = pd.read_csv('data/clean_nhsclinicalstaff.csv', sep=',').iloc[:,-1].values

# Extra rate of hospitalisation due to COVID-19 infection stages
# TODO - find / estimate data on this (unfortunately true rates are hard to get due to many unknown cases)
# Symptom to hospitalisation is 5.76 days on average (Imperial #8)

infToHospitalExtra = np.array([1e-4, 1e-3, 2e-2, 1e-2])

# We do know at least how age affects these risks:

# For calculations see data_cleaning_py.ipynb, calculations from CHESS dataset as per 05 Apr
relativeAdmissionRisk_given_COVID_by_age = np.array([-0.94886625, -0.96332087, -0.86528671, -0.79828999, -0.61535305,
       -0.35214767,  0.12567034,  0.85809052,  3.55950368])

riskOfAEAttandance_by_age = np.array([0.41261361, 0.31560648, 0.3843979 , 0.30475704, 0.26659415,
       0.25203475, 0.24970244, 0.31549102, 0.65181376])





# Build the transition tensor from any non-hospitalised state to a hospitalised state
# (being in home quarantine is assumed to affect only the infection probability [below], not the hospitalisation probability)
# caseIsolationHospitalisationRateAdjustment = 1.

# This function takes as input the number of people in given age and health state, and in any non-hospitalised state
# and returns the number of people staying in the same age and health state,
# but now hospitalised (the rest of people remain in whatever state they were in)

def trFunc_HospitalAdmission(
    ageHospitalisationRateBaseline = ageHospitalisationRateBaseline,
    infToHospitalExtra = infToHospitalExtra,
    ageRelativeExtraAdmissionRiskToCovid = relativeAdmissionRisk_given_COVID_by_age * riskOfAEAttandance_by_age,

    **kwargs
    ):

    # This tensor will pointwise multiply an nAge x nHS slice of the stateTensor
    trTensor_HospitalAdmission = np.zeros((nAge, nHS))

    ageAdjusted_infToHospitalExtra = copy.deepcopy(np.repeat(infToHospitalExtra[np.newaxis],nAge,axis=0))
    for ii in range(ageAdjusted_infToHospitalExtra.shape[1]):
        # Adjust death rate by age dependent disease severity
        ageAdjusted_infToHospitalExtra[:,ii] = adjustRatesByAge_KeepAverageRate(
            infToHospitalExtra[ii],
            ageRelativeAdjustment=ageRelativeExtraAdmissionRiskToCovid
        )

    # Add baseline hospitalisation to all non-dead states
    trTensor_HospitalAdmission[:,:-1] += np.expand_dims(ageHospitalisationRateBaseline,-1)

    # Add COVID-caused hospitalisation to all infected states (TODO: This is summation of rates for independent processes, should be correct, but check)
    trTensor_HospitalAdmission[:,1:(nI+1)] += ageAdjusted_infToHospitalExtra

    return trTensor_HospitalAdmission

# Recovery rates (hospital discharge)
# ------------------------------------

# Higher-than-normal discharge rate for people who recovered (as they were likely to be in hospital mostly due to the virus)
# TODO - check with health experts if this is correct assumption; probably also depends on testing state

def trFunc_HospitalDischarge(
    ageHospitalisationRecoveryRateBaseline = ageHospitalisationRecoveryRateBaseline,
    dischargeDueToCovidRateMultiplier = 3.,

    **kwargs
    ):

    trTensor_HospitalDischarge = np.zeros((nAge, nHS))

    # Baseline discharges apply to all non-symptomatic patients (TODO: take into account testing state!)
    trTensor_HospitalDischarge[:, :3] += ageHospitalisationRecoveryRateBaseline[:,np.newaxis]

    # No discharges for COVID symptomatic people from the hospital until they recover
    # TODO - check with health experts if this is correct assumption; probably also depends on testing state
    trTensor_HospitalDischarge[:, 3:5] = 0.


    trTensor_HospitalDischarge[:, 5:7] = dischargeDueToCovidRateMultiplier * ageHospitalisationRecoveryRateBaseline[:,np.newaxis]

    return trTensor_HospitalDischarge




# TODO - think of how the latest changes (no prenatal care, no elective surgeries, etc) changed the default hospitalisation rate
#trTensor_HospitalAdmission[:,5]

# TODO!!! - adjust disease progression transitions so that it shifts direct death probabilities to hospitalised death probabilities


# ## Disease progression
#
#  - assumed to be strictly age and infection stage dependent distributions (progression rates), doesn't depend on other people
#  - distinct states represent progression, not necessarly time, but only forward progression is allowed, and the inverse of rates represent average number of days in progression
#  - there is a small chance of COVID death from every state, but we assume death is most often preceeded by hospitalisation
#  - there is a chance of recovery (and becoming immunised) from every state
#
# We wish to calibrate these disease progression probabilities to adhere to observed data / earlier models
# - serial interval distribution suggests time-to-transmission of Gamma(6.5 days, 0.62) MODEL [Imperial #13]
# Symptom progression (All params with relatively wide confidence intervals)
# - infect-to-symptom onset is assumed 5 days mean MODEL [AceMod, https://arxiv.org/pdf/2003.10218.pdf]
# - symptom-to-death is 16 days DATA_WEAK [Imperial #8]
# - symptom-to-discharge is 20.5 days DATA_WEAK [Imperial #8]
# - symptom-to-hospitalisation is 5.76 days DATA_WEAK [Imperial #8]
# - hospitalisation-to-recovery is 14.51 days DATA_WEAK [Imperial #8]
# all the above in Imperial #8 is largely age dependent. Raw data available in data/ImperialReport8_subset_international_cases_2020_03_11.csv
#

# In[12]:


# Based on England data (CHESS and NHS England)



# I want a way to keep this as the "average" disease progression, but modify it such that old people have less favorable outcomes (as observed)
# But correspondingly I want people at lower risk to have more favorable outcome on average

# For calculations see data_cleaning_py.ipynb, calculations from NHS England dataset as per 05 Apr
relativeDeathRisk_given_COVID_by_age = np.array([-0.99742186, -0.99728639, -0.98158438, -0.9830432 , -0.82983414,
       -0.84039294,  0.10768979,  0.38432409,  5.13754904])

#ageRelativeDiseaseSeverity = np.array([-0.8, -0.6, -0.3, -0.3, -0.1, 0.1, 0.35, 0.4, 0.5]) # FIXED (above) - this is a guess, find data and fix
#ageRelativeRecoverySpeed = np.array([0.2]*5+[-0.1, -0.2, -0.3, -0.5]) # TODO - this is a guess, find data and fix
ageRelativeRecoverySpeed = np.array([0.]*9) # For now we make it same for everyone, makes calculations easier

# For calculations see data_cleaning_py.ipynb, calculations from NHS England dataset as per 05 Apr
caseFatalityRatioHospital_given_COVID_by_age = np.array([0.00856164, 0.03768844, 0.02321319, 0.04282494, 0.07512237,
       0.12550367, 0.167096  , 0.37953452, 0.45757006])


def trFunc_diseaseProgression(
    # Basic parameters to adhere to
    nonsymptomatic_ratio = 0.86,

    # number of days between measurable events
    infect_to_symptoms = 5.,
    #symptom_to_death = 16.,
    symptom_to_recovery = 10., # 20.5, #unrealiticly long for old people
    symptom_to_hospitalisation = 5.76,
    hospitalisation_to_recovery = 14.51,
    IgG_formation = 15.,

    # Age related parameters
    # for now we'll assume that all hospitalised cases are known (overall 23% of hospitalised COVID patients die. 9% overall case fatality ratio)
    caseFatalityRatioHospital_given_COVID_by_age = caseFatalityRatioHospital_given_COVID_by_age,
    ageRelativeRecoverySpeed = ageRelativeRecoverySpeed,

    # Unknown rates to estimate
    nonsymp_to_recovery = 15.,
    inverse_IS1_IS2 = 4.,


    **kwargs
    ):
    # Now we have all the information to build the age-aware multistage SIR model transition matrix
    # The full transition tensor is a sparse map from the Age x HealthState x isolation state to HealthState,
        # and thus is a 4th order tensor itself, representing a linear mapping
        # from "number of people aged A in health state B and isolation state C to health state D.
    trTensor_diseaseProgression = np.zeros((nAge, nHS, nIso, nHS))


    # Use basic parameters to regularise inputs
    E_IS1 = 1./infect_to_symptoms
    # Numbers nonsymptomatic is assumed to be 86% -> E->IN / E-IS1 = 0.86/0.14
    E_IN = 0.86/0.14 * E_IS1

    # Nonsymptomatic recovery
    IN_R1 = 1./nonsymp_to_recovery

    IS1_IS2  = 1./inverse_IS1_IS2

    IS2_R1 = 1./(symptom_to_recovery-inverse_IS1_IS2)

    R1_R2 = 1./IgG_formation


    # Disease progression matrix # TODO - calibrate (together with transmissionInfectionStage)
    # rows: from-state, cols: to-state (non-symmetric!)
    # - this represent excess deaths only, doesn't contain baseline deaths!

    # Calculate all non-serious cases that do not end up in hospitals.
    # Note that we only have reliable death data from hospitals (NHS England), so we do not model people dieing outside hospitals
    diseaseProgBaseline = np.array([
    # to: E,   IN,   IS1,   IS2,    R1,   R2,   D
        [  0 , E_IN, E_IS1,    0,   0,     0,   0   ], # from E
        [  0,   0,     0,   0,    IN_R1,   0,   0   ], # from IN
        [  0 ,  0,     0, IS1_IS2,  0,     0,    0 ], # from IS1
        [  0 ,  0,     0,    0,  IS2_R1,   0,   0  ], # from IS2
        [  0 ,  0,     0,    0,    0,    R1_R2,  0   ], # from R1
        [  0 ,  0,     0,    0,    0,     0,   0   ], # from R2
        [  0 ,  0,     0,    0,    0,     0,   0   ] # from D
    ])

    ageAdjusted_diseaseProgBaseline = copy.deepcopy(np.repeat(diseaseProgBaseline[np.newaxis],nAge,axis=0))

    # Modify all death and R1 rates:
    for ii in range(ageAdjusted_diseaseProgBaseline.shape[1]):
        # Adjust death rate by age dependent disease severity
        ageAdjusted_diseaseProgBaseline[:,ii,-1] = adjustRatesByAge_KeepAverageRate(
            ageAdjusted_diseaseProgBaseline[0,ii,-1],
            ageRelativeAdjustment=relativeDeathRisk_given_COVID_by_age
        )

        # Adjust recovery rate by age dependent recovery speed
        ageAdjusted_diseaseProgBaseline[:,ii,-3] = adjustRatesByAge_KeepAverageRate(
            ageAdjusted_diseaseProgBaseline[0,ii,-3],
            ageRelativeAdjustment=ageRelativeRecoverySpeed,
            agePopulationRatio=agePopulationRatio
        )

    ageAdjusted_diseaseProgBaseline_Hospital = copy.deepcopy(ageAdjusted_diseaseProgBaseline)
    # Calculate hospitalisation based rates, for which we do have data. Hospitalisation can end up with deaths

    # Make sure that the ratio of recoveries in hospital honour the case fatality ratio appropriately
    # IS2 -> death
    ageAdjusted_diseaseProgBaseline_Hospital[:,3,-1] = (
        # IS2 -> recovery
        ageAdjusted_diseaseProgBaseline_Hospital[:,3,-3] * (
            # multiply by cfr / (1-cfr) to get correct rate towards death
            caseFatalityRatioHospital_given_COVID_by_age/(
                 1 -  caseFatalityRatioHospital_given_COVID_by_age)
        )
    )


    # TODO - time to death might be incorrect overall without an extra delay state, especially for young people

    # Non-hospitalised disease progression
    for i1 in [0,1,3]:
        trTensor_diseaseProgression[:,1:,i1,1:] = ageAdjusted_diseaseProgBaseline

    # hospitalised disease progression
    trTensor_diseaseProgression[:,1:,2,1:] = ageAdjusted_diseaseProgBaseline_Hospital


    return trTensor_diseaseProgression


# ## Testing
#
# In this section we describe multiple types of tests (PCR, antigen and antibody), and estimate their sensitivity and specificity in different health stages. These are thought to be the same for patients of all ages, and isolation states at this time.
#
# We then model the transitions to other testing states, which are largely policy-based.
#
# To model the current data (up to 03 April 2020):
# - only PCR tests have been done in the UK
# - PCR tests are thought to be carried out almost exclusively on symptomatic patients, to determine if their symptoms are caused by SARS-CoV2 or some other infection (this helps us determine the baseline ILI symptoms in practice, to predict true negative rates of the tests given the SARS-infected vs non-SARS-infected (but ILI symptom producing) populations).
#
# One aim of this complete model is to enable policy makers to make decisions now, based on predicted test availability in the future, therefore most testing-related concerns will be hypotheticals. That said, we aim to accurately model the tests' capabilities based on extensive literature research, and also aim to bring stable policy-level outcomes despite the actual numbers may be inaccurate.
#
# Two important questions answered by integrating this section into the epidemiology model above will be:
#
#     1. In what ratio we should produce antibody and antigen lateral flow immunoassay tests? They require the same production capabilities and reagents, there is a question ideally suited to the policy making level
#
#     2. At what level of testing capabilities (PCR, antigen and antibody) can the country lessen the complete lockdown, without risking lives or overburdening the NHS?
#
#
#
# API:
#
# - trFunc_testing(stateTensor, t, policyFunc, testSpecifications, trFunc_testCapacity):
#     - This is the main transition rate function, it returns transition rates from and to all testing states
#
# - policyFunc
#     - Returns a testing policy about what states are tested with how many of which test
#
# - testSpecifications
#     - Details the FPR/FNR of individual tests given the health state
#
# - trFunc_testCapacity(t)
#     - outputs how many tests are available at time t of the different test types modelled

# In[13]:


# Test parameters
# ---------------


# assumptions about practical (not theoretical, see discrapancy in PCR!) parameters of tests
# TODO - but particular data and references from lit (or estimates based on previous similar tests)

# TODO - MANUAL! - this function is VERY specific to current health state setup, and needs to be manually edited if number of health states change
def inpFunc_testSpecifications(
    PCR_FNR_I1_to_R2 = np.array([ 0.9,  0.4, 0.15, 0.35, 0.5, 0.8]),
    PCR_FPR = 0.01,
    antigen_FNR_I1_to_R2 = np.array([ 0.95, 0.6, 0.35, 0.45, 0.6, 0.9]),
    antigen_FPR = 0.1,
    antibody_FNR_I1_to_R2 = np.array([0.99, 0.85, 0.8, 0.65, 0.3, 0.05]),
    antibody_FPR_S_to_I4 =  np.array([0.05, 0.04, 0.03, 0.02, 0.01])
    ):


    testSpecifications = pd.DataFrame(
    columns=["Name"],#, "Infection stage"],#, "Sensitivity", "Specificity"],

    data = (
        ["PCR"] * nHS +
        ["Antigen"] * (nHS) +
        ["Antibody"] * (nHS))
    )

    testSpecifications['OutputTestState'] = [1]*nHS + [1]*nHS + [2]*nHS # what information state does a pos test transition you to.

    testSpecifications['TruePosHealthState'] = [np.arange(1,nI+1)]*nHS + [np.arange(1,nI+1)]*nHS + [np.arange(nI+1,nI+nR+1)]*nHS # what information state does a pos test transition you to.

    # In some health states some people are true negatives and some are true positives! (No, makes litte sense to use, just account for it in FPR? Only matters for test makers...)
    # testSpecifications['AmbiguousPosHealthState'] = [np.arange(nI+1, nI+nR+1)]*nHS + [np.arange(nI+1, nI+nR+1)]*nHS + [np.arange(1, nI+1)]*nHS # what information state does a pos test transition you to.

    testSpecifications['InputHealthState'] = list(np.tile(range(nHS),3))

    # These numbers below are "defaults" illustrating the concept, but are modified by the inputs!!!

    testSpecifications['FalseNegativeRate'] = [ # ratio of positive (infected / immune) people missed by the test
        # For each health stage:
        #  S -> I1 (asymp) -> I2 (mild symp) -> I3 (symp, sick) -> I4 (symp, less sick) -> R1 / R2 (IgM, IgG avail) -> D

        # PCR
            0.,   0.9,            0.4,           0.15,                0.35,              0.5, 0.8,   0.,

        # Antigen
            0.,   0.95,           0.6,           0.35,                0.45,              0.6, 0.9,   0.,

        # Antibody
            0.,   0.99,           0.85,          0.8,                 0.65,              0.3, 0.05,  0.
    ]


    testSpecifications.loc[1:6,'FalseNegativeRate'] = PCR_FNR_I1_to_R2
    testSpecifications.loc[9:14,'FalseNegativeRate'] = antigen_FNR_I1_to_R2
    testSpecifications.loc[17:22,'FalseNegativeRate'] = antibody_FNR_I1_to_R2



    testSpecifications['FalsePositiveRate'] = [ # ratio of negative (non-infected or not immune) people deemed positive by the test
        # PCR
        0.01, 0.,0.,0.,0., 0.01, 0.01, 0.,

        # Antigen
        0.1, 0.,0.,0.,0., 0.1, 0.1, 0.,

        # Antibody
        0.05, 0.04, 0.03, 0.02, 0.01, 0., 0., 0.
    ]

    testSpecifications.loc[0,'FalsePositiveRate'] = PCR_FPR
    testSpecifications.loc[5:6,'FalsePositiveRate'] = PCR_FPR
    testSpecifications.loc[8,'FalsePositiveRate'] = antigen_FPR
    testSpecifications.loc[13:14,'FalsePositiveRate'] = antigen_FPR
    testSpecifications.loc[16:20,'FalsePositiveRate'] = antibody_FPR_S_to_I4

    return testSpecifications


# In[14]:


inpFunc_testSpecifications()


# In[15]:


# TODO - think if we should introdce an "autopsy test" posthumously, categorising people as tested after death?
# How is this done, is there data on its sens/spec?

# Testing capacity
# ----------------

# Assumptions about the testing capacity available at day d of the simulation

# For PCR - we will model this (for now, for fitting we'll plug in real data!), as the sum of two sigmoids:
#   - initial stage of PHE ramping up its limited capacity (parameterised by total capacity, inflection day and slope of ramp-up)
#   - second stage of non-PHE labs joining in and ramping up capacity (this hasn't happened yet, but expected soon! same parameterisation)

# For the antigen / antibody tests we define a single sigmoidal capacity curve (starting later than PCR, but with potentially much higher total capacity)
# We further define a ratio between the production of the two, due to them requiring the same capabilities.


def trFunc_testCapacity(
    realTime, # time within simulation (day)

    # PCR capacity - initial
    testCapacity_pcr_phe_total = 1e4,
    testCapacity_pcr_phe_inflexday = pd.to_datetime("2020-03-25", format="%Y-%m-%d"),
    testCapacity_pcr_phe_inflexslope = 5.,

    # PCR capacity - increased
    testCapacity_pcr_country_total = 1e5,
    testCapacity_pcr_country_inflexday = pd.to_datetime("2020-04-25", format="%Y-%m-%d"),
    testCapacity_pcr_country_inflexslope = 10,

    # Antibody / antigen capacity
    testCapacity_antibody_country_firstday = pd.to_datetime("2020-04-25", format="%Y-%m-%d"),

    testCapacity_antibody_country_total = 5e6,
    testCapacity_antibody_country_inflexday = pd.to_datetime("2020-05-20", format="%Y-%m-%d"),
    testCapacity_antibody_country_inflexslope = 20,

    testCapacity_antigenratio_country = 0.7,

    **kwargs

):

    # Returns a dictionary with test names and number available at day "t"

    outPCR = (
        #phe phase
        testCapacity_pcr_phe_total * expit((realTime-testCapacity_pcr_phe_inflexday).days/testCapacity_pcr_phe_inflexslope)
        +
        #whole country phase
        testCapacity_pcr_country_total * expit((realTime-testCapacity_pcr_country_inflexday).days/testCapacity_pcr_country_inflexslope)
    )


    if realTime<testCapacity_antibody_country_firstday:
        outAntiTotal = 0.
    else:
        outAntiTotal = (
            testCapacity_antibody_country_total * expit((realTime-testCapacity_antibody_country_inflexday).days/testCapacity_antibody_country_inflexslope)
        )

    return {
        "PCR": outPCR,
        "Antigen": outAntiTotal*testCapacity_antigenratio_country,
        "Antibody": outAntiTotal*(1-testCapacity_antigenratio_country)
    }



# Real life data on test capacity and who got tested
# ---------------------------------------------------

df_CHESS = pd.read_csv("/mnt/efs/data/CHESS_Aggregate20200417.csv").drop(0)
df_CHESS.index = pd.to_datetime(df_CHESS["DateOfAdmission"].values,format="%d-%m-%Y")

# Ignore too old and too recent data points
df_CHESS = df_CHESS.sort_index().drop("DateOfAdmission", axis=1).query('20200309 <= index <= '+CONST_DATA_CUTOFF_DATE)

# Get number of tests per age group
df_CHESS_numTests = df_CHESS.loc[:,df_CHESS.columns.str.startswith("AllAdmittedPatientsTestedForCOVID19")]

# Change age groups to reflect our groupings
df_CHESS_numTests_regroup = pd.DataFrame(data = regroup_by_age(
    inp = df_CHESS_numTests.to_numpy().T,
    fromAgeSplits=np.concatenate([np.array([1,5,15,25]),np.arange(45,85+1,10)]),
    toAgeSplits=np.arange(10,80+1,10)
).T)

df_CHESS_numTests_regroup.index = df_CHESS_numTests.index

def inpFunc_testingDataCHESS_PCR(
    realTime,
    realTestData = df_CHESS_numTests_regroup,
    **kwargs
    ):

    def nearest(items, pivot):
        return min(items, key=lambda x: abs(x - pivot))


    return df_CHESS_numTests_regroup.loc[nearest(df_CHESS_numTests_regroup.index, pd.to_datetime(realTime, format="%Y-%m-%d"))]






# In[16]:


df_CHESS_numTests_regroup


# In[17]:


# Symptom parameters
# ------------------


# Estimating the baseline ILI-symptoms from earlier studies as well as the success rate of COVID-19 tests

# ILI rate estimate from 2018-19 PHE Surveillance of influenza and other respiratory viruses in the UK report:
# https://assets.publishing.service.gov.uk/government/uploads/system/uploads/attachment_data/file/839350/Surveillance_of_influenza_and_other_respiratory_viruses_in_the_UK_2018_to_2019-FINAL.pdf

# TODO - get actual seasonal symptom rate predictions (from 2020 non-SARS respiratory viruses, this data surely exists)
 # (daily rate estimate from Figure 8 of the report)


# Respiratory diagnosis on hospital admissions (not just ILI, all, TODO - get only ILI?)
# NHS Hosp episode statistics 2018-19, page 12 https://files.digital.nhs.uk/F2/E70669/hosp-epis-stat-admi-summ-rep-2018-19-rep.pdf
# In hospital: 1.1 million respiratory episodes out of 17.1 million total episodes

def f_symptoms_nonCOVID(
    realTime,
    symptomsIliRCGP = 15./100000., # Symptom rate in general non-hospitalised population
    symptomsRespInHospitalFAEs = 1.1/17.1, # Symptom rate in hospitalised population

    **kwargs):
    """
    This function defines the non-COVID ILI symptoms rate in the population at a given t time
    """


    # TODO, add extra data etc as input. For now:
    return (symptomsIliRCGP, symptomsRespInHospitalFAEs)



# In[18]:


# Distribute tests amongst (a given subset of) symptomatic people
def distTestsSymp(people, testsAvailable, noncovid_sympRatio, symp_HS = range(3,5), alreadyTestedRate = None):
    """
    distribute tests amongst symptomatic people
    people is nAge x nHS-1 x ... (excluding dead)
    """

    # Calculate noncovid, but symptomatic people
    peopleSymp = copy.deepcopy(people)
    peopleSymp[:, :min(symp_HS)] *= noncovid_sympRatio
    peopleSymp[:, max(symp_HS):] *= noncovid_sympRatio

    # Subtract already tested people
    if alreadyTestedRate is not None:
        peopleSymp -= people*alreadyTestedRate




    # Check if we already tested everyone with a different test
    if np.sum(peopleSymp)<1e-6:  # avoid numerical instabilities
        return (0.,0.)

    testedRatio = min(1., testsAvailable/np.sum(peopleSymp))


    return (
        # test rate
        testedRatio * (peopleSymp/(people+1e-6)), # avoid dividing by zero
        # tests used to achieve this
        testedRatio * np.sum(peopleSymp)
    )


# In[19]:


# Testing policies (how to distribute available tests)
# ----------------------------------------------------

# Estimate at any one time how many people are getting tested (with which tests) from which health states
def policyFunc_testing_symptomaticOnly(
    stateTensor,
    realTime,

    # Test types (names correspoding to testSpecifications)
    testTypes, # = ["PCR", "Antigen", "Antibody"],

    # Test Capacity (dict with names above and numbers available on day t)
    testsAvailable, # = trFunc_testCapacity(t)

    # OPTIONAL ARGUMENTS (may be different for different policy functions, should come with defaults!)
    antibody_testing_policy = "hospworker_then_random",
    # This has these values (for now), {"none", "hospworker_then_random", "virus_positive_only", "virus_positive_only_hospworker_first"}

    # Baseline symptoms
    f_symptoms_nonCOVID = f_symptoms_nonCOVID,

    distributeRemainingToRandom = True,
    return_testsAvailable_remaining = False,

    **kwargs
    ):
    """
    Returns a rate distribution of available test types over age, health and isolation states
    (although age assumed not to matter here)
    """

    # Output nAge x nHS x nIso x nTest x len(testTypes) tensor
    out_testRate = np.zeros(stateTensor.shape+(len(testTypes),))

    # Testing capacity is testsAvailable

    # Get sympom ratio. [0] - general, [1] - hospitalised
    cur_noncovid_sympRatio = f_symptoms_nonCOVID(realTime, **kwargs["f_symptoms_nonCOVID_params"])





    # PCR testing
    # -----------

    # Hospitalised people get priority over PCR tests
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,2,0], # hospitalised non-positive people, exclude tested and dead people
        testsAvailable = testsAvailable["PCR"],
        noncovid_sympRatio = cur_noncovid_sympRatio[1]
    )

    out_testRate[:,:-1,2,0, testTypes.index("PCR")] += testRate
    testsAvailable["PCR"] -= testsUsed

    # Prioritise hospital workers next:
    # TODO: check if we should do this? In UK policy there was a 15% max for hospital worker testing until ~2 April...
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,3,0],
        testsAvailable = testsAvailable["PCR"],
        noncovid_sympRatio= cur_noncovid_sympRatio[0]
    )

    out_testRate[:,:-1,3,0, testTypes.index("PCR")] += testRate
    testsAvailable["PCR"] -= testsUsed

    # Distribute PCRs left over the other populations
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:2,0],
        testsAvailable = testsAvailable["PCR"],
        noncovid_sympRatio= cur_noncovid_sympRatio[0]
    )

    out_testRate[:,:-1,:2,0, testTypes.index("PCR")] += testRate
    testsAvailable["PCR"] -= testsUsed

    if distributeRemainingToRandom:
        # Distribute PCRs left over the other populations
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:,0],
            testsAvailable = testsAvailable["PCR"],
            noncovid_sympRatio= 1.,
            alreadyTestedRate= out_testRate[:,:-1,:,0, testTypes.index("PCR")]
        )

        out_testRate[:,:-1,:,0, testTypes.index("PCR")] += testRate
        testsAvailable["PCR"] -= testsUsed


    # Antigen testing
    # ---------------

    # Hospitalised people get priority over PCR tests
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,2,0], # hospitalised non-positive people, exclude tested and dead people
        testsAvailable = testsAvailable["Antigen"],
        noncovid_sympRatio= cur_noncovid_sympRatio[1],
        alreadyTestedRate=out_testRate[:,:-1,2, 0, testTypes.index("PCR")]
    )

    out_testRate[:,:-1,2,0, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed

    # Prioritise hospital workers next:
    # TODO: check if we should do this? In UK policy there was a 15% max for hospital worker testing until ~2 April...
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,3,0],
        testsAvailable = testsAvailable["Antigen"],
        noncovid_sympRatio= cur_noncovid_sympRatio[0],
        alreadyTestedRate=out_testRate[:,:-1,3, 0, testTypes.index("PCR")]
    )

    out_testRate[:,:-1,3,0, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed

    # Distribute Antigen tests left over the other symptomatic people
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:2,0],
        testsAvailable = testsAvailable["Antigen"],
        noncovid_sympRatio= cur_noncovid_sympRatio[0],
        alreadyTestedRate=out_testRate[:,:-1,:2, 0, testTypes.index("PCR")]
    )

    out_testRate[:,:-1,:2,0, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed

    if distributeRemainingToRandom:
        # Distribute antigen tests left over the other non-symptmatic populations
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:,0],
            testsAvailable = testsAvailable["Antigen"],
            noncovid_sympRatio= 1.,
            alreadyTestedRate= out_testRate[:,:-1,:,0, :].sum(-1)
        )

        out_testRate[:,:-1,:,0, testTypes.index("Antigen")] += testRate
        testsAvailable["Antigen"] -= testsUsed


    # Antibody testing
    # ----------------

    if antibody_testing_policy == "hospworker_then_random":

        # For now: give to hospital workers first, not taking into account previous tests or symptoms
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,3,:2],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically workers get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,3,:2, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed

        # Afterwards let's just distribute randomly in the rest of the population
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:3,:2],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically people get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,:3,:2, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed

    if antibody_testing_policy == "virus_positive_only_hospworker_first":

        # For now: give to hospital workers first, not taking into account previous tests or symptoms
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,3,1],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically workers get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,3,1, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed

        # Afterwards let's just distribute randomly in the rest of the population
        # TODO: Maybe prioratise people who tested positive for the virus before???
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:3,1],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically people get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,:3,1, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed


    if antibody_testing_policy == "virus_positive_only":

        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:,1],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically people get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,:,1, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed

    if antibody_testing_policy == "none":
        out_testRate += 0.
        testsAvailable["Antibody"] -= 0.



    if return_testsAvailable_remaining:
        return out_testRate, testsAvailable

    return out_testRate


# In[20]:


# Define reTesting policy(s) (ie give tests to people in non-0 test states!)
def policyFunc_testing_massTesting_with_reTesting(
    stateTensor,
    realTime,

    # Test types (names correspoding to testSpecifications)
    testTypes, # = ["PCR", "Antigen", "Antibody"],

    # Test Capacity (dict with names above and numbers available on day t)
    testsAvailable, # = trFunc_testCapacity(t)

    # OPTIONAL ARGUMENTS (may be different for different policy functions, should come with defaults!)

    basic_policyFunc = policyFunc_testing_symptomaticOnly,
    # This basic policy will:
    # - do PCRs on symptomatic hospitalised people
    # - do PCRs on symptomatic hospital staff
    # - do PCRs on symptomatic non-hospitalised people
    # If PCRs run out at any stage, we use antigen tests with same priorisation

    # Afterwards given fractions of remaining antigen tests are distributed amongst people given these ratios and their earlier testing status:
    #retesting_antigen_viruspos_ratio = 0.1, # find virus false positives
    # UPDATE <- retesting viruspos is same ratio is normal testing, as long as they're not in quarantine already!
    retesting_antigen_immunepos_ratio = 0.05, # find immunity false positives
    # The rest of antigen tests are given out randomly

    # Antibody tests are used primarily on people who tested positive for the virus
    #  (set in basic_policyFunc!, use "virus_positive_only_hospworker_first"!)
    # Afterwards we can use the remaining on either random people (dangerous with many false positives!)
    # or for retesting people with already positive immune tests to make sure they're still immune,
    # controlled by this ratio:
    retesting_antibody_immunepos_ratio = 1.,

    #distributeRemainingToRandom = True, # TODO - otherwise stockpile for future, how?
    return_testsAvailable_remaining = False,

    **kwargs
    ):

    # Output nAge x nHS x nIso x nTest x len(testTypes) tensor
    out_testRate = np.zeros(stateTensor.shape+(len(testTypes),))


    # First distribute tests to symptomatic people as usual:

    # inpArgs change to not distributing tests randomly:
    basic_policyFunc_params_modified = copy.deepcopy(kwargs["basic_policyFunc_params"])
    basic_policyFunc_params_modified["distributeRemainingToRandom"] = False
    basic_policyFunc_params_modified["return_testsAvailable_remaining"] = True


    # Run the basic policy function with these modified parameters
    out_testRate, testsAvailable = basic_policyFunc(
        stateTensor,
        realTime = realTime,
        testTypes = testTypes,
        testsAvailable = testsAvailable,
        **basic_policyFunc_params_modified
    )


    # We assume PCRs tend to run out done on symptomatic people in 0 Test state, so no retesting via PCR.


    # Antigen testing
    # ---------------

    # Retesting immune positive people
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:,2:], # immune positive people
        testsAvailable = testsAvailable["Antigen"] * retesting_antigen_immunepos_ratio,
        noncovid_sympRatio= 1., # set to 1. for ignoring symptom vs non-symptom
    )

    out_testRate[:,:-1,:,2:, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed


    # Distribute antigen tests left over the other non-symptmatic populations
    # UPDATE <- here we use tests equally distributed among people with negative or positive previous virus tests,
    # as long as they are in non-quarantined state (isoState 0) # TODO - hospital worker testing???
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,0,:2], # non-quarantined virus positive people
        testsAvailable = testsAvailable["Antigen"],
        noncovid_sympRatio= 1.,
        alreadyTestedRate= out_testRate[:,:-1,0,:2, testTypes.index("Antigen")] + out_testRate[:,:-1,0,:2, testTypes.index("PCR")]
    )

    out_testRate[:,:-1,0,:2, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed


    # Antibody testing
    # -----------------
    # Retesting antibody positive people
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:,2:], # virus positive people
        testsAvailable = testsAvailable["Antibody"] * retesting_antibody_immunepos_ratio,
        noncovid_sympRatio= 1., # set to 1. for ignoring symptom vs non-symptom
    )


    # Afterwards let's just distribute randomly in the rest of the population
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:,:2],
        testsAvailable = testsAvailable["Antibody"],
        noncovid_sympRatio= 1., # basically people get antibody tested regardless of symptoms
        alreadyTestedRate= out_testRate[:,:-1,:,:2, testTypes.index("Antibody")]
    )

    out_testRate[:,:-1,:,:2, testTypes.index("Antibody")] += testRate
    testsAvailable["Antibody"] -= testsUsed


    if return_testsAvailable_remaining:
        return out_testRate, testsAvailable

    return out_testRate


# In[21]:


def trFunc_testing(
    stateTensor,
    t,
    realStartDate,
    #policyFunc = policyFunc_testing_symptomaticOnly,
    policyFunc = policyFunc_testing_massTesting_with_reTesting,
    inpFunc_testSpecifications = inpFunc_testSpecifications,
    trFunc_testCapacity = trFunc_testCapacity,
    inpFunc_realData_testCapacity = inpFunc_testingDataCHESS_PCR,
    **kwargs
    ):
    """
    Returns a tensor of rates transitioning to tested states
    """
    trTensor_testTransitions = np.zeros((nAge, nHS, nIso, nTest, nTest))


    testSpecifications = inpFunc_testSpecifications(**kwargs["inpFunc_testSpecifications_params"])

    testTypes = list(set(testSpecifications["Name"]))


    # Check if we have real data on the administered tests

    # Add the current data on within-hospital PCRs carried out already
    curDate = pd.to_datetime(realStartDate, format="%Y-%m-%d") + pd.to_timedelta(int(t), unit="D")
    realData_closest = inpFunc_realData_testCapacity(realTime = curDate, **kwargs["inpFunc_realData_testCapacity_params"])

    if realData_closest.name == curDate: # We do have data, just fill it in
        testsAdministeredRate = np.zeros(stateTensor.shape+(len(testTypes),))



        # TODO - fix this very hacky solution accessing symptomatic ratio as a subfunc of the policy func
        noncovid_sympRatio = kwargs["policyFunc_params"]["basic_policyFunc_params"]["f_symptoms_nonCOVID"](curDate, **kwargs["policyFunc_params"]["basic_policyFunc_params"]["f_symptoms_nonCOVID_params"])
        noncovid_sympRatio = noncovid_sympRatio[1] # Use hospitalised patient symptom ratio
        symptomaticRatePerDiseaseState = np.array([noncovid_sympRatio]*stateTensor.shape[1])
        symptomaticRatePerDiseaseState[3:-(nR+1)] = 1. # set the symptomatic ratio of symptomatic states to 1
        symptomaticPeoplePerDiseaseStateInHospital = stateTensor[:,:-1, 2, 0] * np.expand_dims(symptomaticRatePerDiseaseState[:-1], axis=0)

        testsAdministeredRate[:,:-1, 2, 0, testTypes.index("PCR")] += (
            np.expand_dims(realData_closest.to_numpy(),1) # true number of tests on given day per age group
            *
            (symptomaticPeoplePerDiseaseStateInHospital / np.sum(symptomaticPeoplePerDiseaseStateInHospital, axis=-1, keepdims=True))
            # Calculate in what ratio we distribute the tests to people along disease states based on symptomatic (age is given in data!)
        )/(stateTensor[:,:-1, 2, 0]+1e-10) # Divide by total people in each state to get testing rate



    else: # we don't have data, follow our assumed availability and policy curves

        # policyFunc returns stateTensor x testTypes tensor of test administration rates
        testsAdministeredRate = policyFunc(
            stateTensor,
            realTime = curDate,
            testTypes = testTypes,
            testsAvailable = trFunc_testCapacity(realTime = curDate, **kwargs["trFunc_testCapacity_params"]),
            **kwargs["policyFunc_params"]
        )



    # Compute the transition ratio to tested states, given the administered tests

    for testType in testTypes:
        # Get the appropriate slices from testsAdmin. and testSpecs
        curTestSpecs = testSpecifications[testSpecifications["Name"]==testType]

        for curTS in range(nTest):
            # Set output positive test state based on current test state
            if curTS == int(curTestSpecs["OutputTestState"].values[0]):
                # already positive for the given test
                outTS_pos = curTS
            elif curTS == 3:
                # If already positive for both, stay positive
                outTS_pos = 3
            else:
                # Transition 0->1, 0->2, 1->2, 1->3 or 2->3
                outTS_pos = curTS + int(curTestSpecs["OutputTestState"].values[0])


            # Where do we go after negative test based on where we are now?
            if curTS == 0:
                # Negatives stay negatives
                outTS_neg = 0
            elif curTS == 3:
                # go to only virus or antibody positive from both positive
                outTS_neg = 3-int(curTestSpecs["OutputTestState"].values[0])
            elif curTS == int(curTestSpecs["OutputTestState"].values[0]):
                # go to 0 if tested for the one you're positive for
                outTS_neg = 0
            else:
                # stay where you are if you test negative for the one you didnt have anyway
                outTS_neg = curTS



            # Get the transition rates based on current health states
            for curHS in range(nHS):
                # Add the true positives * (1-FNR)
                if curHS in curTestSpecs["TruePosHealthState"].values[0]:
                    trTensor_testTransitions[:,curHS,:, curTS, outTS_pos] += (
                        testsAdministeredRate[:,curHS,:,curTS,testTypes.index(testType)] *
                        (1-curTestSpecs[curTestSpecs["InputHealthState"] == curHS]["FalseNegativeRate"].values[0])
                    )

                else:
                # Add the false positives * FPR
                    trTensor_testTransitions[:,curHS,:,curTS, outTS_pos] += (
                        testsAdministeredRate[:,curHS,:,curTS,testTypes.index(testType)] *
                        curTestSpecs[curTestSpecs["InputHealthState"] == curHS]["FalsePositiveRate"].values[0]
                    )


                # Add the false negatives (FNR)
                if curHS in curTestSpecs["TruePosHealthState"].values[0]:
                    trTensor_testTransitions[:,curHS,:,curTS,outTS_neg] += (
                        testsAdministeredRate[:,curHS,:,curTS,testTypes.index(testType)] *
                        curTestSpecs[curTestSpecs["InputHealthState"] == curHS]["FalseNegativeRate"].values[0]
                    )

                else:
                # Add the true negatives * (1-FNR)
                    trTensor_testTransitions[:,curHS,:,curTS,outTS_neg] += (
                        testsAdministeredRate[:,curHS,:,curTS,testTypes.index(testType)] *
                        curTestSpecs[curTestSpecs["InputHealthState"] == curHS]["FalsePositiveRate"].values[0]
                    )



    return trTensor_testTransitions#, testsAdministeredRate


# ## Quarantine policies
#
# This section describes alternatives to the social distancing by full lockdown (that is implemented as a change in the socialMixing matrices).
#
# One alternative is case isolation, either by hospitalisation or by home isolation. We will assume that all non-symptomatic people who test positive are home isolated along with families for nDaysInIsolation days. Symptomatic people have a chance of being immediately hospitalised instead of sent into home isolation

# In[22]:


def trFunc_quarantine_caseIsolation(
    trTensor_complete,
    t,
    trTensor_testing, # This is used to establish who gets tests and how many of those end up positive.
    nDaysInHomeIsolation = 14.,
    timeToIsolation = 0.5, # (days) time from testing positive to actually getting isolated
    # On average this many people get hospitalised (compared to home isolation), but modulated by age (TODO: values > 1? clip for now..)
    symptomHospitalisedRate_ageAdjusted = np.clip(
        adjustRatesByAge_KeepAverageRate(0.3, ageRelativeAdjustment=relativeAdmissionRisk_given_COVID_by_age),
        0.,1.),
    symptomaticHealthStates = [3,4], # TODO - define this in global variable and just pass here!
    **kwargs
    ):
    """
    This function redistributes testing rates, so they dont only create a testing state update, but also an isolation state update
    """

    trTensor_quarantineRate = np.zeros(stateTensor.shape+(nIso,))

    trTensor_freshlyVirusPositiveRate_inIso0 = copy.deepcopy(trTensor_testing[:,:,0,:2,1])
    trTensor_freshlyBothPositiveRate_inIso0 = copy.deepcopy(trTensor_testing[:,:,0,2:,3])


    for curHS in range(stateTensor.shape[1]-1): # ignore dead
        if curHS in symptomaticHealthStates:
            # Send a fraction of people (normal) who are symptomatic and tested positive to hospital, based on their age
            trTensor_quarantineRate[:,curHS,0,:2,2] += (
                (1./timeToIsolation)*symptomHospitalisedRate_ageAdjusted[:,np.newaxis]
                *
                trTensor_freshlyVirusPositiveRate_inIso0[:,curHS]
            )
            trTensor_quarantineRate[:,curHS,0,2:,2] += (
                (1./timeToIsolation)*symptomHospitalisedRate_ageAdjusted[:,np.newaxis]
                *
                trTensor_freshlyBothPositiveRate_inIso0[:,curHS]
            )
            # The rest to home isolation
            trTensor_quarantineRate[:,curHS,0,:2,1] += (
                (1./timeToIsolation)*(1.-symptomHospitalisedRate_ageAdjusted[:,np.newaxis])
                *
                trTensor_freshlyVirusPositiveRate_inIso0[:,curHS]
            )
            trTensor_quarantineRate[:,curHS,0,2:,1] += (
                (1./timeToIsolation)*(1.-symptomHospitalisedRate_ageAdjusted[:,np.newaxis])
                *
                trTensor_freshlyBothPositiveRate_inIso0[:,curHS]
            )

        else:
            # Send all non-symptomatic (normal) who tested freshly positive to home isolation
            trTensor_quarantineRate[:,curHS,0,:2,1] += (
                1./timeToIsolation
                *
                trTensor_freshlyVirusPositiveRate_inIso0[:,curHS]
            )
            trTensor_quarantineRate[:,curHS,0,2:,1] += (
                1./timeToIsolation
                *
                trTensor_freshlyBothPositiveRate_inIso0[:,curHS]
            )

    # Release people from home isolation after isolation period
    trTensor_quarantineRate[:,:,1,:,0] = 1./nDaysInHomeIsolation

    # Hospitalised people are assumed to be released after recovery, with normal rates (TODO: think if this is correct)

    # TODO!!! - importantly, hospital workers are not being home isolated / hospitalised under this policy.
    # How to keep track of hospital workers who get hospitalised or home isolated themselves,
    # such that they get back to being hospital workers afterwards?
    # A simple (slightly incorrect) solution would be to just implement a non-specific "pull" from isoState=0 people to hospital workers to fill up the missing people?
    # But the rate of this pull would be impossible to compute and would still be incorrect. Gotta think more on this.


    # Update the whole tensor accordingly
    # Make a copy for safety:
    out_trTensor_complete = copy.deepcopy(trTensor_complete)

    # First remove all the iso 0->0, test 0,1->1, 2,3->3 transitions (as they're all either hospitalised or sent to home isolation)
    out_trTensor_complete[:,:,0,:2,:,0,1] = 0.
    out_trTensor_complete[:,:,0,2:,:,0,3] = 0.



    # Newly virus positive, newly home-isolated, diagonal in disease state transition
    np.einsum('ijkj->ijk',
        out_trTensor_complete[:,:,0,:2,:,1,1])[:] = trTensor_quarantineRate[:,:,0,:2,1]
    np.einsum('ijkj->ijk',
        out_trTensor_complete[:,:,0,2:,:,1,3])[:] = trTensor_quarantineRate[:,:,0,2:,1]

    # Newly virus positive, newly hospitalised, diagonal in disease state transition
    np.einsum('ijkj->ijk',
        out_trTensor_complete[:,:,0,:2,:,2,1])[:] = trTensor_quarantineRate[:,:,0,:2,2]
    np.einsum('ijkj->ijk',
        out_trTensor_complete[:,:,0,2:,:,2,3])[:] = trTensor_quarantineRate[:,:,0,2:,2]

    # Home isolated people are "let go" after nDaysInHomeIsolation, without changing disease or testing state
    # (TODO: represent multiple testing / needing negative tests to let go, etc - hard problem!)
    # (UPDATE: multiple testing have now been represented, but for now we'll still let go people based on fixed time rather than negative test, to save tests!)
    np.einsum('ijkjk->ijk',
        out_trTensor_complete[:,:,1,:,:,0,:])[:] = trTensor_quarantineRate[:,:,1,:,0]


    # Return the full updated tensor (so NOT += outside, but actually =)
    return out_trTensor_complete



# ## Full simulation function

# In[23]:


# Function that computes the right side of the non-lin model ODE
def dydt_Complete(t,
                  stateTensor_flattened, # Might be double the normal size (as first dimension) _withNewOnlyCopy, if debugReturnNewPerDay

                  realStartDate = pd.to_datetime("2020-02-20", format="%Y-%m-%d"),

                  # debug
                  debugTransition = False,
                  debugTimestep = False,
                  debugReturnNewPerDay = True, # Now implemented by default into state iteration

                  # Dimensions
                  nAge=nAge, nHS=nHS, nI=nI, nR=nR, nIso=nIso, nTest=nTest,

                  # Input functions and tensors
                  # ----------------------------

                  # Health state updates
                  trFunc_diseaseProgression = trFunc_diseaseProgression,
                  trFunc_newInfections = trFunc_newInfections_Complete,

                  # Initial incoming travel-based infections (before restrictions)
                  trFunc_travelInfectionRate_ageAdjusted = trFunc_travelInfectionRate_ageAdjusted,

                  # Hospitalisation and recovery
                  trFunc_HospitalAdmission = trFunc_HospitalAdmission,
                  trFunc_HospitalDischarge = trFunc_HospitalDischarge,

                  # Policy changes (on social distancing for now) (TODO - possibly make more changes)
                  tStartSocialDistancing = pd.to_datetime("2020-03-23", format="%Y-%m-%d"),
                  tStopSocialDistancing = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),

                  tStartImmunityPassports = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),
                  tStopImmunityPassports = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),

                  tStartQuarantineCaseIsolation = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),
                  tStopQuarantineCaseIsolation = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),
                  trFunc_quarantine = trFunc_quarantine_caseIsolation,

                  # Testing
                  trFunc_testing = trFunc_testing,
                  #policyFunc_testing = policyFunc_testing_symptomaticOnly,
                  #testSpecifications = testSpecifications,
                  #trFunc_testCapacity = trFunc_testCapacity,
                  #trFunc_testCapacity_param_testCapacity_antigenratio_country = 0.3

                  **kwargs

):

    if debugTimestep:
        print(t)

    # Initialise return
    if debugReturnNewPerDay: # the input has 2 copies of the state tensor, second copy being the cumulative incomings
        stateTensor = np.reshape(stateTensor_flattened, [2, nAge, nHS, nIso, nTest])[0]
    else:
        stateTensor = np.reshape(stateTensor_flattened, [nAge, nHS, nIso, nTest])

    dydt = np.zeros_like(stateTensor)

    # Initialise the full transition tensor
    trTensor_complete = np.zeros((nAge, nHS, nIso, nTest, nHS, nIso, nTest))


    # Disease condition updates
    # ---------------------------
    trTensor_diseaseProgression = trFunc_diseaseProgression(**kwargs["trFunc_diseaseProgression_params"])

    # Get disease condition updates with no isolation or test transition ("diagonal along those")
    for k1 in [0,1,2,3]:
        np.einsum('ijlml->ijlm',
            trTensor_complete[:,:,k1,:,:,k1,:])[:] += np.expand_dims(
                trTensor_diseaseProgression[:,:,k1,:]
                ,[2]) # all non-hospitalised disease progression is same


    # Compute new infections (0->1 in HS) with no isolation or test transition ("diagonal along those")
    cur_policySocialDistancing = (
                    t >= (tStartSocialDistancing - realStartDate).days
                )*(
                    t <   (tStopSocialDistancing - realStartDate).days
                )
    cur_policyImmunityPassports = (
                    t >= (tStartImmunityPassports - realStartDate).days
                )*(
                    t <   (tStopImmunityPassports - realStartDate).days
                )
    np.einsum('iklkl->ikl',
        trTensor_complete[:,0,:,:,1,:,:])[:] += (
            trFunc_newInfections(
                stateTensor,
                policySocialDistancing = cur_policySocialDistancing,
                policyImmunityPassports = cur_policyImmunityPassports,
                **kwargs["trFunc_newInfections_params"]
            ))

    # Also add new infected from travelling of healthy people, based on time-within-simulation (this is correct with all (0,0) states, as tested or isolated people dont travel)
    trTensor_complete[:,0,0,0,1,0,0] += trFunc_travelInfectionRate_ageAdjusted(t, **kwargs["trFunc_travelInfectionRate_ageAdjusted_params"])


    # Hospitalisation state updates
    # -----------------------

    # Hospitalisation and recovery rates
    # We assume for now that these only depend on age and disease progression, not on testing state
    # (TODO - update this given new policies)

    # The disease and testing states don't change due to hospitalisation.
    # Hospital staff is treated as already hospitalised from all aspects expect social mixing, should suffice for now
    # TODO - Could try to devise a scheme in which hospital staff gets hospitalised and some recoveries from hospitalised state go back to hospital staff.
    # TODO - same issue with hospital staff home isolating; that's probably more important question!
    for k1 in [0,1]:
         np.einsum('ijljl->ijl',
            trTensor_complete[:,:,k1,:,:,2,:])[:] += np.expand_dims(
             trFunc_HospitalAdmission(**kwargs["trFunc_HospitalAdmission_params"]),[2])

    # Add recovery from hospital rates
    # TODO - again here (for now) we assume all discharged people go back to "normal state" instead of home isolation, have to think more on this
    np.einsum('ijljl->ijl',
            trTensor_complete[:,:,2,:,:,0,:])[:] += np.expand_dims(
                 trFunc_HospitalDischarge(**kwargs["trFunc_HospitalDischarge_params"]),[2])





    # Testing state updates
    # ---------------------

    # trFunc_testing returns a stateTensor x testStates output
    #      after the policyFunc assigns tests that are evaluated according to testSpecifications

    # Diagonal (no transitions) in age, health state and isolation state
    # (for now, probably TODO: testing positive correlates with new hospitalisation!)
    trTensor_testing = trFunc_testing(
                                            stateTensor,
                                            t,
                                            realStartDate,
                                            **kwargs["trFunc_testing_params"]
                                        )

    np.einsum('ijkljkm->ijklm',
            trTensor_complete)[:] += trTensor_testing


    # Quarantine policy
    # ------------------

    # Check if policy is "on"
    if (
            t >= (tStartQuarantineCaseIsolation - realStartDate).days
        )*(
            t <   (tStopQuarantineCaseIsolation - realStartDate).days
        ):
        # New quarantining only happens to people who are transitioning already from untested to virus positive state
        # Therefore here we DO use non-diagonal transitions, and we
        #     redistribute the transtion rates given the testing (which was previously assumed not to create transition in isolation state)
        trTensor_complete = trFunc_quarantine(
                                                trTensor_complete,
                                                t,
                                                trTensor_testing,
                                                **kwargs["trFunc_quarantine_params"]
                                            )




    # Final corrections
    # -----------------



    # TODO: simulate aging and normal birth / death (not terribly important on these time scales, but should be quite simple)


    # Ensure that every "row" sums to 0 by adding to the diagonal (doesn't create new people out of nowhere)
    # Extract (writable) diagonal array and subtract the "row"-sums for each initial state
    np.einsum('ijkljkl->ijkl', trTensor_complete)[:] -= np.einsum('...jkl->...', trTensor_complete)


    # Compute the actual derivatives
    dydt = np.einsum('ijkl,ijklmnp->imnp', stateTensor, trTensor_complete) # contract the HS axis, keep age


    if debugReturnNewPerDay:
        """
            If this is true, instead of returning the real dydt,
            return only the positive "incoming" number of people to each state, so we can track "new cases"
            This needs some approximations, as follows:
                1. Take the normal transition tensor (with rates potentially > 0)
                2. From all states re-normalise the outgoing rates to sum at most to 1
                    (if they were less, keep it, if larger, then this represents
                    “in this day, all people will leave this state, in these ratios to these states”)
                3. Multiply only these outgoing rates with the current state
                    (so the result wont keep the same number of people as normal,
                    but only represent the “new incomings” for each state)
        """

        trTensor_complete_newOnly = copy.deepcopy(trTensor_complete)

        # TODO - Think - this is probably unnecessary actually, artifically reduces "new" rates?
#         # Devide each row by the absolute diagonal rate (that is the sum of the row), but only if its larger than 1
#         trTensor_complete_newOnly /= (
#             np.expand_dims(
#                 np.clip(np.abs(np.einsum('ijkljkl->ijkl', trTensor_complete_newOnly)), a_min=1., a_max=np.inf),
#                 axis=[4,5,6]
#             )
#         )

        # Set the diagonals to zero (no preservation, no outgoing, will end up being the incoming only)
        np.einsum('ijkljkl->ijkl', trTensor_complete_newOnly)[:] = 0.

        dydt_newOnly = np.einsum('ijkl,ijklmnp->imnp', stateTensor, trTensor_complete_newOnly)

        dydt = np.stack([dydt, dydt_newOnly], axis=0)


    if debugTransition:
        return np.reshape(dydt, -1), trTensor_complete

    return np.reshape(dydt, -1)



# ## Initialise and run the model

# In[24]:


# Initialise state
stateTensor_init = copy.deepcopy(stateTensor)

# Populate
stateTensor_init[:,0,0,0] = agePopulationTotal

# Move hospital staff to working in hospital
stateTensor_init[:,0,0,0] -= ageNhsClinicalStaffPopulationRatio * agePopulationTotal
stateTensor_init[:,0,3,0] += ageNhsClinicalStaffPopulationRatio * agePopulationTotal

# Move people to hospital according to baseline occupation (move only from normal people, not hospital staff!)
stateTensor_init[:,0,2,0] += initBaselineHospitalOccupancyEquilibriumAgeRatio * stateTensor_init[:,0,0,0]
stateTensor_init[:,0,0,0] -= initBaselineHospitalOccupancyEquilibriumAgeRatio * stateTensor_init[:,0,0,0]




# Infect some young adults/middle-aged people
# stateTensor_init[2:4,0,0,0] -= 1000.#/np.sum(agePopulationTotal)
# stateTensor_init[2:4,1,0,0] += 1000.#/np.sum(agePopulationTotal)

# BETTER! - People get infected by travel in early stages!


# In[25]:

#!/usr/bin/env python
# coding: utf-8

# # COVID-19 model for policy makers in the United Kingdom

#  <p align="center">
#  <img src="images/dynamicalModel.png" width="70%">
#  </p>

# We use an extended [SEIR model](https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#The_SEIR_model) to capture available information both about the disease progression, as well as how accessible various disease states are by testing. Being tested might cause a transition in the <span style="display: inline-block;background-color:#FFFCCC">Testing State</span>; the rate of such a transition depends both on the <span style="display: inline-block;background-color:#D1E2FF">Health State</span> as well the parameters of the test used.
#
# Due to age being both an extra risk factor as well as a potential way for decision makers to clearly apply policies to only parts of the total population, we include it directly in our model, and most transition rates are strongly influenced by <span style="display: inline-block;background-color:#FFD4D1">Age State</span>.
#
# Finally, the main policy making tool as well as conundrum nowadays is the implementation of quarantining and social distancing in order to keep hospitals and medical staff under tolerable levels of pressure. We represent <span style="display: inline-block;background-color:#C2EDC0">Isolation States</span> to investigate the effects of various interventions on policy targets relating to hospitalisation rates and economic freedom, while describing the different health outcomes via the <span style="display: inline-block;background-color:#D1E2FF">Health States</span>.
#

# ### Health states and Disease progression
#
# Susceptible people may be Exposed to the virus by mixing with other infected people (E,A,I<sub>1</sub> and I<sub>2</sub>).
#
# They may get through the infection Asymptomatic and eventually Recover, or become symptomatic and more Infectious, spreading the disease, and potentially Dying or Recovering.
#
# Recovered people develop more effective antibodies against the virus, and are considered immune<sup>*</sup>.
#
#
# | State | Description | Testing |
# | ----- | ----------- | -------- |
# | S | Susceptible | Negative |
# | E | Exposed | Very weakly virus positive
# | A | Asymptomatic | Weakly virus positive
# | I<sub>1</sub> | Symptomatic early | Strongly virus positive
# | I<sub>2</sub> | Symptomatic late | Medium virus positive <br>Weakly IgM antibody positive
# | R<sub>1</sub> | Recovered early | IgM antibody positive
# | R<sub>2</sub> | Recovered late | IgM/IgG antibody positive
# | D | COVID-related death | May be virus or antibody positive
#
#
# <sub><sup>*</sup>We plan to consider partial / short-term immunity, see further discussion in [Continued research](#header_contd).</sub>

# # Code
#
# The below consists of the implementation of the model described above, with extended comments including assumptions, sources of data and areas to improve upon

# In[1]:


# Use available data up until this day; cutoff is important due to more recent data being less complete.
CONST_DATA_CUTOFF_DATE = "20200414"


# # Packages and Helper functions
#
# To preserve the single-notebook formulation, we include all packages used as well as subfunctions here
#
# To skip to the start of model implementation, <a href="#modelImplementation">click here</a>!

# In[2]:


# Packages

# Basic packages
import numpy as np

from scipy import integrate, stats, spatial
from scipy.special import expit, binom

import pandas as pd
import xlrd # help read excel files directly from source into pandas

import copy
import warnings

# Building parameter/computation graph
import inspect
from collections import OrderedDict

# OS/filesystem tools
import time
from datetime import datetime
import random
import string
import os
import shutil
import sys
import cloudpickle

# Distributed computing tools
import dask
import distributed
from dask.distributed import Client
from dask.distributed import as_completed
import itertools


# In[3]:


# Regroup various age-group representations into our internal one, and vice versa
def regroup_by_age(
    inp, # first dimension is ages, others dont matter.
    fromAgeSplits, toAgeSplits, maxAge=100., maxAgeWeight = 5.):
    fromAgeSplits = np.concatenate([np.array([0]), fromAgeSplits, np.array([maxAge])]) # Add a zero at beginning for calculations
    toAgeSplits = np.concatenate([np.array([0]), toAgeSplits, np.array([maxAge])]) # Add inf at end for calculations
    def getOverlap(a, b):
        return max(0, min(a[1], b[1]) - max(a[0], b[0]))
    out = np.zeros((len(toAgeSplits)-1,)+inp.shape[1:])
    for from_ind in range(1, len(fromAgeSplits)):
        # Redistribute to the new bins by calculating how many years in from_ind-1:from_ind falls into each output bin
        cur_out_distribution = (
        [getOverlap(toAgeSplits[cur_to_ind-1:cur_to_ind+1],fromAgeSplits[from_ind-1:from_ind+1])  for cur_to_ind in range(1, len(toAgeSplits))]
        )

        if cur_out_distribution[-1] > 0:
            cur_out_distribution[-1] = maxAgeWeight # Define the relative number of ages if we have to distribute between second to last and last age groups

        cur_out_distribution = cur_out_distribution/np.sum(cur_out_distribution)

        for to_ind in range(len(out)):
            out[to_ind] += cur_out_distribution[to_ind] * inp[from_ind-1]

    return out


# PARAMETER DICTIONARIES AND TABLES
# -----------------------------------------------------------------------------------------

# Build the nested parameter/computation graph of a single function.
def build_paramDict(cur_func):
    """
    This function iterates through all inputs of a function,
    and saves the default argument names and values into a dictionary.

    If any of the default arguments are functions themselves, then recursively (depth-first) adds an extra field to
    the dictionary, named <funcName + "_params">, that contains its inputs and arguments.

    The output of this function can then be passed as a "kwargs" object to the highest level function,
    which will then pass the parameter values to the lower dictionary levels appropriately
    """

    paramDict = OrderedDict()

    allArgs = inspect.getfullargspec(cur_func)

    # Check if there are any default parameters, if no, just return empty dict
    if allArgs.defaults is None:
        return paramDict


    for argname, argval in zip(allArgs.args[-len(allArgs.defaults):], allArgs.defaults):
        # Save the default argument
        paramDict[argname] = argval
        # If the default argument is a function, inspect it for further

        if callable(argval):
            # print(argname)
            paramDict[argname+"_params"] = build_paramDict(argval)

    return paramDict


# Do a mapping between dictionary and parameter table row and vice versa (for convenient use)

# Flatten the dictionary into a table with a single row (but many column):
def paramDict_toTable(paramDict):
    paramTable = pd.DataFrame()
    def paramDictRecurseIter(cur_table, cur_dict, preString):
        # Iterate through the dictionary to find all keys not ending in "_params",
        # and add them to the table with name <preString + key>
        #
        # If the key doesn end in "_params", then append the key to preString, in call this function on the value (that is a dict)
        for key, value in cur_dict.items():
            if key.endswith("_params"):
                paramDictRecurseIter(cur_table, value, preString+key+"_")
            else:
                paramTable[preString+key] = [value]

        # For the rare case where we want to keep an empty dictionary, the above for cycle doesn't keep it
        if len(cur_dict)==0:
            paramTable[preString] = [OrderedDict()]

        return cur_table

    return paramDictRecurseIter(paramTable, paramDict, preString="")


def paramTable_toDict(paramTable, defaultDict=None):
    # enable to pass a default dict (if paramTable is incomplete), in which we'll just add / overwrite the values
    paramDict = defaultDict if defaultDict is not None else OrderedDict()
    def placeArgInDictRecurse(argName, argVal, cur_dict):
        # Find all "_params_" in the argName, and for each step more and more down in the dictionary
        strloc = argName.find("_params_")
        if strloc == -1:
            # We're at the correct level of dictionary
            cur_dict[argName] = argVal
            return cur_dict
        else:
            # step to the next level of dictionary
            nextKey = argName[:strloc+len("_params_")-1]
            nextArgName = argName[strloc+len("_params_"):]
            if not nextKey in cur_dict:
                cur_dict[nextKey] = OrderedDict()
            placeArgInDictRecurse(nextArgName, argVal, cur_dict[nextKey])

        return cur_dict

    for key in paramTable.columns:
        paramDict = placeArgInDictRecurse(key, paramTable.at[0,key], paramDict)

    return paramDict


# <a name="modelImplementation"></a>
# # Model implementation
#
# The model is governed by these two main tensors:
# - State tensor: a 4th order tensor containing axes:
#     - Age groups
#     - Health states
#     - Isolation states
#     - Testing states
#
#     In our extended state tensor formulation, we also keep track of not just people currently in each state, but also people newly arriving to each state, as a large number of available datasets refer to "new cases" or "new hospitalisations" each day, rather than current state occupancies normally represented by ODE solutions.
#
#
# - Dynamic state transition rate tensor
#     - Rates that govern all allowed transitions between states
#     - Dynamically recomputed every iteration, based on number of infected, their current social mixing and number and types of tests available, amongst other variables.
#     - Briefly:
#         - No transition between age groups
#         - No transitions between testing states without performing a test
#         - No transitions into S or out of D and R_IgG health states
#     - Allowed transitions are as showcased in the model image above
#     - Represented by a 7th order, sparse tensor, containing all transitions except age (unimportant due to relatively short time scales compared to coarse age grouping)
#

# ```
# NOTICE
# ------
# THE "MODEL IMPLEMENTATION" SECTION CONTAINS A LARGE NUMBER OF PARAMETER VALUES SET TO A DEFAULT VALUE.
# THESE ARE LARGELY INFORMED BY DATA, BUT NOT FIXED!
# THEY ARE VARIED DURING THE FITTING OF THE MODEL, ACCORDING TO HOW UNCERTAIN WE ARE IN THE PARAMETER
# ```
#
# The priors are defined <a href="#defineEnsemblePriors">below the model</a>. Although many of our uncertain / weak assumptions are signalled by "TODO" comments, we feel that the overall conclusions would not be affected by finding better parameter values, especially given our fitting and exploration approach.
#

# ### The state tensor

# In[4]:


# State Dimensions

# Health states (S, E and D are fixed to 1 dimension)
nI_symp = 2 # number of sympyomatic infected states
nI = 2+nI_symp # number of total infected states (disease stages), the +2 are Exposed and I_nonsymptomatic
nR = 2 # number of recovery states (antibody development post-disease, IgM and IgG are two stages)
nHS = 2+nI+nR # number of total health states, the +2: S, D are suspectible and dead

# Age groups (risk groups)
nAge = 9 # In accordance w Imperial #13 report (0-9, 10-19, ... 70-79, 80+)

# Isolation states
nIso = 4 # None/distancing, Case isolation, Hospitalised, Hospital staff

# Testing states
nTest = 4 # untested/negative, Virus positive, Antibody positive, Both positive


stateTensor = np.zeros((nAge, nHS, nIso, nTest))


# ### Transition rates (1 / day)
#
# The full transition rate structure is an 8th order tensor,
# mapping from any 4D state in the state tensor, to any other 4D state in the state tensor
#
# However, many of these transitions are non-sensical (ie a 3 year old cannot suddenly become 72, or a dead person be infected again), therefore during the construction of the full model below, we fill in the rates on all "allowed" transitions.
#
# We attempt to do so based on existing data either describing particular rates (like COVID-related hospitalisation),
# or data that helps constrain the product or ratios of multiple rates (such as the R0, or the case fatality ratio [noting this latter depends on testing policy and test availability]).
#
# Further, to massively reduce the number of weakly constrained parameters, we will approximate many of the weakly correlated transition rates as rank 1 (uncorrelated) matrices. For example the rate of hospitalisation for a patient at a given age and stage of infection will be computed as a product of two indepent rates, one based purely on the age (older people are generally more at risk of hospitalisation), and the other purely on how far the patient has progressed into the disease. This allows us to estimate more of required parameters from available published data.
#
# There of course still is a lot of uncertainty about how the virus behaves, and all of the data that we use is likely uncomplete and noisy. In order to better represent the things we do not know, we use advanced machine learning techniques, and investigate many possible scenarios (settings of parameters) and for all analyses we retain all plausible scenarios (various parameter settings that explain the available data well enough).
#
# Any policies we suggest for the near future are investigated for all plausible scenarios, such that decision makers know how likely each one will work as expected in these uncertain times. We further note that as we progress further into the pandemic, the number of plausible scenarios reduces more and more, enabling us to see the way out clearer and clearer.
#
#

# In[5]:


# Population (data from Imperial #13 ages.csv/UK)
agePopulationTotal = 1000.*np.array([8044.056,7642.473,8558.707,9295.024,8604.251,9173.465,7286.777,5830.635,3450.616])
#agePopulationTotal = 1000.*pd.read_csv("https://raw.githubusercontent.com/ImperialCollegeLondon/covid19model/master/data/ages.csv").iloc[3].values[2:]
# Currently: let's work with england population only instead of full UK, as NHS England + CHESS data is much clearer than other regions
agePopulationTotal *= 55.98/66.27 # (google england/uk population 2018, assuming age dist is similar)
agePopulationRatio = agePopulationTotal/np.sum(agePopulationTotal)

# Helper function to adjust average rates to age-aware rates
def adjustRatesByAge_KeepAverageRate(rate, ageRelativeAdjustment, agePopulationRatio=agePopulationRatio, maxOutRate=10):
    """This is a helper function and wont be picked up as a model parameter!"""
    if rate == 0:
        return np.zeros_like(ageRelativeAdjustment)
    if rate >= maxOutRate:
        warnings.warn("covidTesting::adjustRatesByAge_KeepAverageRate Input rate {} > maxOutRate {}, returning input rates".format(rate, maxOutRate))
        return rate*np.ones_like(ageRelativeAdjustment)
    out = np.zeros_like(ageRelativeAdjustment)
    out[0] = maxOutRate+1 # just to start the while loop below
    while np.sum(out>=maxOutRate)>0:
        corrFactor = np.sum(agePopulationRatio/(1+ageRelativeAdjustment))
        out =  rate * (1+ageRelativeAdjustment) * corrFactor
        if np.sum(out>=maxOutRate)>0:
            warnings.warn("covidTesting::adjustRatesByAge_KeepAverageRate Adjusted rate larger than {} encountered, reducing ageAdjustment variance by 10%".format(maxOutRate))
            #print(out)
            tmp_mean = np.mean(ageRelativeAdjustment)
            ageRelativeAdjustment = tmp_mean + np.sqrt(0.9)*(ageRelativeAdjustment-tmp_mean)
    return out


# ## Getting infected

# In[7]:


# Getting infected
# ------------------
# We wish to calibrate overall infection rates to match
# - previous R0 estimates,
# - available age-attack-ratios,



# Age-dependent mixing affects state transition S -> I1 (data available eg Imperial #13 report)
# The mixing-related data is nowhere to be found!
# This is an Age x Age symmetric matrix, showing which groups mix with which other ones.
# Data from DOI: 10.1097/EDE.0000000000001047 via http://www.socialcontactdata.org/tools/ interactive tool in data folder
# This is assumed to be contacts per day (but may need to be time-rescaled)
ageSocialMixingBaseline = pd.read_csv('data/socialcontactdata_UK_Mossong2008_social_contact_matrix.csv', sep=',').iloc[:,1:].values
ageSocialMixingDistancing = pd.read_csv('data/socialcontactdata_UK_Mossong2008_social_contact_matrix_with_distancing.csv', sep=',').iloc[:,1:].values

# Symmetrise these matrices (not sure why they aren't symmetric)
ageSocialMixingBaseline = (ageSocialMixingBaseline+ageSocialMixingBaseline.T)/2.
ageSocialMixingDistancing = (ageSocialMixingDistancing+ageSocialMixingDistancing.T)/2.

# For simplicity, let's assume scenario of perfect isolation in state-issued home quarantine, see commented below for alternatives
ageSocialMixingIsolation = np.zeros_like(ageSocialMixingBaseline)
#isolationEffectComparedToDistancing = 3. # TODO - find better numbers for proper isolation mixing estimation!
#ageSocialMixingIsolation = ageSocialMixingBaseline/(isolationEffectComparedToDistancing * np.mean(ageSocialMixingBaseline/ageSocialMixingDistancing))



# For the S->I1 transition we also need a product mapping,
# as the AS->AI1 rate is variable and depend on all AI via social mixing (ages) and transmission rates (I stages)
# this vector is nI long only, calibrated together with other variables to reproduce overall R0
# These numbers should represent rate of transmission given contact [will be multiplied by social mixing matrices]
transmissionInfectionStage = np.array([0.001, 0.1, 0.6, 0.5]) # We vary this during model fitting



# In[8]:


# Getting Infected in the Hospital
# ---------------------------------------

# The general experience is that infections spread faster in a hospital environment,
# we capture this intuition with an age-independent but increased "social Mixing" amongst hospital patients and staff

# TODO - This requires further data-driven calibration!

# Capture S->I1 within hospital, given the number of total infected inside hospitals
elevatedMixingRatioInHospital = 3. # TODO - fact check this number, atm just set based on intuition
# Called "Nosocomial viral infection", some data: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5414085/
# HAP: hospital acquired pneumonia, apparently quite common
# more data: https://cmr.asm.org/content/14/3/528
# on covid-19: https://www.thelancet.com/journals/lanpub/article/PIIS2468-2667(20)30073-6/fulltext "Nosocomial infection risk among health-care workers and patients has been identified as a research gap to be prioritised in the next few months by WHO."
withinHospitalSocialMixing = elevatedMixingRatioInHospital * np.sum(np.dot(agePopulationRatio, ageSocialMixingBaseline))


# In[9]:


# Also add new infected from travelling based on time-within-simulation

# TODO - get real travel data to make these numbers more realistic. For now based on the following assumptions:
# - people's age distribution in travel is square of the usual age distribution
# - travel rates declined from a base rate as a sigmoid due to border closures, with given mean and slope
# - infection rates due to travel are modelled as a gamma pdf over time, with given peak value, loc, and scale parameter
def trFunc_travelInfectionRate_ageAdjusted(
    t, # Time (int, in days) within simulation

    travelMaxTime = 200,
    travelBaseRate = 5e-4, # How many people normally travel back to the country per day # TODO - get data
    travelDecline_mean = 15.,
    travelDecline_slope = 1.,

    travelInfection_peak = 1e-1,
    travelInfection_maxloc = 10.,
    travelInfection_shape = 2.,

    **kwargs
):

    tmpTime = np.arange(travelMaxTime)
    # nAge x T TODO get some realistic data on this
    travelAgeRateByTime = travelBaseRate * np.outer(agePopulationRatio, 1-expit((tmpTime-travelDecline_mean)/travelDecline_slope))


    # 1 x T TODO get some realistic data on this, maybe make it age weighted
    travelContractionRateByTime = stats.gamma.pdf(tmpTime, a=travelInfection_shape, loc=0., scale=travelInfection_maxloc/(travelInfection_shape-1))
    travelContractionRateByTime = travelInfection_peak*travelContractionRateByTime/np.max(travelContractionRateByTime)


    if t >= travelAgeRateByTime.shape[-1]:
        return np.zeros(travelAgeRateByTime.shape[0])
    else:
        return travelAgeRateByTime[:,int(t)] * travelContractionRateByTime[int(t)]


# In[10]:


# Overall new infections include within quarantine and hospital infections
# ------------------------------------------------------------------------

def trFunc_newInfections_Complete(
    stateTensor,
    policySocialDistancing, # True / False, no default because it's important to know which one we use at any moment!
    policyImmunityPassports, # True / False, no default because it's important to know which one we use at any moment!
    ageSocialMixingBaseline = ageSocialMixingBaseline,
    ageSocialMixingDistancing = ageSocialMixingDistancing,
    ageSocialMixingIsolation = ageSocialMixingIsolation,
    withinHospitalSocialMixing = withinHospitalSocialMixing,
    transmissionInfectionStage = transmissionInfectionStage,

    **kwargs
):
    """
    All new infections, given infected people on all different isolation states (normal, home, hospital)
    We use the following assumptions:

    - Infectiousness only depends on infection stage, not age or location

    - Hospitalised people are assumed to only mix with other hospitalised people (this includes staff!),
    in a non-age-dependent manner: withinHospitalSocialMixing

    If policySocialDistancing is True
    - Non- and home-isolated people mix with non- and home isolated via ageSocialMixingDistancing (averaging interactions)

    If policySocialDistancing is False, we assume home-isolation is taken more seriously, but with little effect on non-isolated people
    - Non-isolated people mix with each other via ageSocialMixingBaseline, and with home-isolated people via ageSocialMixingIsolation
    - Home-isolated people do not mix with each other

    This separation will help disentangle the effects of simply a blanket lessening of social distancing
    (keeping the policy True but with less effective ageSocialMixingDistancing matrix),
    vs case isolation (policy = False, but with serious ageSocialMixingIsolation)
    """

    ageIsoContractionRate = np.zeros((nAge, nIso, nTest))


    # Add non-hospital infections
    #--------------------------------

    curNonIsolatedSocialMixing = ageSocialMixingDistancing if policySocialDistancing else ageSocialMixingBaseline

    # Add baseline interactions only between non-isolated people
    for k1 in [0,3]:
        for k2 in [0,3]:
            ageIsoContractionRate[:,k1,:] += np.expand_dims(
                np.matmul(
                    curNonIsolatedSocialMixing,
                    np.einsum('ijl,j->i',
                        stateTensor[:,1:(nI+1),k2,:], transmissionInfectionStage) # all infected in non-isolation
                ),
                axis=1
            )

    if policyImmunityPassports:
        # If the immunity passports policy is on, everyone who tested antibody positive, can roam freely
        # Therefore replace the interactions between people with testingState = 2 with ageSocialMixingBaseline
        # we do this by using the distributive property of matrix multiplication, and adding extra interactions
        # "ageSocialMixingBaseline"-"curNonIsolatedSocialMixing" with each other (this is zero if no social distancing!)
        # TODO - this is a bit hacky?, but probably correct - double check though!
        for k1 in [0,3]:
            for k2 in [0,3]:
                ageIsoContractionRate[:,k1,2:] += np.matmul(
                        ageSocialMixingBaseline-curNonIsolatedSocialMixing,
                        np.einsum('ijk,j->ik',
                            stateTensor[:,1:(nI+1),k2,2:], transmissionInfectionStage) # all infected in non-isolation
                    )

    # Add isolation interactions only between isolated and non-isolated people
    # non-isolated contracting it from isolated
    for k1 in [0,3]:
        ageIsoContractionRate[:,k1,:] += np.expand_dims(
            np.matmul(
                ageSocialMixingIsolation,
                np.einsum('ijl,j->i',
                    stateTensor[:,1:(nI+1),1,:], transmissionInfectionStage) # all infected in isolation
            ),
            axis=1
        )

    # isolated contracting it from non-isolated
    for k1 in [0,3]:
        ageIsoContractionRate[:,1,:] += np.expand_dims(
            np.matmul(
                ageSocialMixingIsolation,
                np.einsum('ijl,j->i',
                    stateTensor[:,1:(nI+1),k1,:], transmissionInfectionStage) # all infected in non-hospital, non-isolation
            ),
            axis = 1
        )

        # isolated cannot contracting it from another isolated


    # Add in-hospital infections (of hospitalised patients, and staff)
    #--------------------------------
    # (TODO - within hospitals we probably want to take into effect the testing state;
    #      tested people are better isolated and there's less mixing)

    ageIsoContractionRate[:,2:,:] += np.expand_dims(
            withinHospitalSocialMixing *
            np.einsum('ijkl,j->i',
                stateTensor[:,1:(nI+1),2:,:], transmissionInfectionStage), # all infected in hospital (sick or working)
        axis = (1,2))


    return ageIsoContractionRate/np.sum(stateTensor) # Normalise the rate by total population



# ## Hospitalisation and hospital staff
#
# Disease progression in severe cases naturally leads to hospitalisation before death.
# One of the important policy questions we wish to estimate is how many people at any one time would require a hospital bed during the treatment of their disease.
#
# Hospitalisation is generally a simple situation modeling wise. People with either symptomatic infection (I3-...In states), or for other sicknesses (baseline hospitalisation) end up in hospital. People in S health state may return to non-hospitalised S state, however people in (informed, see later) I state generally remain in hospital until they recovered or dead.
#
# Home quarantine / social distancing is a different situation. Unlike other reports, here we do not (yet) wish to disentagle the effects of individual quarantine operations (school closures, working from home, social distancing), but rather investigate the effects of current full lockdown (coming into effect on 24 March in the UK), versus testing-based informed individual quarantining.
#
# Numerically:
#
# - People in home isolation change their social mixing patterns. The overall social mixing matrix between people in no isolation and home has been estimated via the http://www.socialcontactdata.org/tools/ software, see details in the data_cleaning notebook, this will determine the S->I transition overall.
#
# - People in hospitals (sick) dramatically reduce their contacts outside the hospital, but increase the chance of transmission within the hospitalised community. For the purpose of this simulation, hospital staff will also in effect be suspecitble to higher risk of infection due to "hospitalised" patients and they also keep their normal interaction.
#
# - Reported numbers regarding pressure on the health system will report both COVID-19 and non-COVID-19 patients
#

# In[11]:


# Getting Hospitalised
# ---------------------------------------

# Describe the transitions to-from hospitals
# Note that this implementation will assume that hospitalisation takes an extra day,
# due to the discrete nature of the simulation, might need to be re-thought.
# -> if simulation of a single day is done in multiple steps (first disease progression, then potential hospitalisation),
#.    then this problem is avoided. Can do the same with testing.

# Further we assume that hospitalisation does not change health state,
# but if happens in a non-S state, then it persists until R1 or D
# (this may need to be relaxed for early untested I states, where the hospitalisation is not COVID-related).

# Hospitalisation mainly depends on disease severity
# Baseline hospitalisation rate (Data from Scotland: https://www.isdscotland.org/Health-Topics/Hospital-Care/Publications/Acute-Hospital-Publication/data-summary/)
#hospitalisationRateBaseline = 261278./(91.*(5.425*10**6)) # hospitalisation / (period in days * population) -> frac of pop hospitalised per day
#hospitalisationRecoveryRateBaseline = 1./4.2 # inverse of mean length of stay in days

# Larger data driver approaches, with age distribution, see data_cleaning_R.ipynb for details
ageHospitalisationRateBaseline = pd.read_csv('data/clean_hosp-epis-stat-admi-summ-rep-2015-16-rep_table_6.csv', sep=',').iloc[:,-1].values
ageHospitalisationRecoveryRateBaseline = 1./pd.read_csv('data/clean_10641_LoS_age_provider_suppressed.csv', sep=',').iloc[:,-1].values

# Calculate initial hospitalisation (occupancy), that will be used to initialise the model
initBaselineHospitalOccupancyEquilibriumAgeRatio = ageHospitalisationRateBaseline/(ageHospitalisationRateBaseline+ageHospitalisationRecoveryRateBaseline)


# Take into account the NHS work-force in hospitals that for our purposes count as "hospitalised S" population,
# also unaffected by quarantine measures
ageNhsClinicalStaffPopulationRatio = pd.read_csv('data/clean_nhsclinicalstaff.csv', sep=',').iloc[:,-1].values

# Extra rate of hospitalisation due to COVID-19 infection stages
# TODO - find / estimate data on this (unfortunately true rates are hard to get due to many unknown cases)
# Symptom to hospitalisation is 5.76 days on average (Imperial #8)

infToHospitalExtra = np.array([1e-4, 1e-3, 2e-2, 1e-2])

# We do know at least how age affects these risks:

# For calculations see data_cleaning_py.ipynb, calculations from CHESS dataset as per 05 Apr
relativeAdmissionRisk_given_COVID_by_age = np.array([-0.94886625, -0.96332087, -0.86528671, -0.79828999, -0.61535305,
       -0.35214767,  0.12567034,  0.85809052,  3.55950368])

riskOfAEAttandance_by_age = np.array([0.41261361, 0.31560648, 0.3843979 , 0.30475704, 0.26659415,
       0.25203475, 0.24970244, 0.31549102, 0.65181376])





# Build the transition tensor from any non-hospitalised state to a hospitalised state
# (being in home quarantine is assumed to affect only the infection probability [below], not the hospitalisation probability)
# caseIsolationHospitalisationRateAdjustment = 1.

# This function takes as input the number of people in given age and health state, and in any non-hospitalised state
# and returns the number of people staying in the same age and health state,
# but now hospitalised (the rest of people remain in whatever state they were in)

def trFunc_HospitalAdmission(
    ageHospitalisationRateBaseline = ageHospitalisationRateBaseline,
    infToHospitalExtra = infToHospitalExtra,
    ageRelativeExtraAdmissionRiskToCovid = relativeAdmissionRisk_given_COVID_by_age * riskOfAEAttandance_by_age,

    **kwargs
    ):

    # This tensor will pointwise multiply an nAge x nHS slice of the stateTensor
    trTensor_HospitalAdmission = np.zeros((nAge, nHS))

    ageAdjusted_infToHospitalExtra = copy.deepcopy(np.repeat(infToHospitalExtra[np.newaxis],nAge,axis=0))
    for ii in range(ageAdjusted_infToHospitalExtra.shape[1]):
        # Adjust death rate by age dependent disease severity
        ageAdjusted_infToHospitalExtra[:,ii] = adjustRatesByAge_KeepAverageRate(
            infToHospitalExtra[ii],
            ageRelativeAdjustment=ageRelativeExtraAdmissionRiskToCovid
        )

    # Add baseline hospitalisation to all non-dead states
    trTensor_HospitalAdmission[:,:-1] += np.expand_dims(ageHospitalisationRateBaseline,-1)

    # Add COVID-caused hospitalisation to all infected states (TODO: This is summation of rates for independent processes, should be correct, but check)
    trTensor_HospitalAdmission[:,1:(nI+1)] += ageAdjusted_infToHospitalExtra

    return trTensor_HospitalAdmission

# Recovery rates (hospital discharge)
# ------------------------------------

# Higher-than-normal discharge rate for people who recovered (as they were likely to be in hospital mostly due to the virus)
# TODO - check with health experts if this is correct assumption; probably also depends on testing state

def trFunc_HospitalDischarge(
    ageHospitalisationRecoveryRateBaseline = ageHospitalisationRecoveryRateBaseline,
    dischargeDueToCovidRateMultiplier = 3.,

    **kwargs
    ):

    trTensor_HospitalDischarge = np.zeros((nAge, nHS))

    # Baseline discharges apply to all non-symptomatic patients (TODO: take into account testing state!)
    trTensor_HospitalDischarge[:, :3] += ageHospitalisationRecoveryRateBaseline[:,np.newaxis]

    # No discharges for COVID symptomatic people from the hospital until they recover
    # TODO - check with health experts if this is correct assumption; probably also depends on testing state
    trTensor_HospitalDischarge[:, 3:5] = 0.


    trTensor_HospitalDischarge[:, 5:7] = dischargeDueToCovidRateMultiplier * ageHospitalisationRecoveryRateBaseline[:,np.newaxis]

    return trTensor_HospitalDischarge




# TODO - think of how the latest changes (no prenatal care, no elective surgeries, etc) changed the default hospitalisation rate
#trTensor_HospitalAdmission[:,5]

# TODO!!! - adjust disease progression transitions so that it shifts direct death probabilities to hospitalised death probabilities


# ## Disease progression
#
#  - assumed to be strictly age and infection stage dependent distributions (progression rates), doesn't depend on other people
#  - distinct states represent progression, not necessarly time, but only forward progression is allowed, and the inverse of rates represent average number of days in progression
#  - there is a small chance of COVID death from every state, but we assume death is most often preceeded by hospitalisation
#  - there is a chance of recovery (and becoming immunised) from every state
#
# We wish to calibrate these disease progression probabilities to adhere to observed data / earlier models
# - serial interval distribution suggests time-to-transmission of Gamma(6.5 days, 0.62) MODEL [Imperial #13]
# Symptom progression (All params with relatively wide confidence intervals)
# - infect-to-symptom onset is assumed 5 days mean MODEL [AceMod, https://arxiv.org/pdf/2003.10218.pdf]
# - symptom-to-death is 16 days DATA_WEAK [Imperial #8]
# - symptom-to-discharge is 20.5 days DATA_WEAK [Imperial #8]
# - symptom-to-hospitalisation is 5.76 days DATA_WEAK [Imperial #8]
# - hospitalisation-to-recovery is 14.51 days DATA_WEAK [Imperial #8]
# all the above in Imperial #8 is largely age dependent. Raw data available in data/ImperialReport8_subset_international_cases_2020_03_11.csv
#

# In[12]:


# Based on England data (CHESS and NHS England)



# I want a way to keep this as the "average" disease progression, but modify it such that old people have less favorable outcomes (as observed)
# But correspondingly I want people at lower risk to have more favorable outcome on average

# For calculations see data_cleaning_py.ipynb, calculations from NHS England dataset as per 05 Apr
relativeDeathRisk_given_COVID_by_age = np.array([-0.99742186, -0.99728639, -0.98158438, -0.9830432 , -0.82983414,
       -0.84039294,  0.10768979,  0.38432409,  5.13754904])

#ageRelativeDiseaseSeverity = np.array([-0.8, -0.6, -0.3, -0.3, -0.1, 0.1, 0.35, 0.4, 0.5]) # FIXED (above) - this is a guess, find data and fix
#ageRelativeRecoverySpeed = np.array([0.2]*5+[-0.1, -0.2, -0.3, -0.5]) # TODO - this is a guess, find data and fix
ageRelativeRecoverySpeed = np.array([0.]*9) # For now we make it same for everyone, makes calculations easier

# For calculations see data_cleaning_py.ipynb, calculations from NHS England dataset as per 05 Apr
caseFatalityRatioHospital_given_COVID_by_age = np.array([0.00856164, 0.03768844, 0.02321319, 0.04282494, 0.07512237,
       0.12550367, 0.167096  , 0.37953452, 0.45757006])


def trFunc_diseaseProgression(
    # Basic parameters to adhere to
    nonsymptomatic_ratio = 0.86,

    # number of days between measurable events
    infect_to_symptoms = 5.,
    #symptom_to_death = 16.,
    symptom_to_recovery = 10., # 20.5, #unrealiticly long for old people
    symptom_to_hospitalisation = 5.76,
    hospitalisation_to_recovery = 14.51,
    IgG_formation = 15.,

    # Age related parameters
    # for now we'll assume that all hospitalised cases are known (overall 23% of hospitalised COVID patients die. 9% overall case fatality ratio)
    caseFatalityRatioHospital_given_COVID_by_age = caseFatalityRatioHospital_given_COVID_by_age,
    ageRelativeRecoverySpeed = ageRelativeRecoverySpeed,

    # Unknown rates to estimate
    nonsymp_to_recovery = 15.,
    inverse_IS1_IS2 = 4.,


    **kwargs
    ):
    # Now we have all the information to build the age-aware multistage SIR model transition matrix
    # The full transition tensor is a sparse map from the Age x HealthState x isolation state to HealthState,
        # and thus is a 4th order tensor itself, representing a linear mapping
        # from "number of people aged A in health state B and isolation state C to health state D.
    trTensor_diseaseProgression = np.zeros((nAge, nHS, nIso, nHS))


    # Use basic parameters to regularise inputs
    E_IS1 = 1./infect_to_symptoms
    # Numbers nonsymptomatic is assumed to be 86% -> E->IN / E-IS1 = 0.86/0.14
    E_IN = 0.86/0.14 * E_IS1

    # Nonsymptomatic recovery
    IN_R1 = 1./nonsymp_to_recovery

    IS1_IS2  = 1./inverse_IS1_IS2

    IS2_R1 = 1./(symptom_to_recovery-inverse_IS1_IS2)

    R1_R2 = 1./IgG_formation


    # Disease progression matrix # TODO - calibrate (together with transmissionInfectionStage)
    # rows: from-state, cols: to-state (non-symmetric!)
    # - this represent excess deaths only, doesn't contain baseline deaths!

    # Calculate all non-serious cases that do not end up in hospitals.
    # Note that we only have reliable death data from hospitals (NHS England), so we do not model people dieing outside hospitals
    diseaseProgBaseline = np.array([
    # to: E,   IN,   IS1,   IS2,    R1,   R2,   D
        [  0 , E_IN, E_IS1,    0,   0,     0,   0   ], # from E
        [  0,   0,     0,   0,    IN_R1,   0,   0   ], # from IN
        [  0 ,  0,     0, IS1_IS2,  0,     0,    0 ], # from IS1
        [  0 ,  0,     0,    0,  IS2_R1,   0,   0  ], # from IS2
        [  0 ,  0,     0,    0,    0,    R1_R2,  0   ], # from R1
        [  0 ,  0,     0,    0,    0,     0,   0   ], # from R2
        [  0 ,  0,     0,    0,    0,     0,   0   ] # from D
    ])

    ageAdjusted_diseaseProgBaseline = copy.deepcopy(np.repeat(diseaseProgBaseline[np.newaxis],nAge,axis=0))

    # Modify all death and R1 rates:
    for ii in range(ageAdjusted_diseaseProgBaseline.shape[1]):
        # Adjust death rate by age dependent disease severity
        ageAdjusted_diseaseProgBaseline[:,ii,-1] = adjustRatesByAge_KeepAverageRate(
            ageAdjusted_diseaseProgBaseline[0,ii,-1],
            ageRelativeAdjustment=relativeDeathRisk_given_COVID_by_age
        )

        # Adjust recovery rate by age dependent recovery speed
        ageAdjusted_diseaseProgBaseline[:,ii,-3] = adjustRatesByAge_KeepAverageRate(
            ageAdjusted_diseaseProgBaseline[0,ii,-3],
            ageRelativeAdjustment=ageRelativeRecoverySpeed,
            agePopulationRatio=agePopulationRatio
        )

    ageAdjusted_diseaseProgBaseline_Hospital = copy.deepcopy(ageAdjusted_diseaseProgBaseline)
    # Calculate hospitalisation based rates, for which we do have data. Hospitalisation can end up with deaths

    # Make sure that the ratio of recoveries in hospital honour the case fatality ratio appropriately
    # IS2 -> death
    ageAdjusted_diseaseProgBaseline_Hospital[:,3,-1] = (
        # IS2 -> recovery
        ageAdjusted_diseaseProgBaseline_Hospital[:,3,-3] * (
            # multiply by cfr / (1-cfr) to get correct rate towards death
            caseFatalityRatioHospital_given_COVID_by_age/(
                 1 -  caseFatalityRatioHospital_given_COVID_by_age)
        )
    )


    # TODO - time to death might be incorrect overall without an extra delay state, especially for young people

    # Non-hospitalised disease progression
    for i1 in [0,1,3]:
        trTensor_diseaseProgression[:,1:,i1,1:] = ageAdjusted_diseaseProgBaseline

    # hospitalised disease progression
    trTensor_diseaseProgression[:,1:,2,1:] = ageAdjusted_diseaseProgBaseline_Hospital


    return trTensor_diseaseProgression


# ## Testing
#
# In this section we describe multiple types of tests (PCR, antigen and antibody), and estimate their sensitivity and specificity in different health stages. These are thought to be the same for patients of all ages, and isolation states at this time.
#
# We then model the transitions to other testing states, which are largely policy-based.
#
# To model the current data (up to 03 April 2020):
# - only PCR tests have been done in the UK
# - PCR tests are thought to be carried out almost exclusively on symptomatic patients, to determine if their symptoms are caused by SARS-CoV2 or some other infection (this helps us determine the baseline ILI symptoms in practice, to predict true negative rates of the tests given the SARS-infected vs non-SARS-infected (but ILI symptom producing) populations).
#
# One aim of this complete model is to enable policy makers to make decisions now, based on predicted test availability in the future, therefore most testing-related concerns will be hypotheticals. That said, we aim to accurately model the tests' capabilities based on extensive literature research, and also aim to bring stable policy-level outcomes despite the actual numbers may be inaccurate.
#
# Two important questions answered by integrating this section into the epidemiology model above will be:
#
#     1. In what ratio we should produce antibody and antigen lateral flow immunoassay tests? They require the same production capabilities and reagents, there is a question ideally suited to the policy making level
#
#     2. At what level of testing capabilities (PCR, antigen and antibody) can the country lessen the complete lockdown, without risking lives or overburdening the NHS?
#
#
#
# API:
#
# - trFunc_testing(stateTensor, t, policyFunc, testSpecifications, trFunc_testCapacity):
#     - This is the main transition rate function, it returns transition rates from and to all testing states
#
# - policyFunc
#     - Returns a testing policy about what states are tested with how many of which test
#
# - testSpecifications
#     - Details the FPR/FNR of individual tests given the health state
#
# - trFunc_testCapacity(t)
#     - outputs how many tests are available at time t of the different test types modelled

# In[13]:


# Test parameters
# ---------------


# assumptions about practical (not theoretical, see discrapancy in PCR!) parameters of tests
# TODO - but particular data and references from lit (or estimates based on previous similar tests)

# TODO - MANUAL! - this function is VERY specific to current health state setup, and needs to be manually edited if number of health states change
def inpFunc_testSpecifications(
    PCR_FNR_I1_to_R2 = np.array([ 0.9,  0.4, 0.15, 0.35, 0.5, 0.8]),
    PCR_FPR = 0.01,
    antigen_FNR_I1_to_R2 = np.array([ 0.95, 0.6, 0.35, 0.45, 0.6, 0.9]),
    antigen_FPR = 0.1,
    antibody_FNR_I1_to_R2 = np.array([0.99, 0.85, 0.8, 0.65, 0.3, 0.05]),
    antibody_FPR_S_to_I4 =  np.array([0.05, 0.04, 0.03, 0.02, 0.01])
    ):


    testSpecifications = pd.DataFrame(
    columns=["Name"],#, "Infection stage"],#, "Sensitivity", "Specificity"],

    data = (
        ["PCR"] * nHS +
        ["Antigen"] * (nHS) +
        ["Antibody"] * (nHS))
    )

    testSpecifications['OutputTestState'] = [1]*nHS + [1]*nHS + [2]*nHS # what information state does a pos test transition you to.

    testSpecifications['TruePosHealthState'] = [np.arange(1,nI+1)]*nHS + [np.arange(1,nI+1)]*nHS + [np.arange(nI+1,nI+nR+1)]*nHS # what information state does a pos test transition you to.

    # In some health states some people are true negatives and some are true positives! (No, makes litte sense to use, just account for it in FPR? Only matters for test makers...)
    # testSpecifications['AmbiguousPosHealthState'] = [np.arange(nI+1, nI+nR+1)]*nHS + [np.arange(nI+1, nI+nR+1)]*nHS + [np.arange(1, nI+1)]*nHS # what information state does a pos test transition you to.

    testSpecifications['InputHealthState'] = list(np.tile(range(nHS),3))

    # These numbers below are "defaults" illustrating the concept, but are modified by the inputs!!!

    testSpecifications['FalseNegativeRate'] = [ # ratio of positive (infected / immune) people missed by the test
        # For each health stage:
        #  S -> I1 (asymp) -> I2 (mild symp) -> I3 (symp, sick) -> I4 (symp, less sick) -> R1 / R2 (IgM, IgG avail) -> D

        # PCR
            0.,   0.9,            0.4,           0.15,                0.35,              0.5, 0.8,   0.,

        # Antigen
            0.,   0.95,           0.6,           0.35,                0.45,              0.6, 0.9,   0.,

        # Antibody
            0.,   0.99,           0.85,          0.8,                 0.65,              0.3, 0.05,  0.
    ]


    testSpecifications.loc[1:6,'FalseNegativeRate'] = PCR_FNR_I1_to_R2
    testSpecifications.loc[9:14,'FalseNegativeRate'] = antigen_FNR_I1_to_R2
    testSpecifications.loc[17:22,'FalseNegativeRate'] = antibody_FNR_I1_to_R2



    testSpecifications['FalsePositiveRate'] = [ # ratio of negative (non-infected or not immune) people deemed positive by the test
        # PCR
        0.01, 0.,0.,0.,0., 0.01, 0.01, 0.,

        # Antigen
        0.1, 0.,0.,0.,0., 0.1, 0.1, 0.,

        # Antibody
        0.05, 0.04, 0.03, 0.02, 0.01, 0., 0., 0.
    ]

    testSpecifications.loc[0,'FalsePositiveRate'] = PCR_FPR
    testSpecifications.loc[5:6,'FalsePositiveRate'] = PCR_FPR
    testSpecifications.loc[8,'FalsePositiveRate'] = antigen_FPR
    testSpecifications.loc[13:14,'FalsePositiveRate'] = antigen_FPR
    testSpecifications.loc[16:20,'FalsePositiveRate'] = antibody_FPR_S_to_I4

    return testSpecifications


# In[14]:


inpFunc_testSpecifications()


# In[15]:


# TODO - think if we should introdce an "autopsy test" posthumously, categorising people as tested after death?
# How is this done, is there data on its sens/spec?

# Testing capacity
# ----------------

# Assumptions about the testing capacity available at day d of the simulation

# For PCR - we will model this (for now, for fitting we'll plug in real data!), as the sum of two sigmoids:
#   - initial stage of PHE ramping up its limited capacity (parameterised by total capacity, inflection day and slope of ramp-up)
#   - second stage of non-PHE labs joining in and ramping up capacity (this hasn't happened yet, but expected soon! same parameterisation)

# For the antigen / antibody tests we define a single sigmoidal capacity curve (starting later than PCR, but with potentially much higher total capacity)
# We further define a ratio between the production of the two, due to them requiring the same capabilities.


def trFunc_testCapacity(
    realTime, # time within simulation (day)

    # PCR capacity - initial
    testCapacity_pcr_phe_total = 1e4,
    testCapacity_pcr_phe_inflexday = pd.to_datetime("2020-03-25", format="%Y-%m-%d"),
    testCapacity_pcr_phe_inflexslope = 5.,

    # PCR capacity - increased
    testCapacity_pcr_country_total = 1e5,
    testCapacity_pcr_country_inflexday = pd.to_datetime("2020-04-25", format="%Y-%m-%d"),
    testCapacity_pcr_country_inflexslope = 10,

    # Antibody / antigen capacity
    testCapacity_antibody_country_firstday = pd.to_datetime("2020-04-25", format="%Y-%m-%d"),

    testCapacity_antibody_country_total = 5e6,
    testCapacity_antibody_country_inflexday = pd.to_datetime("2020-05-20", format="%Y-%m-%d"),
    testCapacity_antibody_country_inflexslope = 20,

    testCapacity_antigenratio_country = 0.7,

    **kwargs

):

    # Returns a dictionary with test names and number available at day "t"

    outPCR = (
        #phe phase
        testCapacity_pcr_phe_total * expit((realTime-testCapacity_pcr_phe_inflexday).days/testCapacity_pcr_phe_inflexslope)
        +
        #whole country phase
        testCapacity_pcr_country_total * expit((realTime-testCapacity_pcr_country_inflexday).days/testCapacity_pcr_country_inflexslope)
    )


    if realTime<testCapacity_antibody_country_firstday:
        outAntiTotal = 0.
    else:
        outAntiTotal = (
            testCapacity_antibody_country_total * expit((realTime-testCapacity_antibody_country_inflexday).days/testCapacity_antibody_country_inflexslope)
        )

    return {
        "PCR": outPCR,
        "Antigen": outAntiTotal*testCapacity_antigenratio_country,
        "Antibody": outAntiTotal*(1-testCapacity_antigenratio_country)
    }



# Real life data on test capacity and who got tested
# ---------------------------------------------------

df_CHESS = pd.read_csv("/mnt/efs/data/CHESS_Aggregate20200417.csv").drop(0)
df_CHESS.index = pd.to_datetime(df_CHESS["DateOfAdmission"].values,format="%d-%m-%Y")

# Ignore too old and too recent data points
df_CHESS = df_CHESS.sort_index().drop("DateOfAdmission", axis=1).query('20200309 <= index <= '+CONST_DATA_CUTOFF_DATE)

# Get number of tests per age group
df_CHESS_numTests = df_CHESS.loc[:,df_CHESS.columns.str.startswith("AllAdmittedPatientsTestedForCOVID19")]

# Change age groups to reflect our groupings
df_CHESS_numTests_regroup = pd.DataFrame(data = regroup_by_age(
    inp = df_CHESS_numTests.to_numpy().T,
    fromAgeSplits=np.concatenate([np.array([1,5,15,25]),np.arange(45,85+1,10)]),
    toAgeSplits=np.arange(10,80+1,10)
).T)

df_CHESS_numTests_regroup.index = df_CHESS_numTests.index

def inpFunc_testingDataCHESS_PCR(
    realTime,
    realTestData = df_CHESS_numTests_regroup,
    **kwargs
    ):

    def nearest(items, pivot):
        return min(items, key=lambda x: abs(x - pivot))


    return df_CHESS_numTests_regroup.loc[nearest(df_CHESS_numTests_regroup.index, pd.to_datetime(realTime, format="%Y-%m-%d"))]






# In[16]:


df_CHESS_numTests_regroup


# In[17]:


# Symptom parameters
# ------------------


# Estimating the baseline ILI-symptoms from earlier studies as well as the success rate of COVID-19 tests

# ILI rate estimate from 2018-19 PHE Surveillance of influenza and other respiratory viruses in the UK report:
# https://assets.publishing.service.gov.uk/government/uploads/system/uploads/attachment_data/file/839350/Surveillance_of_influenza_and_other_respiratory_viruses_in_the_UK_2018_to_2019-FINAL.pdf

# TODO - get actual seasonal symptom rate predictions (from 2020 non-SARS respiratory viruses, this data surely exists)
 # (daily rate estimate from Figure 8 of the report)


# Respiratory diagnosis on hospital admissions (not just ILI, all, TODO - get only ILI?)
# NHS Hosp episode statistics 2018-19, page 12 https://files.digital.nhs.uk/F2/E70669/hosp-epis-stat-admi-summ-rep-2018-19-rep.pdf
# In hospital: 1.1 million respiratory episodes out of 17.1 million total episodes

def f_symptoms_nonCOVID(
    realTime,
    symptomsIliRCGP = 15./100000., # Symptom rate in general non-hospitalised population
    symptomsRespInHospitalFAEs = 1.1/17.1, # Symptom rate in hospitalised population

    **kwargs):
    """
    This function defines the non-COVID ILI symptoms rate in the population at a given t time
    """


    # TODO, add extra data etc as input. For now:
    return (symptomsIliRCGP, symptomsRespInHospitalFAEs)



# In[18]:


# Distribute tests amongst (a given subset of) symptomatic people
def distTestsSymp(people, testsAvailable, noncovid_sympRatio, symp_HS = range(3,5), alreadyTestedRate = None):
    """
    distribute tests amongst symptomatic people
    people is nAge x nHS-1 x ... (excluding dead)
    """

    # Calculate noncovid, but symptomatic people
    peopleSymp = copy.deepcopy(people)
    peopleSymp[:, :min(symp_HS)] *= noncovid_sympRatio
    peopleSymp[:, max(symp_HS):] *= noncovid_sympRatio

    # Subtract already tested people
    if alreadyTestedRate is not None:
        peopleSymp -= people*alreadyTestedRate




    # Check if we already tested everyone with a different test
    if np.sum(peopleSymp)<1e-6:  # avoid numerical instabilities
        return (0.,0.)

    testedRatio = min(1., testsAvailable/np.sum(peopleSymp))


    return (
        # test rate
        testedRatio * (peopleSymp/(people+1e-6)), # avoid dividing by zero
        # tests used to achieve this
        testedRatio * np.sum(peopleSymp)
    )


# In[19]:


# Testing policies (how to distribute available tests)
# ----------------------------------------------------

# Estimate at any one time how many people are getting tested (with which tests) from which health states
def policyFunc_testing_symptomaticOnly(
    stateTensor,
    realTime,

    # Test types (names correspoding to testSpecifications)
    testTypes, # = ["PCR", "Antigen", "Antibody"],

    # Test Capacity (dict with names above and numbers available on day t)
    testsAvailable, # = trFunc_testCapacity(t)

    # OPTIONAL ARGUMENTS (may be different for different policy functions, should come with defaults!)
    antibody_testing_policy = "hospworker_then_random",
    # This has these values (for now), {"none", "hospworker_then_random", "virus_positive_only", "virus_positive_only_hospworker_first"}

    # Baseline symptoms
    f_symptoms_nonCOVID = f_symptoms_nonCOVID,

    distributeRemainingToRandom = True,
    return_testsAvailable_remaining = False,

    **kwargs
    ):
    """
    Returns a rate distribution of available test types over age, health and isolation states
    (although age assumed not to matter here)
    """

    # Output nAge x nHS x nIso x nTest x len(testTypes) tensor
    out_testRate = np.zeros(stateTensor.shape+(len(testTypes),))

    # Testing capacity is testsAvailable

    # Get sympom ratio. [0] - general, [1] - hospitalised
    cur_noncovid_sympRatio = f_symptoms_nonCOVID(realTime, **kwargs["f_symptoms_nonCOVID_params"])





    # PCR testing
    # -----------

    # Hospitalised people get priority over PCR tests
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,2,0], # hospitalised non-positive people, exclude tested and dead people
        testsAvailable = testsAvailable["PCR"],
        noncovid_sympRatio = cur_noncovid_sympRatio[1]
    )

    out_testRate[:,:-1,2,0, testTypes.index("PCR")] += testRate
    testsAvailable["PCR"] -= testsUsed

    # Prioritise hospital workers next:
    # TODO: check if we should do this? In UK policy there was a 15% max for hospital worker testing until ~2 April...
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,3,0],
        testsAvailable = testsAvailable["PCR"],
        noncovid_sympRatio= cur_noncovid_sympRatio[0]
    )

    out_testRate[:,:-1,3,0, testTypes.index("PCR")] += testRate
    testsAvailable["PCR"] -= testsUsed

    # Distribute PCRs left over the other populations
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:2,0],
        testsAvailable = testsAvailable["PCR"],
        noncovid_sympRatio= cur_noncovid_sympRatio[0]
    )

    out_testRate[:,:-1,:2,0, testTypes.index("PCR")] += testRate
    testsAvailable["PCR"] -= testsUsed

    if distributeRemainingToRandom:
        # Distribute PCRs left over the other populations
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:,0],
            testsAvailable = testsAvailable["PCR"],
            noncovid_sympRatio= 1.,
            alreadyTestedRate= out_testRate[:,:-1,:,0, testTypes.index("PCR")]
        )

        out_testRate[:,:-1,:,0, testTypes.index("PCR")] += testRate
        testsAvailable["PCR"] -= testsUsed


    # Antigen testing
    # ---------------

    # Hospitalised people get priority over PCR tests
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,2,0], # hospitalised non-positive people, exclude tested and dead people
        testsAvailable = testsAvailable["Antigen"],
        noncovid_sympRatio= cur_noncovid_sympRatio[1],
        alreadyTestedRate=out_testRate[:,:-1,2, 0, testTypes.index("PCR")]
    )

    out_testRate[:,:-1,2,0, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed

    # Prioritise hospital workers next:
    # TODO: check if we should do this? In UK policy there was a 15% max for hospital worker testing until ~2 April...
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,3,0],
        testsAvailable = testsAvailable["Antigen"],
        noncovid_sympRatio= cur_noncovid_sympRatio[0],
        alreadyTestedRate=out_testRate[:,:-1,3, 0, testTypes.index("PCR")]
    )

    out_testRate[:,:-1,3,0, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed

    # Distribute Antigen tests left over the other symptomatic people
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:2,0],
        testsAvailable = testsAvailable["Antigen"],
        noncovid_sympRatio= cur_noncovid_sympRatio[0],
        alreadyTestedRate=out_testRate[:,:-1,:2, 0, testTypes.index("PCR")]
    )

    out_testRate[:,:-1,:2,0, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed

    if distributeRemainingToRandom:
        # Distribute antigen tests left over the other non-symptmatic populations
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:,0],
            testsAvailable = testsAvailable["Antigen"],
            noncovid_sympRatio= 1.,
            alreadyTestedRate= out_testRate[:,:-1,:,0, :].sum(-1)
        )

        out_testRate[:,:-1,:,0, testTypes.index("Antigen")] += testRate
        testsAvailable["Antigen"] -= testsUsed


    # Antibody testing
    # ----------------

    if antibody_testing_policy == "hospworker_then_random":

        # For now: give to hospital workers first, not taking into account previous tests or symptoms
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,3,:2],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically workers get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,3,:2, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed

        # Afterwards let's just distribute randomly in the rest of the population
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:3,:2],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically people get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,:3,:2, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed

    if antibody_testing_policy == "virus_positive_only_hospworker_first":

        # For now: give to hospital workers first, not taking into account previous tests or symptoms
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,3,1],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically workers get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,3,1, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed

        # Afterwards let's just distribute randomly in the rest of the population
        # TODO: Maybe prioratise people who tested positive for the virus before???
        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:3,1],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically people get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,:3,1, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed


    if antibody_testing_policy == "virus_positive_only":

        testRate, testsUsed = distTestsSymp(
            people = stateTensor[:,:-1,:,1],
            testsAvailable = testsAvailable["Antibody"],
            noncovid_sympRatio= 1. # basically people get antibody tested regardless of symptoms
        )

        out_testRate[:,:-1,:,1, testTypes.index("Antibody")] += testRate
        testsAvailable["Antibody"] -= testsUsed

    if antibody_testing_policy == "none":
        out_testRate += 0.
        testsAvailable["Antibody"] -= 0.



    if return_testsAvailable_remaining:
        return out_testRate, testsAvailable

    return out_testRate


# In[20]:


# Define reTesting policy(s) (ie give tests to people in non-0 test states!)
def policyFunc_testing_massTesting_with_reTesting(
    stateTensor,
    realTime,

    # Test types (names correspoding to testSpecifications)
    testTypes, # = ["PCR", "Antigen", "Antibody"],

    # Test Capacity (dict with names above and numbers available on day t)
    testsAvailable, # = trFunc_testCapacity(t)

    # OPTIONAL ARGUMENTS (may be different for different policy functions, should come with defaults!)

    basic_policyFunc = policyFunc_testing_symptomaticOnly,
    # This basic policy will:
    # - do PCRs on symptomatic hospitalised people
    # - do PCRs on symptomatic hospital staff
    # - do PCRs on symptomatic non-hospitalised people
    # If PCRs run out at any stage, we use antigen tests with same priorisation

    # Afterwards given fractions of remaining antigen tests are distributed amongst people given these ratios and their earlier testing status:
    #retesting_antigen_viruspos_ratio = 0.1, # find virus false positives
    # UPDATE <- retesting viruspos is same ratio is normal testing, as long as they're not in quarantine already!
    retesting_antigen_immunepos_ratio = 0.05, # find immunity false positives
    # The rest of antigen tests are given out randomly

    # Antibody tests are used primarily on people who tested positive for the virus
    #  (set in basic_policyFunc!, use "virus_positive_only_hospworker_first"!)
    # Afterwards we can use the remaining on either random people (dangerous with many false positives!)
    # or for retesting people with already positive immune tests to make sure they're still immune,
    # controlled by this ratio:
    retesting_antibody_immunepos_ratio = 1.,

    #distributeRemainingToRandom = True, # TODO - otherwise stockpile for future, how?
    return_testsAvailable_remaining = False,

    **kwargs
    ):

    # Output nAge x nHS x nIso x nTest x len(testTypes) tensor
    out_testRate = np.zeros(stateTensor.shape+(len(testTypes),))


    # First distribute tests to symptomatic people as usual:

    # inpArgs change to not distributing tests randomly:
    basic_policyFunc_params_modified = copy.deepcopy(kwargs["basic_policyFunc_params"])
    basic_policyFunc_params_modified["distributeRemainingToRandom"] = False
    basic_policyFunc_params_modified["return_testsAvailable_remaining"] = True


    # Run the basic policy function with these modified parameters
    out_testRate, testsAvailable = basic_policyFunc(
        stateTensor,
        realTime = realTime,
        testTypes = testTypes,
        testsAvailable = testsAvailable,
        **basic_policyFunc_params_modified
    )


    # We assume PCRs tend to run out done on symptomatic people in 0 Test state, so no retesting via PCR.


    # Antigen testing
    # ---------------

    # Retesting immune positive people
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:,2:], # immune positive people
        testsAvailable = testsAvailable["Antigen"] * retesting_antigen_immunepos_ratio,
        noncovid_sympRatio= 1., # set to 1. for ignoring symptom vs non-symptom
    )

    out_testRate[:,:-1,:,2:, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed


    # Distribute antigen tests left over the other non-symptmatic populations
    # UPDATE <- here we use tests equally distributed among people with negative or positive previous virus tests,
    # as long as they are in non-quarantined state (isoState 0) # TODO - hospital worker testing???
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,0,:2], # non-quarantined virus positive people
        testsAvailable = testsAvailable["Antigen"],
        noncovid_sympRatio= 1.,
        alreadyTestedRate= out_testRate[:,:-1,0,:2, testTypes.index("Antigen")] + out_testRate[:,:-1,0,:2, testTypes.index("PCR")]
    )

    out_testRate[:,:-1,0,:2, testTypes.index("Antigen")] += testRate
    testsAvailable["Antigen"] -= testsUsed


    # Antibody testing
    # -----------------
    # Retesting antibody positive people
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:,2:], # virus positive people
        testsAvailable = testsAvailable["Antibody"] * retesting_antibody_immunepos_ratio,
        noncovid_sympRatio= 1., # set to 1. for ignoring symptom vs non-symptom
    )


    # Afterwards let's just distribute randomly in the rest of the population
    testRate, testsUsed = distTestsSymp(
        people = stateTensor[:,:-1,:,:2],
        testsAvailable = testsAvailable["Antibody"],
        noncovid_sympRatio= 1., # basically people get antibody tested regardless of symptoms
        alreadyTestedRate= out_testRate[:,:-1,:,:2, testTypes.index("Antibody")]
    )

    out_testRate[:,:-1,:,:2, testTypes.index("Antibody")] += testRate
    testsAvailable["Antibody"] -= testsUsed


    if return_testsAvailable_remaining:
        return out_testRate, testsAvailable

    return out_testRate


# In[21]:


def trFunc_testing(
    stateTensor,
    t,
    realStartDate,
    #policyFunc = policyFunc_testing_symptomaticOnly,
    policyFunc = policyFunc_testing_massTesting_with_reTesting,
    inpFunc_testSpecifications = inpFunc_testSpecifications,
    trFunc_testCapacity = trFunc_testCapacity,
    inpFunc_realData_testCapacity = inpFunc_testingDataCHESS_PCR,
    **kwargs
    ):
    """
    Returns a tensor of rates transitioning to tested states
    """
    trTensor_testTransitions = np.zeros((nAge, nHS, nIso, nTest, nTest))


    testSpecifications = inpFunc_testSpecifications(**kwargs["inpFunc_testSpecifications_params"])

    testTypes = list(set(testSpecifications["Name"]))


    # Check if we have real data on the administered tests

    # Add the current data on within-hospital PCRs carried out already
    curDate = pd.to_datetime(realStartDate, format="%Y-%m-%d") + pd.to_timedelta(int(t), unit="D")
    realData_closest = inpFunc_realData_testCapacity(realTime = curDate, **kwargs["inpFunc_realData_testCapacity_params"])

    if realData_closest.name == curDate: # We do have data, just fill it in
        testsAdministeredRate = np.zeros(stateTensor.shape+(len(testTypes),))



        # TODO - fix this very hacky solution accessing symptomatic ratio as a subfunc of the policy func
        noncovid_sympRatio = kwargs["policyFunc_params"]["basic_policyFunc_params"]["f_symptoms_nonCOVID"](curDate, **kwargs["policyFunc_params"]["basic_policyFunc_params"]["f_symptoms_nonCOVID_params"])
        noncovid_sympRatio = noncovid_sympRatio[1] # Use hospitalised patient symptom ratio
        symptomaticRatePerDiseaseState = np.array([noncovid_sympRatio]*stateTensor.shape[1])
        symptomaticRatePerDiseaseState[3:-(nR+1)] = 1. # set the symptomatic ratio of symptomatic states to 1
        symptomaticPeoplePerDiseaseStateInHospital = stateTensor[:,:-1, 2, 0] * np.expand_dims(symptomaticRatePerDiseaseState[:-1], axis=0)

        testsAdministeredRate[:,:-1, 2, 0, testTypes.index("PCR")] += (
            np.expand_dims(realData_closest.to_numpy(),1) # true number of tests on given day per age group
            *
            (symptomaticPeoplePerDiseaseStateInHospital / np.sum(symptomaticPeoplePerDiseaseStateInHospital, axis=-1, keepdims=True))
            # Calculate in what ratio we distribute the tests to people along disease states based on symptomatic (age is given in data!)
        )/(stateTensor[:,:-1, 2, 0]+1e-10) # Divide by total people in each state to get testing rate



    else: # we don't have data, follow our assumed availability and policy curves

        # policyFunc returns stateTensor x testTypes tensor of test administration rates
        testsAdministeredRate = policyFunc(
            stateTensor,
            realTime = curDate,
            testTypes = testTypes,
            testsAvailable = trFunc_testCapacity(realTime = curDate, **kwargs["trFunc_testCapacity_params"]),
            **kwargs["policyFunc_params"]
        )



    # Compute the transition ratio to tested states, given the administered tests

    for testType in testTypes:
        # Get the appropriate slices from testsAdmin. and testSpecs
        curTestSpecs = testSpecifications[testSpecifications["Name"]==testType]

        for curTS in range(nTest):
            # Set output positive test state based on current test state
            if curTS == int(curTestSpecs["OutputTestState"].values[0]):
                # already positive for the given test
                outTS_pos = curTS
            elif curTS == 3:
                # If already positive for both, stay positive
                outTS_pos = 3
            else:
                # Transition 0->1, 0->2, 1->2, 1->3 or 2->3
                outTS_pos = curTS + int(curTestSpecs["OutputTestState"].values[0])


            # Where do we go after negative test based on where we are now?
            if curTS == 0:
                # Negatives stay negatives
                outTS_neg = 0
            elif curTS == 3:
                # go to only virus or antibody positive from both positive
                outTS_neg = 3-int(curTestSpecs["OutputTestState"].values[0])
            elif curTS == int(curTestSpecs["OutputTestState"].values[0]):
                # go to 0 if tested for the one you're positive for
                outTS_neg = 0
            else:
                # stay where you are if you test negative for the one you didnt have anyway
                outTS_neg = curTS



            # Get the transition rates based on current health states
            for curHS in range(nHS):
                # Add the true positives * (1-FNR)
                if curHS in curTestSpecs["TruePosHealthState"].values[0]:
                    trTensor_testTransitions[:,curHS,:, curTS, outTS_pos] += (
                        testsAdministeredRate[:,curHS,:,curTS,testTypes.index(testType)] *
                        (1-curTestSpecs[curTestSpecs["InputHealthState"] == curHS]["FalseNegativeRate"].values[0])
                    )

                else:
                # Add the false positives * FPR
                    trTensor_testTransitions[:,curHS,:,curTS, outTS_pos] += (
                        testsAdministeredRate[:,curHS,:,curTS,testTypes.index(testType)] *
                        curTestSpecs[curTestSpecs["InputHealthState"] == curHS]["FalsePositiveRate"].values[0]
                    )


                # Add the false negatives (FNR)
                if curHS in curTestSpecs["TruePosHealthState"].values[0]:
                    trTensor_testTransitions[:,curHS,:,curTS,outTS_neg] += (
                        testsAdministeredRate[:,curHS,:,curTS,testTypes.index(testType)] *
                        curTestSpecs[curTestSpecs["InputHealthState"] == curHS]["FalseNegativeRate"].values[0]
                    )

                else:
                # Add the true negatives * (1-FNR)
                    trTensor_testTransitions[:,curHS,:,curTS,outTS_neg] += (
                        testsAdministeredRate[:,curHS,:,curTS,testTypes.index(testType)] *
                        curTestSpecs[curTestSpecs["InputHealthState"] == curHS]["FalsePositiveRate"].values[0]
                    )



    return trTensor_testTransitions#, testsAdministeredRate


# ## Quarantine policies
#
# This section describes alternatives to the social distancing by full lockdown (that is implemented as a change in the socialMixing matrices).
#
# One alternative is case isolation, either by hospitalisation or by home isolation. We will assume that all non-symptomatic people who test positive are home isolated along with families for nDaysInIsolation days. Symptomatic people have a chance of being immediately hospitalised instead of sent into home isolation

# In[22]:


def trFunc_quarantine_caseIsolation(
    trTensor_complete,
    t,
    trTensor_testing, # This is used to establish who gets tests and how many of those end up positive.
    nDaysInHomeIsolation = 14.,
    timeToIsolation = 0.5, # (days) time from testing positive to actually getting isolated
    # On average this many people get hospitalised (compared to home isolation), but modulated by age (TODO: values > 1? clip for now..)
    symptomHospitalisedRate_ageAdjusted = np.clip(
        adjustRatesByAge_KeepAverageRate(0.3, ageRelativeAdjustment=relativeAdmissionRisk_given_COVID_by_age),
        0.,1.),
    symptomaticHealthStates = [3,4], # TODO - define this in global variable and just pass here!
    **kwargs
    ):
    """
    This function redistributes testing rates, so they dont only create a testing state update, but also an isolation state update
    """

    trTensor_quarantineRate = np.zeros(stateTensor.shape+(nIso,))

    trTensor_freshlyVirusPositiveRate_inIso0 = copy.deepcopy(trTensor_testing[:,:,0,:2,1])
    trTensor_freshlyBothPositiveRate_inIso0 = copy.deepcopy(trTensor_testing[:,:,0,2:,3])


    for curHS in range(stateTensor.shape[1]-1): # ignore dead
        if curHS in symptomaticHealthStates:
            # Send a fraction of people (normal) who are symptomatic and tested positive to hospital, based on their age
            trTensor_quarantineRate[:,curHS,0,:2,2] += (
                (1./timeToIsolation)*symptomHospitalisedRate_ageAdjusted[:,np.newaxis]
                *
                trTensor_freshlyVirusPositiveRate_inIso0[:,curHS]
            )
            trTensor_quarantineRate[:,curHS,0,2:,2] += (
                (1./timeToIsolation)*symptomHospitalisedRate_ageAdjusted[:,np.newaxis]
                *
                trTensor_freshlyBothPositiveRate_inIso0[:,curHS]
            )
            # The rest to home isolation
            trTensor_quarantineRate[:,curHS,0,:2,1] += (
                (1./timeToIsolation)*(1.-symptomHospitalisedRate_ageAdjusted[:,np.newaxis])
                *
                trTensor_freshlyVirusPositiveRate_inIso0[:,curHS]
            )
            trTensor_quarantineRate[:,curHS,0,2:,1] += (
                (1./timeToIsolation)*(1.-symptomHospitalisedRate_ageAdjusted[:,np.newaxis])
                *
                trTensor_freshlyBothPositiveRate_inIso0[:,curHS]
            )

        else:
            # Send all non-symptomatic (normal) who tested freshly positive to home isolation
            trTensor_quarantineRate[:,curHS,0,:2,1] += (
                1./timeToIsolation
                *
                trTensor_freshlyVirusPositiveRate_inIso0[:,curHS]
            )
            trTensor_quarantineRate[:,curHS,0,2:,1] += (
                1./timeToIsolation
                *
                trTensor_freshlyBothPositiveRate_inIso0[:,curHS]
            )

    # Release people from home isolation after isolation period
    trTensor_quarantineRate[:,:,1,:,0] = 1./nDaysInHomeIsolation

    # Hospitalised people are assumed to be released after recovery, with normal rates (TODO: think if this is correct)

    # TODO!!! - importantly, hospital workers are not being home isolated / hospitalised under this policy.
    # How to keep track of hospital workers who get hospitalised or home isolated themselves,
    # such that they get back to being hospital workers afterwards?
    # A simple (slightly incorrect) solution would be to just implement a non-specific "pull" from isoState=0 people to hospital workers to fill up the missing people?
    # But the rate of this pull would be impossible to compute and would still be incorrect. Gotta think more on this.


    # Update the whole tensor accordingly
    # Make a copy for safety:
    out_trTensor_complete = copy.deepcopy(trTensor_complete)

    # First remove all the iso 0->0, test 0,1->1, 2,3->3 transitions (as they're all either hospitalised or sent to home isolation)
    out_trTensor_complete[:,:,0,:2,:,0,1] = 0.
    out_trTensor_complete[:,:,0,2:,:,0,3] = 0.



    # Newly virus positive, newly home-isolated, diagonal in disease state transition
    np.einsum('ijkj->ijk',
        out_trTensor_complete[:,:,0,:2,:,1,1])[:] = trTensor_quarantineRate[:,:,0,:2,1]
    np.einsum('ijkj->ijk',
        out_trTensor_complete[:,:,0,2:,:,1,3])[:] = trTensor_quarantineRate[:,:,0,2:,1]

    # Newly virus positive, newly hospitalised, diagonal in disease state transition
    np.einsum('ijkj->ijk',
        out_trTensor_complete[:,:,0,:2,:,2,1])[:] = trTensor_quarantineRate[:,:,0,:2,2]
    np.einsum('ijkj->ijk',
        out_trTensor_complete[:,:,0,2:,:,2,3])[:] = trTensor_quarantineRate[:,:,0,2:,2]

    # Home isolated people are "let go" after nDaysInHomeIsolation, without changing disease or testing state
    # (TODO: represent multiple testing / needing negative tests to let go, etc - hard problem!)
    # (UPDATE: multiple testing have now been represented, but for now we'll still let go people based on fixed time rather than negative test, to save tests!)
    np.einsum('ijkjk->ijk',
        out_trTensor_complete[:,:,1,:,:,0,:])[:] = trTensor_quarantineRate[:,:,1,:,0]


    # Return the full updated tensor (so NOT += outside, but actually =)
    return out_trTensor_complete



# ## Full simulation function

# In[23]:


# Function that computes the right side of the non-lin model ODE
def dydt_Complete(t,
                  stateTensor_flattened, # Might be double the normal size (as first dimension) _withNewOnlyCopy, if debugReturnNewPerDay

                  realStartDate = pd.to_datetime("2020-02-20", format="%Y-%m-%d"),

                  # debug
                  debugTransition = False,
                  debugTimestep = False,
                  debugReturnNewPerDay = True, # Now implemented by default into state iteration

                  # Dimensions
                  nAge=nAge, nHS=nHS, nI=nI, nR=nR, nIso=nIso, nTest=nTest,

                  # Input functions and tensors
                  # ----------------------------

                  # Health state updates
                  trFunc_diseaseProgression = trFunc_diseaseProgression,
                  trFunc_newInfections = trFunc_newInfections_Complete,

                  # Initial incoming travel-based infections (before restrictions)
                  trFunc_travelInfectionRate_ageAdjusted = trFunc_travelInfectionRate_ageAdjusted,

                  # Hospitalisation and recovery
                  trFunc_HospitalAdmission = trFunc_HospitalAdmission,
                  trFunc_HospitalDischarge = trFunc_HospitalDischarge,

                  # Policy changes (on social distancing for now) (TODO - possibly make more changes)
                  tStartSocialDistancing = pd.to_datetime("2020-03-23", format="%Y-%m-%d"),
                  tStopSocialDistancing = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),

                  tStartImmunityPassports = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),
                  tStopImmunityPassports = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),

                  tStartQuarantineCaseIsolation = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),
                  tStopQuarantineCaseIsolation = pd.to_datetime("2025-03-23", format="%Y-%m-%d"),
                  trFunc_quarantine = trFunc_quarantine_caseIsolation,

                  # Testing
                  trFunc_testing = trFunc_testing,
                  #policyFunc_testing = policyFunc_testing_symptomaticOnly,
                  #testSpecifications = testSpecifications,
                  #trFunc_testCapacity = trFunc_testCapacity,
                  #trFunc_testCapacity_param_testCapacity_antigenratio_country = 0.3

                  **kwargs

):

    if debugTimestep:
        print(t)

    # Initialise return
    if debugReturnNewPerDay: # the input has 2 copies of the state tensor, second copy being the cumulative incomings
        stateTensor = np.reshape(stateTensor_flattened, [2, nAge, nHS, nIso, nTest])[0]
    else:
        stateTensor = np.reshape(stateTensor_flattened, [nAge, nHS, nIso, nTest])

    dydt = np.zeros_like(stateTensor)

    # Initialise the full transition tensor
    trTensor_complete = np.zeros((nAge, nHS, nIso, nTest, nHS, nIso, nTest))


    # Disease condition updates
    # ---------------------------
    trTensor_diseaseProgression = trFunc_diseaseProgression(**kwargs["trFunc_diseaseProgression_params"])

    # Get disease condition updates with no isolation or test transition ("diagonal along those")
    for k1 in [0,1,2,3]:
        np.einsum('ijlml->ijlm',
            trTensor_complete[:,:,k1,:,:,k1,:])[:] += np.expand_dims(
                trTensor_diseaseProgression[:,:,k1,:]
                ,[2]) # all non-hospitalised disease progression is same


    # Compute new infections (0->1 in HS) with no isolation or test transition ("diagonal along those")
    cur_policySocialDistancing = (
                    t >= (tStartSocialDistancing - realStartDate).days
                )*(
                    t <   (tStopSocialDistancing - realStartDate).days
                )
    cur_policyImmunityPassports = (
                    t >= (tStartImmunityPassports - realStartDate).days
                )*(
                    t <   (tStopImmunityPassports - realStartDate).days
                )
    np.einsum('iklkl->ikl',
        trTensor_complete[:,0,:,:,1,:,:])[:] += (
            trFunc_newInfections(
                stateTensor,
                policySocialDistancing = cur_policySocialDistancing,
                policyImmunityPassports = cur_policyImmunityPassports,
                **kwargs["trFunc_newInfections_params"]
            ))

    # Also add new infected from travelling of healthy people, based on time-within-simulation (this is correct with all (0,0) states, as tested or isolated people dont travel)
    trTensor_complete[:,0,0,0,1,0,0] += trFunc_travelInfectionRate_ageAdjusted(t, **kwargs["trFunc_travelInfectionRate_ageAdjusted_params"])


    # Hospitalisation state updates
    # -----------------------

    # Hospitalisation and recovery rates
    # We assume for now that these only depend on age and disease progression, not on testing state
    # (TODO - update this given new policies)

    # The disease and testing states don't change due to hospitalisation.
    # Hospital staff is treated as already hospitalised from all aspects expect social mixing, should suffice for now
    # TODO - Could try to devise a scheme in which hospital staff gets hospitalised and some recoveries from hospitalised state go back to hospital staff.
    # TODO - same issue with hospital staff home isolating; that's probably more important question!
    for k1 in [0,1]:
         np.einsum('ijljl->ijl',
            trTensor_complete[:,:,k1,:,:,2,:])[:] += np.expand_dims(
             trFunc_HospitalAdmission(**kwargs["trFunc_HospitalAdmission_params"]),[2])

    # Add recovery from hospital rates
    # TODO - again here (for now) we assume all discharged people go back to "normal state" instead of home isolation, have to think more on this
    np.einsum('ijljl->ijl',
            trTensor_complete[:,:,2,:,:,0,:])[:] += np.expand_dims(
                 trFunc_HospitalDischarge(**kwargs["trFunc_HospitalDischarge_params"]),[2])





    # Testing state updates
    # ---------------------

    # trFunc_testing returns a stateTensor x testStates output
    #      after the policyFunc assigns tests that are evaluated according to testSpecifications

    # Diagonal (no transitions) in age, health state and isolation state
    # (for now, probably TODO: testing positive correlates with new hospitalisation!)
    trTensor_testing = trFunc_testing(
                                            stateTensor,
                                            t,
                                            realStartDate,
                                            **kwargs["trFunc_testing_params"]
                                        )

    np.einsum('ijkljkm->ijklm',
            trTensor_complete)[:] += trTensor_testing


    # Quarantine policy
    # ------------------

    # Check if policy is "on"
    if (
            t >= (tStartQuarantineCaseIsolation - realStartDate).days
        )*(
            t <   (tStopQuarantineCaseIsolation - realStartDate).days
        ):
        # New quarantining only happens to people who are transitioning already from untested to virus positive state
        # Therefore here we DO use non-diagonal transitions, and we
        #     redistribute the transtion rates given the testing (which was previously assumed not to create transition in isolation state)
        trTensor_complete = trFunc_quarantine(
                                                trTensor_complete,
                                                t,
                                                trTensor_testing,
                                                **kwargs["trFunc_quarantine_params"]
                                            )




    # Final corrections
    # -----------------



    # TODO: simulate aging and normal birth / death (not terribly important on these time scales, but should be quite simple)


    # Ensure that every "row" sums to 0 by adding to the diagonal (doesn't create new people out of nowhere)
    # Extract (writable) diagonal array and subtract the "row"-sums for each initial state
    np.einsum('ijkljkl->ijkl', trTensor_complete)[:] -= np.einsum('...jkl->...', trTensor_complete)


    # Compute the actual derivatives
    dydt = np.einsum('ijkl,ijklmnp->imnp', stateTensor, trTensor_complete) # contract the HS axis, keep age


    if debugReturnNewPerDay:
        """
            If this is true, instead of returning the real dydt,
            return only the positive "incoming" number of people to each state, so we can track "new cases"
            This needs some approximations, as follows:
                1. Take the normal transition tensor (with rates potentially > 0)
                2. From all states re-normalise the outgoing rates to sum at most to 1
                    (if they were less, keep it, if larger, then this represents
                    “in this day, all people will leave this state, in these ratios to these states”)
                3. Multiply only these outgoing rates with the current state
                    (so the result wont keep the same number of people as normal,
                    but only represent the “new incomings” for each state)
        """

        trTensor_complete_newOnly = copy.deepcopy(trTensor_complete)

        # TODO - Think - this is probably unnecessary actually, artifically reduces "new" rates?
#         # Devide each row by the absolute diagonal rate (that is the sum of the row), but only if its larger than 1
#         trTensor_complete_newOnly /= (
#             np.expand_dims(
#                 np.clip(np.abs(np.einsum('ijkljkl->ijkl', trTensor_complete_newOnly)), a_min=1., a_max=np.inf),
#                 axis=[4,5,6]
#             )
#         )

        # Set the diagonals to zero (no preservation, no outgoing, will end up being the incoming only)
        np.einsum('ijkljkl->ijkl', trTensor_complete_newOnly)[:] = 0.

        dydt_newOnly = np.einsum('ijkl,ijklmnp->imnp', stateTensor, trTensor_complete_newOnly)

        dydt = np.stack([dydt, dydt_newOnly], axis=0)


    if debugTransition:
        return np.reshape(dydt, -1), trTensor_complete

    return np.reshape(dydt, -1)



# ## Initialise and run the model

# In[24]:


# Initialise state
stateTensor_init = copy.deepcopy(stateTensor)

# Populate
stateTensor_init[:,0,0,0] = agePopulationTotal

# Move hospital staff to working in hospital
stateTensor_init[:,0,0,0] -= ageNhsClinicalStaffPopulationRatio * agePopulationTotal
stateTensor_init[:,0,3,0] += ageNhsClinicalStaffPopulationRatio * agePopulationTotal

# Move people to hospital according to baseline occupation (move only from normal people, not hospital staff!)
stateTensor_init[:,0,2,0] += initBaselineHospitalOccupancyEquilibriumAgeRatio * stateTensor_init[:,0,0,0]
stateTensor_init[:,0,0,0] -= initBaselineHospitalOccupancyEquilibriumAgeRatio * stateTensor_init[:,0,0,0]




# Infect some young adults/middle-aged people
# stateTensor_init[2:4,0,0,0] -= 1000.#/np.sum(agePopulationTotal)
# stateTensor_init[2:4,1,0,0] += 1000.#/np.sum(agePopulationTotal)

# BETTER! - People get infected by travel in early stages!


# In[25]:
#!/usr/bin/env python # coding: utf-8 # # COVID-19 model for policy makers in the United Kingdom #  <p align="center"> #  <img src="images/dynamicalModel.png" width="70%"> #  </p> # We use an extended [SEIR model](https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#The_SEIR_model) to capture available information both about the disease progression, as well as how accessible various disease states are by testing. Being tested might cause a transition in the <span style="display: inline-block;background-color:#FFFCCC">Testing State</span>; the rate of such a transition depends both on the <span style="display: inline-block;background-color:#D1E2FF">Health State</span> as well the parameters of the test used. # # Due to age being both an extra risk factor as well as a potential way for decision makers to clearly apply policies to only parts of the total population, we include it directly in our model, and most transition rates are strongly influenced by <span style="display: inline-block;background-color:#FFD4D1">Age State</span>. # # Finally, the main policy making tool as well as conundrum nowadays is the implementation of quarantining and social distancing in order to keep hospitals and medical staff under tolerable levels of pressure. We represent <span style="display: inline-block;background-color:#C2EDC0">Isolation States</span> to investigate the effects of various interventions on policy targets relating to hospitalisation rates and economic freedom, while describing the different health outcomes via the <span style="display: inline-block;background-color:#D1E2FF">Health States</span>. # # ### Health states and Disease progression # # Susceptible people may be Exposed to the virus by mixing with other infected people (E,A,I<sub>1</sub> and I<sub>2</sub>). # # They may get through the infection Asymptomatic and eventually Recover, or become symptomatic and more Infectious, spreading the disease, and potentially Dying or Recovering. # # Recovered people develop more effective antibodies against the virus, and are considered immune<sup>*</sup>. # # # | State | Description | Testing | # | ----- | ----------- | -------- | # | S | Susceptible | Negative | # | E | Exposed | Very weakly virus positive # | A | Asymptomatic | Weakly virus positive # | I<sub>1</sub> | Symptomatic early | Strongly virus positive # | I<sub>2</sub> | Symptomatic late | Medium virus positive <br>Weakly IgM antibody positive # | R<sub>1</sub> | Recovered early | IgM antibody positive # | R<sub>2</sub> | Recovered late | IgM/IgG antibody positive # | D | COVID-related death | May be virus or antibody positive # # # <sub><sup>*</sup>We plan to consider partial / short-term immunity, see further discussion in [Continued research](#header_contd).</sub> # # Code # # The below consists of the implementation of the model described above, with extended comments including assumptions, sources of data and areas to improve upon # In[1]: # Use available data up until this day; cutoff is important due to more recent data being less complete. # # Packages and Helper functions # # To preserve the single-notebook formulation, we include all packages used as well as subfunctions here # # To skip to the start of model implementation, <a href="#modelImplementation">click here</a>! # In[2]: # Packages # Basic packages # help read excel files directly from source into pandas # Building parameter/computation graph # OS/filesystem tools # Distributed computing tools # In[3]: # Regroup various age-group representations into our internal one, and vice versa # first dimension is ages, others dont matter. # Add a zero at beginning for calculations # Add inf at end for calculations # Redistribute to the new bins by calculating how many years in from_ind-1:from_ind falls into each output bin # Define the relative number of ages if we have to distribute between second to last and last age groups # PARAMETER DICTIONARIES AND TABLES # ----------------------------------------------------------------------------------------- # Build the nested parameter/computation graph of a single function. This function iterates through all inputs of a function,     and saves the default argument names and values into a dictionary.      If any of the default arguments are functions themselves, then recursively (depth-first) adds an extra field to     the dictionary, named <funcName + "_params">, that contains its inputs and arguments.      The output of this function can then be passed as a "kwargs" object to the highest level function,     which will then pass the parameter values to the lower dictionary levels appropriately # Check if there are any default parameters, if no, just return empty dict # Save the default argument # If the default argument is a function, inspect it for further # print(argname) # Do a mapping between dictionary and parameter table row and vice versa (for convenient use) # Flatten the dictionary into a table with a single row (but many column): # Iterate through the dictionary to find all keys not ending in "_params", # and add them to the table with name <preString + key> # # If the key doesn end in "_params", then append the key to preString, in call this function on the value (that is a dict) # For the rare case where we want to keep an empty dictionary, the above for cycle doesn't keep it # enable to pass a default dict (if paramTable is incomplete), in which we'll just add / overwrite the values # Find all "_params_" in the argName, and for each step more and more down in the dictionary # We're at the correct level of dictionary # step to the next level of dictionary # <a name="modelImplementation"></a> # # Model implementation # # The model is governed by these two main tensors: # - State tensor: a 4th order tensor containing axes: #     - Age groups #     - Health states #     - Isolation states #     - Testing states # #     In our extended state tensor formulation, we also keep track of not just people currently in each state, but also people newly arriving to each state, as a large number of available datasets refer to "new cases" or "new hospitalisations" each day, rather than current state occupancies normally represented by ODE solutions. # # # - Dynamic state transition rate tensor #     - Rates that govern all allowed transitions between states #     - Dynamically recomputed every iteration, based on number of infected, their current social mixing and number and types of tests available, amongst other variables. #     - Briefly: #         - No transition between age groups #         - No transitions between testing states without performing a test #         - No transitions into S or out of D and R_IgG health states #     - Allowed transitions are as showcased in the model image above #     - Represented by a 7th order, sparse tensor, containing all transitions except age (unimportant due to relatively short time scales compared to coarse age grouping) # # ``` # NOTICE # ------ # THE "MODEL IMPLEMENTATION" SECTION CONTAINS A LARGE NUMBER OF PARAMETER VALUES SET TO A DEFAULT VALUE. # THESE ARE LARGELY INFORMED BY DATA, BUT NOT FIXED! # THEY ARE VARIED DURING THE FITTING OF THE MODEL, ACCORDING TO HOW UNCERTAIN WE ARE IN THE PARAMETER # ``` # # The priors are defined <a href="#defineEnsemblePriors">below the model</a>. Although many of our uncertain / weak assumptions are signalled by "TODO" comments, we feel that the overall conclusions would not be affected by finding better parameter values, especially given our fitting and exploration approach. # # ### The state tensor # In[4]: # State Dimensions # Health states (S, E and D are fixed to 1 dimension) # number of sympyomatic infected states # number of total infected states (disease stages), the +2 are Exposed and I_nonsymptomatic # number of recovery states (antibody development post-disease, IgM and IgG are two stages) # number of total health states, the +2: S, D are suspectible and dead # Age groups (risk groups) # In accordance w Imperial #13 report (0-9, 10-19, ... 70-79, 80+) # Isolation states # None/distancing, Case isolation, Hospitalised, Hospital staff # Testing states # untested/negative, Virus positive, Antibody positive, Both positive # ### Transition rates (1 / day) # # The full transition rate structure is an 8th order tensor, # mapping from any 4D state in the state tensor, to any other 4D state in the state tensor # # However, many of these transitions are non-sensical (ie a 3 year old cannot suddenly become 72, or a dead person be infected again), therefore during the construction of the full model below, we fill in the rates on all "allowed" transitions. # # We attempt to do so based on existing data either describing particular rates (like COVID-related hospitalisation), # or data that helps constrain the product or ratios of multiple rates (such as the R0, or the case fatality ratio [noting this latter depends on testing policy and test availability]). # # Further, to massively reduce the number of weakly constrained parameters, we will approximate many of the weakly correlated transition rates as rank 1 (uncorrelated) matrices. For example the rate of hospitalisation for a patient at a given age and stage of infection will be computed as a product of two indepent rates, one based purely on the age (older people are generally more at risk of hospitalisation), and the other purely on how far the patient has progressed into the disease. This allows us to estimate more of required parameters from available published data. # # There of course still is a lot of uncertainty about how the virus behaves, and all of the data that we use is likely uncomplete and noisy. In order to better represent the things we do not know, we use advanced machine learning techniques, and investigate many possible scenarios (settings of parameters) and for all analyses we retain all plausible scenarios (various parameter settings that explain the available data well enough). # # Any policies we suggest for the near future are investigated for all plausible scenarios, such that decision makers know how likely each one will work as expected in these uncertain times. We further note that as we progress further into the pandemic, the number of plausible scenarios reduces more and more, enabling us to see the way out clearer and clearer. # # # In[5]: # Population (data from Imperial #13 ages.csv/UK) #agePopulationTotal = 1000.*pd.read_csv("https://raw.githubusercontent.com/ImperialCollegeLondon/covid19model/master/data/ages.csv").iloc[3].values[2:] # Currently: let's work with england population only instead of full UK, as NHS England + CHESS data is much clearer than other regions # (google england/uk population 2018, assuming age dist is similar) # Helper function to adjust average rates to age-aware rates This is a helper function and wont be picked up as a model parameter! # just to start the while loop below #print(out) # ## Getting infected # In[7]: # Getting infected # ------------------ # We wish to calibrate overall infection rates to match # - previous R0 estimates, # - available age-attack-ratios, # Age-dependent mixing affects state transition S -> I1 (data available eg Imperial #13 report) # The mixing-related data is nowhere to be found! # This is an Age x Age symmetric matrix, showing which groups mix with which other ones. # Data from DOI: 10.1097/EDE.0000000000001047 via http://www.socialcontactdata.org/tools/ interactive tool in data folder # This is assumed to be contacts per day (but may need to be time-rescaled) # Symmetrise these matrices (not sure why they aren't symmetric) # For simplicity, let's assume scenario of perfect isolation in state-issued home quarantine, see commented below for alternatives #isolationEffectComparedToDistancing = 3. # TODO - find better numbers for proper isolation mixing estimation! #ageSocialMixingIsolation = ageSocialMixingBaseline/(isolationEffectComparedToDistancing * np.mean(ageSocialMixingBaseline/ageSocialMixingDistancing)) # For the S->I1 transition we also need a product mapping, # as the AS->AI1 rate is variable and depend on all AI via social mixing (ages) and transmission rates (I stages) # this vector is nI long only, calibrated together with other variables to reproduce overall R0 # These numbers should represent rate of transmission given contact [will be multiplied by social mixing matrices] # We vary this during model fitting # In[8]: # Getting Infected in the Hospital # --------------------------------------- # The general experience is that infections spread faster in a hospital environment, # we capture this intuition with an age-independent but increased "social Mixing" amongst hospital patients and staff # TODO - This requires further data-driven calibration! # Capture S->I1 within hospital, given the number of total infected inside hospitals # TODO - fact check this number, atm just set based on intuition # Called "Nosocomial viral infection", some data: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5414085/ # HAP: hospital acquired pneumonia, apparently quite common # more data: https://cmr.asm.org/content/14/3/528 # on covid-19: https://www.thelancet.com/journals/lanpub/article/PIIS2468-2667(20)30073-6/fulltext "Nosocomial infection risk among health-care workers and patients has been identified as a research gap to be prioritised in the next few months by WHO." # In[9]: # Also add new infected from travelling based on time-within-simulation # TODO - get real travel data to make these numbers more realistic. For now based on the following assumptions: # - people's age distribution in travel is square of the usual age distribution # - travel rates declined from a base rate as a sigmoid due to border closures, with given mean and slope # - infection rates due to travel are modelled as a gamma pdf over time, with given peak value, loc, and scale parameter # Time (int, in days) within simulation # How many people normally travel back to the country per day # TODO - get data # nAge x T TODO get some realistic data on this # 1 x T TODO get some realistic data on this, maybe make it age weighted # In[10]: # Overall new infections include within quarantine and hospital infections # ------------------------------------------------------------------------ # True / False, no default because it's important to know which one we use at any moment! # True / False, no default because it's important to know which one we use at any moment! All new infections, given infected people on all different isolation states (normal, home, hospital)     We use the following assumptions:      - Infectiousness only depends on infection stage, not age or location      - Hospitalised people are assumed to only mix with other hospitalised people (this includes staff!),     in a non-age-dependent manner: withinHospitalSocialMixing      If policySocialDistancing is True     - Non- and home-isolated people mix with non- and home isolated via ageSocialMixingDistancing (averaging interactions)      If policySocialDistancing is False, we assume home-isolation is taken more seriously, but with little effect on non-isolated people     - Non-isolated people mix with each other via ageSocialMixingBaseline, and with home-isolated people via ageSocialMixingIsolation     - Home-isolated people do not mix with each other      This separation will help disentangle the effects of simply a blanket lessening of social distancing     (keeping the policy True but with less effective ageSocialMixingDistancing matrix),     vs case isolation (policy = False, but with serious ageSocialMixingIsolation) # Add non-hospital infections #-------------------------------- # Add baseline interactions only between non-isolated people # all infected in non-isolation # If the immunity passports policy is on, everyone who tested antibody positive, can roam freely # Therefore replace the interactions between people with testingState = 2 with ageSocialMixingBaseline # we do this by using the distributive property of matrix multiplication, and adding extra interactions # "ageSocialMixingBaseline"-"curNonIsolatedSocialMixing" with each other (this is zero if no social distancing!) # TODO - this is a bit hacky?, but probably correct - double check though! # all infected in non-isolation # Add isolation interactions only between isolated and non-isolated people # non-isolated contracting it from isolated # all infected in isolation # isolated contracting it from non-isolated # all infected in non-hospital, non-isolation # isolated cannot contracting it from another isolated # Add in-hospital infections (of hospitalised patients, and staff) #-------------------------------- # (TODO - within hospitals we probably want to take into effect the testing state; #      tested people are better isolated and there's less mixing) # all infected in hospital (sick or working) # Normalise the rate by total population # ## Hospitalisation and hospital staff # # Disease progression in severe cases naturally leads to hospitalisation before death. # One of the important policy questions we wish to estimate is how many people at any one time would require a hospital bed during the treatment of their disease. # # Hospitalisation is generally a simple situation modeling wise. People with either symptomatic infection (I3-...In states), or for other sicknesses (baseline hospitalisation) end up in hospital. People in S health state may return to non-hospitalised S state, however people in (informed, see later) I state generally remain in hospital until they recovered or dead. # # Home quarantine / social distancing is a different situation. Unlike other reports, here we do not (yet) wish to disentagle the effects of individual quarantine operations (school closures, working from home, social distancing), but rather investigate the effects of current full lockdown (coming into effect on 24 March in the UK), versus testing-based informed individual quarantining. # # Numerically: # # - People in home isolation change their social mixing patterns. The overall social mixing matrix between people in no isolation and home has been estimated via the http://www.socialcontactdata.org/tools/ software, see details in the data_cleaning notebook, this will determine the S->I transition overall. # # - People in hospitals (sick) dramatically reduce their contacts outside the hospital, but increase the chance of transmission within the hospitalised community. For the purpose of this simulation, hospital staff will also in effect be suspecitble to higher risk of infection due to "hospitalised" patients and they also keep their normal interaction. # # - Reported numbers regarding pressure on the health system will report both COVID-19 and non-COVID-19 patients # # In[11]: # Getting Hospitalised # --------------------------------------- # Describe the transitions to-from hospitals # Note that this implementation will assume that hospitalisation takes an extra day, # due to the discrete nature of the simulation, might need to be re-thought. # -> if simulation of a single day is done in multiple steps (first disease progression, then potential hospitalisation), #.    then this problem is avoided. Can do the same with testing. # Further we assume that hospitalisation does not change health state, # but if happens in a non-S state, then it persists until R1 or D # (this may need to be relaxed for early untested I states, where the hospitalisation is not COVID-related). # Hospitalisation mainly depends on disease severity # Baseline hospitalisation rate (Data from Scotland: https://www.isdscotland.org/Health-Topics/Hospital-Care/Publications/Acute-Hospital-Publication/data-summary/) #hospitalisationRateBaseline = 261278./(91.*(5.425*10**6)) # hospitalisation / (period in days * population) -> frac of pop hospitalised per day #hospitalisationRecoveryRateBaseline = 1./4.2 # inverse of mean length of stay in days # Larger data driver approaches, with age distribution, see data_cleaning_R.ipynb for details # Calculate initial hospitalisation (occupancy), that will be used to initialise the model # Take into account the NHS work-force in hospitals that for our purposes count as "hospitalised S" population, # also unaffected by quarantine measures # Extra rate of hospitalisation due to COVID-19 infection stages # TODO - find / estimate data on this (unfortunately true rates are hard to get due to many unknown cases) # Symptom to hospitalisation is 5.76 days on average (Imperial #8) # We do know at least how age affects these risks: # For calculations see data_cleaning_py.ipynb, calculations from CHESS dataset as per 05 Apr # Build the transition tensor from any non-hospitalised state to a hospitalised state # (being in home quarantine is assumed to affect only the infection probability [below], not the hospitalisation probability) # caseIsolationHospitalisationRateAdjustment = 1. # This function takes as input the number of people in given age and health state, and in any non-hospitalised state # and returns the number of people staying in the same age and health state, # but now hospitalised (the rest of people remain in whatever state they were in) # This tensor will pointwise multiply an nAge x nHS slice of the stateTensor # Adjust death rate by age dependent disease severity # Add baseline hospitalisation to all non-dead states # Add COVID-caused hospitalisation to all infected states (TODO: This is summation of rates for independent processes, should be correct, but check) # Recovery rates (hospital discharge) # ------------------------------------ # Higher-than-normal discharge rate for people who recovered (as they were likely to be in hospital mostly due to the virus) # TODO - check with health experts if this is correct assumption; probably also depends on testing state # Baseline discharges apply to all non-symptomatic patients (TODO: take into account testing state!) # No discharges for COVID symptomatic people from the hospital until they recover # TODO - check with health experts if this is correct assumption; probably also depends on testing state # TODO - think of how the latest changes (no prenatal care, no elective surgeries, etc) changed the default hospitalisation rate #trTensor_HospitalAdmission[:,5] # TODO!!! - adjust disease progression transitions so that it shifts direct death probabilities to hospitalised death probabilities # ## Disease progression # #  - assumed to be strictly age and infection stage dependent distributions (progression rates), doesn't depend on other people #  - distinct states represent progression, not necessarly time, but only forward progression is allowed, and the inverse of rates represent average number of days in progression #  - there is a small chance of COVID death from every state, but we assume death is most often preceeded by hospitalisation #  - there is a chance of recovery (and becoming immunised) from every state # # We wish to calibrate these disease progression probabilities to adhere to observed data / earlier models # - serial interval distribution suggests time-to-transmission of Gamma(6.5 days, 0.62) MODEL [Imperial #13] # Symptom progression (All params with relatively wide confidence intervals) # - infect-to-symptom onset is assumed 5 days mean MODEL [AceMod, https://arxiv.org/pdf/2003.10218.pdf] # - symptom-to-death is 16 days DATA_WEAK [Imperial #8] # - symptom-to-discharge is 20.5 days DATA_WEAK [Imperial #8] # - symptom-to-hospitalisation is 5.76 days DATA_WEAK [Imperial #8] # - hospitalisation-to-recovery is 14.51 days DATA_WEAK [Imperial #8] # all the above in Imperial #8 is largely age dependent. Raw data available in data/ImperialReport8_subset_international_cases_2020_03_11.csv # # In[12]: # Based on England data (CHESS and NHS England) # I want a way to keep this as the "average" disease progression, but modify it such that old people have less favorable outcomes (as observed) # But correspondingly I want people at lower risk to have more favorable outcome on average # For calculations see data_cleaning_py.ipynb, calculations from NHS England dataset as per 05 Apr #ageRelativeDiseaseSeverity = np.array([-0.8, -0.6, -0.3, -0.3, -0.1, 0.1, 0.35, 0.4, 0.5]) # FIXED (above) - this is a guess, find data and fix #ageRelativeRecoverySpeed = np.array([0.2]*5+[-0.1, -0.2, -0.3, -0.5]) # TODO - this is a guess, find data and fix # For now we make it same for everyone, makes calculations easier # For calculations see data_cleaning_py.ipynb, calculations from NHS England dataset as per 05 Apr # Basic parameters to adhere to # number of days between measurable events #symptom_to_death = 16., # 20.5, #unrealiticly long for old people # Age related parameters # for now we'll assume that all hospitalised cases are known (overall 23% of hospitalised COVID patients die. 9% overall case fatality ratio) # Unknown rates to estimate # Now we have all the information to build the age-aware multistage SIR model transition matrix # The full transition tensor is a sparse map from the Age x HealthState x isolation state to HealthState, # and thus is a 4th order tensor itself, representing a linear mapping # from "number of people aged A in health state B and isolation state C to health state D. # Use basic parameters to regularise inputs # Numbers nonsymptomatic is assumed to be 86% -> E->IN / E-IS1 = 0.86/0.14 # Nonsymptomatic recovery # Disease progression matrix # TODO - calibrate (together with transmissionInfectionStage) # rows: from-state, cols: to-state (non-symmetric!) # - this represent excess deaths only, doesn't contain baseline deaths! # Calculate all non-serious cases that do not end up in hospitals. # Note that we only have reliable death data from hospitals (NHS England), so we do not model people dieing outside hospitals # to: E,   IN,   IS1,   IS2,    R1,   R2,   D # from E # from IN # from IS1 # from IS2 # from R1 # from R2 # from D # Modify all death and R1 rates: # Adjust death rate by age dependent disease severity # Adjust recovery rate by age dependent recovery speed # Calculate hospitalisation based rates, for which we do have data. Hospitalisation can end up with deaths # Make sure that the ratio of recoveries in hospital honour the case fatality ratio appropriately # IS2 -> death # IS2 -> recovery # multiply by cfr / (1-cfr) to get correct rate towards death # TODO - time to death might be incorrect overall without an extra delay state, especially for young people # Non-hospitalised disease progression # hospitalised disease progression # ## Testing # # In this section we describe multiple types of tests (PCR, antigen and antibody), and estimate their sensitivity and specificity in different health stages. These are thought to be the same for patients of all ages, and isolation states at this time. # # We then model the transitions to other testing states, which are largely policy-based. # # To model the current data (up to 03 April 2020): # - only PCR tests have been done in the UK # - PCR tests are thought to be carried out almost exclusively on symptomatic patients, to determine if their symptoms are caused by SARS-CoV2 or some other infection (this helps us determine the baseline ILI symptoms in practice, to predict true negative rates of the tests given the SARS-infected vs non-SARS-infected (but ILI symptom producing) populations). # # One aim of this complete model is to enable policy makers to make decisions now, based on predicted test availability in the future, therefore most testing-related concerns will be hypotheticals. That said, we aim to accurately model the tests' capabilities based on extensive literature research, and also aim to bring stable policy-level outcomes despite the actual numbers may be inaccurate. # # Two important questions answered by integrating this section into the epidemiology model above will be: # #     1. In what ratio we should produce antibody and antigen lateral flow immunoassay tests? They require the same production capabilities and reagents, there is a question ideally suited to the policy making level # #     2. At what level of testing capabilities (PCR, antigen and antibody) can the country lessen the complete lockdown, without risking lives or overburdening the NHS? # # # # API: # # - trFunc_testing(stateTensor, t, policyFunc, testSpecifications, trFunc_testCapacity): #     - This is the main transition rate function, it returns transition rates from and to all testing states # # - policyFunc #     - Returns a testing policy about what states are tested with how many of which test # # - testSpecifications #     - Details the FPR/FNR of individual tests given the health state # # - trFunc_testCapacity(t) #     - outputs how many tests are available at time t of the different test types modelled # In[13]: # Test parameters # --------------- # assumptions about practical (not theoretical, see discrapancy in PCR!) parameters of tests # TODO - but particular data and references from lit (or estimates based on previous similar tests) # TODO - MANUAL! - this function is VERY specific to current health state setup, and needs to be manually edited if number of health states change #, "Infection stage"],#, "Sensitivity", "Specificity"], # what information state does a pos test transition you to. # what information state does a pos test transition you to. # In some health states some people are true negatives and some are true positives! (No, makes litte sense to use, just account for it in FPR? Only matters for test makers...) # testSpecifications['AmbiguousPosHealthState'] = [np.arange(nI+1, nI+nR+1)]*nHS + [np.arange(nI+1, nI+nR+1)]*nHS + [np.arange(1, nI+1)]*nHS # what information state does a pos test transition you to. # These numbers below are "defaults" illustrating the concept, but are modified by the inputs!!! # ratio of positive (infected / immune) people missed by the test # For each health stage: #  S -> I1 (asymp) -> I2 (mild symp) -> I3 (symp, sick) -> I4 (symp, less sick) -> R1 / R2 (IgM, IgG avail) -> D # PCR # Antigen # Antibody # ratio of negative (non-infected or not immune) people deemed positive by the test # PCR # Antigen # Antibody # In[14]: # In[15]: # TODO - think if we should introdce an "autopsy test" posthumously, categorising people as tested after death? # How is this done, is there data on its sens/spec? # Testing capacity # ---------------- # Assumptions about the testing capacity available at day d of the simulation # For PCR - we will model this (for now, for fitting we'll plug in real data!), as the sum of two sigmoids: #   - initial stage of PHE ramping up its limited capacity (parameterised by total capacity, inflection day and slope of ramp-up) #   - second stage of non-PHE labs joining in and ramping up capacity (this hasn't happened yet, but expected soon! same parameterisation) # For the antigen / antibody tests we define a single sigmoidal capacity curve (starting later than PCR, but with potentially much higher total capacity) # We further define a ratio between the production of the two, due to them requiring the same capabilities. # time within simulation (day) # PCR capacity - initial # PCR capacity - increased # Antibody / antigen capacity # Returns a dictionary with test names and number available at day "t" #phe phase #whole country phase # Real life data on test capacity and who got tested # --------------------------------------------------- # Ignore too old and too recent data points # Get number of tests per age group # Change age groups to reflect our groupings # In[16]: # In[17]: # Symptom parameters # ------------------ # Estimating the baseline ILI-symptoms from earlier studies as well as the success rate of COVID-19 tests # ILI rate estimate from 2018-19 PHE Surveillance of influenza and other respiratory viruses in the UK report: # https://assets.publishing.service.gov.uk/government/uploads/system/uploads/attachment_data/file/839350/Surveillance_of_influenza_and_other_respiratory_viruses_in_the_UK_2018_to_2019-FINAL.pdf # TODO - get actual seasonal symptom rate predictions (from 2020 non-SARS respiratory viruses, this data surely exists) # (daily rate estimate from Figure 8 of the report) # Respiratory diagnosis on hospital admissions (not just ILI, all, TODO - get only ILI?) # NHS Hosp episode statistics 2018-19, page 12 https://files.digital.nhs.uk/F2/E70669/hosp-epis-stat-admi-summ-rep-2018-19-rep.pdf # In hospital: 1.1 million respiratory episodes out of 17.1 million total episodes # Symptom rate in general non-hospitalised population # Symptom rate in hospitalised population This function defines the non-COVID ILI symptoms rate in the population at a given t time # TODO, add extra data etc as input. For now: # In[18]: # Distribute tests amongst (a given subset of) symptomatic people distribute tests amongst symptomatic people     people is nAge x nHS-1 x ... (excluding dead) # Calculate noncovid, but symptomatic people # Subtract already tested people # Check if we already tested everyone with a different test # avoid numerical instabilities # test rate # avoid dividing by zero # tests used to achieve this # In[19]: # Testing policies (how to distribute available tests) # ---------------------------------------------------- # Estimate at any one time how many people are getting tested (with which tests) from which health states # Test types (names correspoding to testSpecifications) # = ["PCR", "Antigen", "Antibody"], # Test Capacity (dict with names above and numbers available on day t) # = trFunc_testCapacity(t) # OPTIONAL ARGUMENTS (may be different for different policy functions, should come with defaults!) # This has these values (for now), {"none", "hospworker_then_random", "virus_positive_only", "virus_positive_only_hospworker_first"} # Baseline symptoms Returns a rate distribution of available test types over age, health and isolation states     (although age assumed not to matter here) # Output nAge x nHS x nIso x nTest x len(testTypes) tensor # Testing capacity is testsAvailable # Get sympom ratio. [0] - general, [1] - hospitalised # PCR testing # ----------- # Hospitalised people get priority over PCR tests # hospitalised non-positive people, exclude tested and dead people # Prioritise hospital workers next: # TODO: check if we should do this? In UK policy there was a 15% max for hospital worker testing until ~2 April... # Distribute PCRs left over the other populations # Distribute PCRs left over the other populations # Antigen testing # --------------- # Hospitalised people get priority over PCR tests # hospitalised non-positive people, exclude tested and dead people # Prioritise hospital workers next: # TODO: check if we should do this? In UK policy there was a 15% max for hospital worker testing until ~2 April... # Distribute Antigen tests left over the other symptomatic people # Distribute antigen tests left over the other non-symptmatic populations # Antibody testing # ---------------- # For now: give to hospital workers first, not taking into account previous tests or symptoms # basically workers get antibody tested regardless of symptoms # Afterwards let's just distribute randomly in the rest of the population # basically people get antibody tested regardless of symptoms # For now: give to hospital workers first, not taking into account previous tests or symptoms # basically workers get antibody tested regardless of symptoms # Afterwards let's just distribute randomly in the rest of the population # TODO: Maybe prioratise people who tested positive for the virus before??? # basically people get antibody tested regardless of symptoms # basically people get antibody tested regardless of symptoms # In[20]: # Define reTesting policy(s) (ie give tests to people in non-0 test states!) # Test types (names correspoding to testSpecifications) # = ["PCR", "Antigen", "Antibody"], # Test Capacity (dict with names above and numbers available on day t) # = trFunc_testCapacity(t) # OPTIONAL ARGUMENTS (may be different for different policy functions, should come with defaults!) # This basic policy will: # - do PCRs on symptomatic hospitalised people # - do PCRs on symptomatic hospital staff # - do PCRs on symptomatic non-hospitalised people # If PCRs run out at any stage, we use antigen tests with same priorisation # Afterwards given fractions of remaining antigen tests are distributed amongst people given these ratios and their earlier testing status: #retesting_antigen_viruspos_ratio = 0.1, # find virus false positives # UPDATE <- retesting viruspos is same ratio is normal testing, as long as they're not in quarantine already! # find immunity false positives # The rest of antigen tests are given out randomly # Antibody tests are used primarily on people who tested positive for the virus #  (set in basic_policyFunc!, use "virus_positive_only_hospworker_first"!) # Afterwards we can use the remaining on either random people (dangerous with many false positives!) # or for retesting people with already positive immune tests to make sure they're still immune, # controlled by this ratio: #distributeRemainingToRandom = True, # TODO - otherwise stockpile for future, how? # Output nAge x nHS x nIso x nTest x len(testTypes) tensor # First distribute tests to symptomatic people as usual: # inpArgs change to not distributing tests randomly: # Run the basic policy function with these modified parameters # We assume PCRs tend to run out done on symptomatic people in 0 Test state, so no retesting via PCR. # Antigen testing # --------------- # Retesting immune positive people # immune positive people # set to 1. for ignoring symptom vs non-symptom # Distribute antigen tests left over the other non-symptmatic populations # UPDATE <- here we use tests equally distributed among people with negative or positive previous virus tests, # as long as they are in non-quarantined state (isoState 0) # TODO - hospital worker testing??? # non-quarantined virus positive people # Antibody testing # ----------------- # Retesting antibody positive people # virus positive people # set to 1. for ignoring symptom vs non-symptom # Afterwards let's just distribute randomly in the rest of the population # basically people get antibody tested regardless of symptoms # In[21]: #policyFunc = policyFunc_testing_symptomaticOnly, Returns a tensor of rates transitioning to tested states # Check if we have real data on the administered tests # Add the current data on within-hospital PCRs carried out already # We do have data, just fill it in # TODO - fix this very hacky solution accessing symptomatic ratio as a subfunc of the policy func # Use hospitalised patient symptom ratio # set the symptomatic ratio of symptomatic states to 1 # true number of tests on given day per age group # Calculate in what ratio we distribute the tests to people along disease states based on symptomatic (age is given in data!) # Divide by total people in each state to get testing rate # we don't have data, follow our assumed availability and policy curves # policyFunc returns stateTensor x testTypes tensor of test administration rates # Compute the transition ratio to tested states, given the administered tests # Get the appropriate slices from testsAdmin. and testSpecs # Set output positive test state based on current test state # already positive for the given test # If already positive for both, stay positive # Transition 0->1, 0->2, 1->2, 1->3 or 2->3 # Where do we go after negative test based on where we are now? # Negatives stay negatives # go to only virus or antibody positive from both positive # go to 0 if tested for the one you're positive for # stay where you are if you test negative for the one you didnt have anyway # Get the transition rates based on current health states # Add the true positives * (1-FNR) # Add the false positives * FPR # Add the false negatives (FNR) # Add the true negatives * (1-FNR) #, testsAdministeredRate # ## Quarantine policies # # This section describes alternatives to the social distancing by full lockdown (that is implemented as a change in the socialMixing matrices). # # One alternative is case isolation, either by hospitalisation or by home isolation. We will assume that all non-symptomatic people who test positive are home isolated along with families for nDaysInIsolation days. Symptomatic people have a chance of being immediately hospitalised instead of sent into home isolation # In[22]: # This is used to establish who gets tests and how many of those end up positive. # (days) time from testing positive to actually getting isolated # On average this many people get hospitalised (compared to home isolation), but modulated by age (TODO: values > 1? clip for now..) # TODO - define this in global variable and just pass here! This function redistributes testing rates, so they dont only create a testing state update, but also an isolation state update # ignore dead # Send a fraction of people (normal) who are symptomatic and tested positive to hospital, based on their age # The rest to home isolation # Send all non-symptomatic (normal) who tested freshly positive to home isolation # Release people from home isolation after isolation period # Hospitalised people are assumed to be released after recovery, with normal rates (TODO: think if this is correct) # TODO!!! - importantly, hospital workers are not being home isolated / hospitalised under this policy. # How to keep track of hospital workers who get hospitalised or home isolated themselves, # such that they get back to being hospital workers afterwards? # A simple (slightly incorrect) solution would be to just implement a non-specific "pull" from isoState=0 people to hospital workers to fill up the missing people? # But the rate of this pull would be impossible to compute and would still be incorrect. Gotta think more on this. # Update the whole tensor accordingly # Make a copy for safety: # First remove all the iso 0->0, test 0,1->1, 2,3->3 transitions (as they're all either hospitalised or sent to home isolation) # Newly virus positive, newly home-isolated, diagonal in disease state transition # Newly virus positive, newly hospitalised, diagonal in disease state transition # Home isolated people are "let go" after nDaysInHomeIsolation, without changing disease or testing state # (TODO: represent multiple testing / needing negative tests to let go, etc - hard problem!) # (UPDATE: multiple testing have now been represented, but for now we'll still let go people based on fixed time rather than negative test, to save tests!) # Return the full updated tensor (so NOT += outside, but actually =) # ## Full simulation function # In[23]: # Function that computes the right side of the non-lin model ODE # Might be double the normal size (as first dimension) _withNewOnlyCopy, if debugReturnNewPerDay # debug # Now implemented by default into state iteration # Dimensions # Input functions and tensors # ---------------------------- # Health state updates # Initial incoming travel-based infections (before restrictions) # Hospitalisation and recovery # Policy changes (on social distancing for now) (TODO - possibly make more changes) # Testing #policyFunc_testing = policyFunc_testing_symptomaticOnly, #testSpecifications = testSpecifications, #trFunc_testCapacity = trFunc_testCapacity, #trFunc_testCapacity_param_testCapacity_antigenratio_country = 0.3 # Initialise return # the input has 2 copies of the state tensor, second copy being the cumulative incomings # Initialise the full transition tensor # Disease condition updates # --------------------------- # Get disease condition updates with no isolation or test transition ("diagonal along those") # all non-hospitalised disease progression is same # Compute new infections (0->1 in HS) with no isolation or test transition ("diagonal along those") # Also add new infected from travelling of healthy people, based on time-within-simulation (this is correct with all (0,0) states, as tested or isolated people dont travel) # Hospitalisation state updates # ----------------------- # Hospitalisation and recovery rates # We assume for now that these only depend on age and disease progression, not on testing state # (TODO - update this given new policies) # The disease and testing states don't change due to hospitalisation. # Hospital staff is treated as already hospitalised from all aspects expect social mixing, should suffice for now # TODO - Could try to devise a scheme in which hospital staff gets hospitalised and some recoveries from hospitalised state go back to hospital staff. # TODO - same issue with hospital staff home isolating; that's probably more important question! # Add recovery from hospital rates # TODO - again here (for now) we assume all discharged people go back to "normal state" instead of home isolation, have to think more on this # Testing state updates # --------------------- # trFunc_testing returns a stateTensor x testStates output #      after the policyFunc assigns tests that are evaluated according to testSpecifications # Diagonal (no transitions) in age, health state and isolation state # (for now, probably TODO: testing positive correlates with new hospitalisation!) # Quarantine policy # ------------------ # Check if policy is "on" # New quarantining only happens to people who are transitioning already from untested to virus positive state # Therefore here we DO use non-diagonal transitions, and we #     redistribute the transtion rates given the testing (which was previously assumed not to create transition in isolation state) # Final corrections # ----------------- # TODO: simulate aging and normal birth / death (not terribly important on these time scales, but should be quite simple) # Ensure that every "row" sums to 0 by adding to the diagonal (doesn't create new people out of nowhere) # Extract (writable) diagonal array and subtract the "row"-sums for each initial state # Compute the actual derivatives # contract the HS axis, keep age If this is true, instead of returning the real dydt,             return only the positive "incoming" number of people to each state, so we can track "new cases"             This needs some approximations, as follows:                 1. Take the normal transition tensor (with rates potentially > 0)                 2. From all states re-normalise the outgoing rates to sum at most to 1                     (if they were less, keep it, if larger, then this represents                     “in this day, all people will leave this state, in these ratios to these states”)                 3. Multiply only these outgoing rates with the current state                     (so the result wont keep the same number of people as normal,                     but only represent the “new incomings” for each state) # TODO - Think - this is probably unnecessary actually, artifically reduces "new" rates? #         # Devide each row by the absolute diagonal rate (that is the sum of the row), but only if its larger than 1 #         trTensor_complete_newOnly /= ( #             np.expand_dims( #                 np.clip(np.abs(np.einsum('ijkljkl->ijkl', trTensor_complete_newOnly)), a_min=1., a_max=np.inf), #                 axis=[4,5,6] #             ) #         ) # Set the diagonals to zero (no preservation, no outgoing, will end up being the incoming only) # ## Initialise and run the model # In[24]: # Initialise state # Populate # Move hospital staff to working in hospital # Move people to hospital according to baseline occupation (move only from normal people, not hospital staff!) # Infect some young adults/middle-aged people # stateTensor_init[2:4,0,0,0] -= 1000.#/np.sum(agePopulationTotal) # stateTensor_init[2:4,1,0,0] += 1000.#/np.sum(agePopulationTotal) # BETTER! - People get infected by travel in early stages! # In[25]: