%md 

You may find this series of notebooks at https://github.com/databricks-industry-solutions/psm. For more information about this solution accelerator, visit https://www.databricks.com/blog/2020/10/20/detecting-at-risk-patients-with-real-world-data.html.

%md

# Data preparation and QC

In this notebook we start by creating cohorts based on our study design using synthea resources that are already loaded into delta (using `1-data-ingest` notebook).

%md

## 0. Initial Setup

First we run `./0-config` notebook to configure our project's environment by setting up the base path, deltalake path, mlflow experiment etc. We then run `./cohort_builder` notebook to import `DeltaEHR` class that is designed to make it easy to create cohorts based on synthea resources.

%md
Now we specify our experiment's main parameters, namely: target medication (intervention under study), target event (event defining cohort entry date), and the target outcome (outcome under study)

%md
We also would like to include information regarding past histroy of comorbidities (for example obesity etc) that can be later used in our propencity score matching

%md
## 1. cohort creation
To make cohort creation and ETL easier, we created a class `DeltaEHR` that makes data pre-processing easier. Using the available methods we create target and outcome cohorts, based on our inclusion/exclusion criteria. First we start by adding our target cohort (patients diagnosed with covid) and outcome cohort (patients who have been admitted to the hospital). Eeach cohort is then automatically added to the collection of cohorts. Each chort contains three columns, patinet id (`PATIENT`), cohort start index date (`START`) and cohort exit date (`STOP`).

%md
next we specify which demographic information to inlclude in the dataset

%md

Now we add cohorts based on prior events (comorbidities, drug exposure etc). For each comorbid condition of interest, we choose a window of time 
to go back and look for any record of diagnosis of, or exposure to, of a condition of interest. This is done simply by using the `add_cohort` method 
defined in `DeltaEHR` class. This method also allows you to specify a washout window (`gate_window`) as a buffer in cases where we want to ensure effects of treatments do not interfere. For example if this is specified to
10 days, then if there is an instance of a comorbidity diagnosis within 10 dyas of the target ourcome, we do not inlcude that event. Note that
if you speciy a negative value for the washout window you can include evnets occuring after the target event (see below)

<img src="https://drive.google.com/uc?export=view&id=1O2OmMrJS97FvX1lzrc4xOck5mii1c5cm" width=700>

%md
Now we add the cohort of patinets that have received the target treatment with 10 days of being diagnosed with covid. This is done the same way as adding cohorts based on historic events, with the difference that in this case we set `hist_winodw=0` and `gate_window=-10`

%md
Optionally you can also add cohorts correspodning to other treatments, for example:

```
meds_test=(
  delta_ehr.tables['medications'].filter("to_date(START) > to_date('2020-01 01')")
  .join(
    delta_ehr.cohorts['covid'].select('PATIENT'),on='PATIENT')
  .join(delta_ehr.cohorts['admission'].select('PATIENT'),on='PATIENT')
  .groupBy('CODE','DESCRIPTION')
  .count()
  .orderBy(desc('count'))
  .limit(20)
  .collect()
)

medications={f"m_{m['CODE']}":m['CODE'] for m in meds_test}
for med,codes in list(medications.items()):
  delta_ehr.add_cohort('medications', med, codes,0,-10, 'covid')
```

%md
we can also add cohort of patients experiencing other symptoms such as blood clots within 20 days of diagnosis with covid

%md
## 2. create dataset
Now we creat the final dataset for our downstream analysis. One of the methods in `DeltaEHR` is `get_cohort_tags()`. This method combines all cohort information in form of columns of indicator functions (cohort membership indicator) and the cohort index date correspoding to each patient id.

%md
## 3. Exploratory Analysis
Now let's take a look at the dataset and look at trends, such as number of individuals diagnosed with the target condition (covid) over time, and demographic trends and other statistics of interest.

%md
Now let's take a look at some of the trends regarding reported blood clot occurances, such as dirtribution of blood clots among covid patients vs others, distribution of the timeframe within which a blood clot is reported, and look at demographic patterns that may be correlated with these timeframes.

%md
Now let's compare admission rates among patinets who have received the treatment and those who have not

%md
as we see, overall the admission rates are lower among those who have received the traget treatmnet, however this can be confunded by many factors, for example it can be the case that younger patients are more likely receive the treatment and also being young less likely being admitted to the hospital. In this case we cannot attribute lower admission rates to the treatment. In the next notebook, we use propencity score matching to correct for such confunding factors. But first, let's write the resulting dataset to delta silver layer.

%md
## 4. Write final dataset to Delta
Now we write the resulting dataset back in the delta lake for our next analysis that is specifically looking into the effect of databrixovir on hospital admissions