# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # RAMP Lockdown Experiments
#
# Run two experiments:
#
#  1. _Baseline_ (uses the [`baseline.yml`](../../model_parameters/baseline.yml) parameters file) that runs the calibrated model under normal conditions.
#  2. _Lockdown earlier_ (uses the [`week_earlier.yml`](../../model_parameters/week_earlier.yml) parameters file) supposes that the UK's national lockdown in March 2020 started a week earlier.

# ## Import modules

# +
import multiprocessing as mp
import numpy as np
import yaml # pyyaml library for reading the parameters.yml file
import os
import pandas as pd
import unittest
import pickle
import copy
import random
import matplotlib.pyplot as plt
import scipy.stats as stats
import re


from microsim.opencl.ramp.run import run_headless
from microsim.opencl.ramp.snapshot_convertor import SnapshotConvertor
from microsim.opencl.ramp.snapshot import Snapshot
from microsim.opencl.ramp.params import Params, IndividualHazardMultipliers, LocationHazardMultipliers
from microsim.opencl.ramp.simulator import Simulator
from microsim.opencl.ramp.disease_statuses import DiseaseStatus

import sys
sys.path.append('..')
#import experiments_functions  # For the ones outside the class
from opencl_runner import OpenCLRunner # Some additional notebook-specific functions required (functions.py)

# Useful for connecting to this kernel
# %connect_info
# -

# ## Setup parameters
#
# Prepare the parameters for the OpenCL model. Can choose between the baseline scenario (by reading [`baseline.yml`](../../model_parameters/baseline.yml) parameters file) or the earlier lockdown_ scenario (by reading [`week_earlier.yml`](../../model_parameters/week_earlier.yml)).
#
# (Note: see the `create_params()` function in [main.py](https://github.com/Urban-Analytics/RAMP-UA/blob/master/microsim/main.py) for an example of how to set parameters in the main codebase). 

# +
#PARAM_FILE = "baseline.yml"
PARAM_FILE = "week_earlier.yml"

PARAMETERS_FILE = os.path.join("..", "../", "model_parameters", PARAM_FILE )

PARAMS = OpenCLRunner.create_parameters(parameters_file=PARAMETERS_FILE)
# -

# ### Get snapshot path
#
# **NB** this is the path to the OpenCL snapshot file generated by running `microsim/main.py`. You need to initilaise the model at least once to create the snapshot. Also, when changing scenario you need to re-initialise the model and create a new snapshot.
#
# The lines below automatically delete any previous snapshots, re-initialise the model, then load the new snapshot.

PARAM_DIR = os.path.join("model_parameters", PARAM_FILE)

PARAMS.place_hazard_multipliers # Check the params file was read

# +
# Now delete any existing snapshots and re-create

# Move into the RAMP-UA root directory 
# %cd '../../'
# Delete any existing caches
# !rm -f "microsim/opencl/snapshots/cache.npz"
# Re-run the model
# !python microsim/main.py -ocl -init -p $PARAM_DIR
# Set the working directory back to the one where this script is located
# %cd 'experiments/calibration/'
# -

OPENCL_DIR = "../../microsim/opencl"
SNAPSHOT_FILEPATH = os.path.join(OPENCL_DIR, "snapshots", "cache.npz")
assert os.path.isfile(SNAPSHOT_FILEPATH), f"Snapshot doesn't exist: {SNAPSHOT_FILEPATH}"

# ## Observation Data
#
# Read the real observations (number of hospital admissions in Devon) that will be used to calibrate the model. See the [README](./observation_data/README.md) for information about how these observations were obtained. They aren't the raw cases, it's actually a model that was fitted to the lagged cases. They need to be made cumulative as this is how they will be compared to the model.

# +
# New per day:
gam_cases = pd.read_csv(os.path.join("..", "..", "gam_cases_new.csv"), header = 0, names=["Day", "Cases"],)

# Cumulative
OBSERVATIONS = pd.DataFrame( {"Day": gam_cases['Day'], "Cases": gam_cases.cumsum()['Cases']} )

assert OBSERVATIONS.tail(1)['Cases'].values[0] == sum(gam_cases['Cases'])

plt.plot(gam_cases['Day'], gam_cases['Cases'])
print(f"Total cases: {sum(gam_cases['Cases'])}")
# -

# ## Run default (manually calibrated) model
#
# This shows what happens with the 'default' (manually calibrated) model

# +
ITERATIONS = 100  # Number of iterations to run for
USE_GPU = False
STORE_DETAILED_COUNTS = True
REPETITIONS = 10
USE_HEALTHIER_POP = False

assert ITERATIONS < len(OBSERVATIONS), \
    f"Have more iterations ({ITERATIONS}) than observations ({len(OBSERVATIONS)})."

# Initialise the class so that its ready to run the model.
# This isn't actually necessary immediately as the `run_opencl_model_multi` function is a static method
# so doesn't read any of the class parameters, but the init is necessary
# for calibration later when some parameters can't be passed to the run function directly
OpenCLRunner.init(
    iterations = ITERATIONS, 
    repetitions = REPETITIONS, 
    observations = OBSERVATIONS,
    use_gpu = USE_GPU,
    use_healthier_pop = USE_HEALTHIER_POP,
    store_detailed_counts = STORE_DETAILED_COUNTS, 
    parameters_file = PARAMETERS_FILE, 
    opencl_dir = OPENCL_DIR, 
    snapshot_filepath = SNAPSHOT_FILEPATH
)
# -

# Results from the manually-calibrated model
manual_results = OpenCLRunner.run_opencl_model_multi(
    repetitions=REPETITIONS,  # Don't use the default, want slightly more robust results
    iterations=ITERATIONS,
    params=PARAMS,
    opencl_dir=OPENCL_DIR, 
    snapshot_filepath=SNAPSHOT_FILEPATH, 
    use_gpu=USE_GPU,
    store_detailed_counts=False,  # Get full info to plot age breakdowns
    multiprocess=False,
    random_ids=False,
    use_healthier_pop=USE_HEALTHIER_POP
)

# +
# Can optionally save/load ('pickle') the results 

# Save
#pkl_out = re.sub( ".yml","",PARAM_FILE)
#pkl_out = pkl_out + "_" + str(REPETITIONS) + ".p"
#pickle.dump(manual_results, open(os.path.join("../../","output", pkl_out), "wb" ) )

# Load
#manual_results = pickle.load( open( os.path.join("../../","output", pkl_out), "rb" ) )
# -

# Get the summary results of of the finished model

manual_summaries = [x[0] for x in manual_results]
# Store the results as they can be useful as hypothetical observations to test some of the calibration algorithms
pseudo_observations = OpenCLRunner.get_cumulative_new_infections(manual_summaries)

# +
matrix_ifr = np.zeros(shape=(REPETITIONS,3))

for rep in range(REPETITIONS):
   matrix_ifr[rep,0] = manual_summaries[rep].total_counts[6][ITERATIONS -1]
   matrix_ifr[rep,1] = manual_summaries[rep].total_counts[5][ITERATIONS -1]
   matrix_ifr[rep,2] = ((manual_summaries[rep].total_counts[6][ITERATIONS -1])/(manual_summaries[rep].total_counts[6][ITERATIONS -1] + manual_summaries[rep].total_counts[5][ITERATIONS -1]))*100
    
np.mean(matrix_ifr[:,2])


# +
matrix_ifr = np.zeros(shape=(ITERATIONS*REPETITIONS,9))

for rep in range(REPETITIONS):
   row_ind = range((rep)*ITERATIONS,(rep+1)*ITERATIONS)
   matrix_ifr[row_ind,0] = rep+1
   matrix_ifr[row_ind,1] = range(1,ITERATIONS+1)
   matrix_ifr[row_ind,2] = manual_summaries[rep].total_counts[0]
   matrix_ifr[row_ind,3] = manual_summaries[rep].total_counts[1]
   matrix_ifr[row_ind,4] = manual_summaries[rep].total_counts[2]
   matrix_ifr[row_ind,5] = manual_summaries[rep].total_counts[3]
   matrix_ifr[row_ind,6] = manual_summaries[rep].total_counts[4]
   matrix_ifr[row_ind,7] = manual_summaries[rep].total_counts[5]
   matrix_ifr[row_ind,8] = manual_summaries[rep].total_counts[6]


csv_out = re.sub( ".yml","",PARAM_FILE)
csv_out = csv_out + "_" + str(REPETITIONS) + ".csv"
csv_path_out = os.path.join("../../","output", csv_out)

# Can optionally save the results as a csv
#np.savetxt(csv_path_out, matrix_ifr, delimiter=",")
# -

# ## Plot output summary data
#
# ### Total counts of disease status

# +
def plot_summaries(summaries, observations=None, plot_type="error_bars"):

    #fig, ax = plt.subplots(1, len(DiseaseStatus), sharey=True)
    fig, ax = plt.subplots(1, 1, figsize=(10,7))
    
    # Work out the number of repetitions and iterations
    iters, reps = _get_iters_and_reps(summaries)
    x = range(iters)
    total_not_susceptible = np.zeros(iters)  # Used to compare to observations
    
    for d, disease_status in enumerate(DiseaseStatus):

        # Calculate the mean and standard deviation
        mean, sd = OpenCLRunner.get_mean_total_counts(summaries, d, get_sd=True)
        
        # Don't plot susceptible or recovered as it hides everytihng else
        if disease_status==DiseaseStatus.Susceptible or disease_status==DiseaseStatus.Recovered:
            continue
        
        if plot_type == "error_bars":
            ax.errorbar(x, mean, sd, label=f"{disease_status}" )
        elif plot_type == "lines":
            for rep in range(reps):
                ax.plot(x, matrix[rep], label=f"{disease_status} {rep}", 
                        color=plt.cm.get_cmap("hsv", len(DiseaseStatus))(d) )

    if observations is not None:
        # Plot the observations (cumulative infections)
        ax.plot(x, observations.loc[:len(x)-1, "Cases"], 
                label=f"Observations (cumulative cases)", color="black", linestyle="dashed")
        # And the total new infections, for comparison
        ax.plot(x, OpenCLRunner.get_cumulative_new_infections(summaries),
               label=f"Total not susceptible ", color="grey", linestyle="dashed")
        
    
    ax.legend() 
    ax.set_title("Disease Status")
    ax.set_xlabel("Iteration")
    ax.set_ylabel("Number of cases")
    #ax.set_ylim(0, 5000)
    #ax.set_xlim(0,30)

def _get_iters_and_reps(summaries):
    reps = len(summaries)
    iters = len(summaries[0].total_counts[0])
    return (iters, reps)


# -

plot_summaries(summaries=manual_summaries, observations=OBSERVATIONS, plot_type="error_bars")

# +
#plot_summaries(summaries=summaries, plot_type="lines")

StarcoderdataJupyterScriptsDedupFiltered