diff --git a/exploration/expl_esm_labels.py b/exploration/expl_esm_labels.py
index a0e706c..996ac3a 100644
--- a/exploration/expl_esm_labels.py
+++ b/exploration/expl_esm_labels.py
@@ -7,7 +7,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.13.0
+#       jupytext_version: 1.14.5
 #   kernelspec:
 #     display_name: straw2analysis
 #     language: python
@@ -15,19 +15,24 @@
 # ---
 
 # %%
-import os
-import sys
 import datetime
 
 import seaborn as sns
 
-nb_dir = os.path.split(os.getcwd())[0]
-if nb_dir not in sys.path:
-    sys.path.append(nb_dir)
 import participants.query_db
-from features.esm import *
-from features.esm_JCQ import *
-from features.esm_SAM import *
+from features.esm import clean_up_esm, get_esm_data, preprocess_esm
+from features.esm_JCQ import reverse_jcq_demand_control_scoring
+from features.esm_SAM import extract_stressful_events
+
+# import os
+# import sys
+# nb_dir = os.path.split(os.getcwd())[0]
+# if nb_dir not in sys.path:
+#     sys.path.append(nb_dir)
+
+
+# %%
+save_figs = True
 
 # %%
 participants_inactive_usernames = participants.query_db.get_usernames(
@@ -60,7 +65,8 @@ df_esm_PANAS_daily_means = (
 )
 
 # %% [markdown]
-# Next, calculate mean, median, and standard deviation across all days for each participant.
+# Next, calculate mean, median,
+# and standard deviation across all days for each participant.
 
 # %%
 df_esm_PANAS_summary_participant = (
@@ -68,8 +74,8 @@ df_esm_PANAS_summary_participant = (
     .agg(["mean", "median", "std"])
     .reset_index(col_level=1)
 )
-df_esm_PANAS_summary_participant.columns = df_esm_PANAS_summary_participant.columns.get_level_values(
-    1
+df_esm_PANAS_summary_participant.columns = (
+    df_esm_PANAS_summary_participant.columns.get_level_values(1)
 )
 df_esm_PANAS_summary_participant[
     "PANAS_subscale"
@@ -78,9 +84,11 @@ df_esm_PANAS_summary_participant[
 )
 
 # %%
-sns.displot(
+fig1 = sns.displot(
     data=df_esm_PANAS_summary_participant, x="mean", hue="PANAS_subscale", binwidth=0.2
 )
+if save_figs:
+    fig1.figure.savefig("PANAS_mean_participant.png", dpi=300)
 
 # %%
 sns.displot(
@@ -91,9 +99,11 @@ sns.displot(
 )
 
 # %%
-sns.displot(
+fig2 = sns.displot(
     data=df_esm_PANAS_summary_participant, x="std", hue="PANAS_subscale", binwidth=0.05
 )
+if save_figs:
+    fig2.figure.savefig("PANAS_std_participant.png", dpi=300)
 
 # %%
 df_esm_PANAS_summary_participant[df_esm_PANAS_summary_participant["std"] < 0.1]
@@ -131,7 +141,9 @@ df_esm_SAM_daily_events = (
 )
 
 # %% [markdown]
-# Calculate the daily mean of YES (1) or NO (0) answers to the question about a stressful events. This is then the daily ratio of EMA sessions that included a stressful event.
+# Calculate the daily mean of YES (1) or NO (0) answers
+# to the question about stressful events.
+# This is then the daily ratio of EMA sessions that included a stressful event.
 
 # %%
 df_esm_SAM_event_summary_participant = (
@@ -139,12 +151,14 @@ df_esm_SAM_event_summary_participant = (
     .agg(["mean", "median", "std"])
     .reset_index(col_level=1)
 )
-df_esm_SAM_event_summary_participant.columns = df_esm_SAM_event_summary_participant.columns.get_level_values(
-    1
+df_esm_SAM_event_summary_participant.columns = (
+    df_esm_SAM_event_summary_participant.columns.get_level_values(1)
 )
 
 # %%
-sns.displot(data=df_esm_SAM_event_summary_participant, x="mean", binwidth=0.1)
+fig6 = sns.displot(data=df_esm_SAM_event_summary_participant, x="mean", binwidth=0.1)
+if save_figs:
+    fig6.figure.savefig("SAM_events_mean_participant.png", dpi=300)
 
 # %%
 sns.displot(data=df_esm_SAM_event_summary_participant, x="std", binwidth=0.05)
@@ -155,7 +169,12 @@ sns.displot(data=df_esm_SAM_event_summary_participant, x="std", binwidth=0.05)
 # %% [markdown]
 # * Example of threat: "Did this event make you feel anxious?"
 # * Example of challenge: "How eager are you to tackle this event?"
-# * Possible answers: 0 - Not at all, 1 - Slightly, 2 - Moderately, 3 - Considerably, 4 - Extremely
+# * Possible answers:
+#   0 - Not at all,
+#   1 - Slightly,
+#   2 - Moderately,
+#   3 - Considerably,
+#   4 - Extremely
 
 # %%
 df_esm_SAM_daily = (
@@ -177,8 +196,8 @@ df_esm_SAM_summary_participant = (
     .agg(["mean", "median", "std"])
     .reset_index(col_level=1)
 )
-df_esm_SAM_summary_participant.columns = df_esm_SAM_summary_participant.columns.get_level_values(
-    1
+df_esm_SAM_summary_participant.columns = (
+    df_esm_SAM_summary_participant.columns.get_level_values(1)
 )
 
 # %%
@@ -203,12 +222,14 @@ sns.displot(
 )
 
 # %%
-sns.displot(
+fig3 = sns.displot(
     data=df_esm_SAM_threat_challenge_summary_participant,
     x="std",
     hue="event_subscale",
     binwidth=0.1,
 )
+if save_figs:
+    fig3.figure.savefig("SAM_std_participant.png", dpi=300)
 
 # %% [markdown]
 # ## Stressfulness of period
@@ -253,8 +274,8 @@ df_esm_JCQ_summary_participant = (
     .agg(["mean", "median", "std"])
     .reset_index(col_level=1)
 )
-df_esm_JCQ_summary_participant.columns = df_esm_JCQ_summary_participant.columns.get_level_values(
-    1
+df_esm_JCQ_summary_participant.columns = (
+    df_esm_JCQ_summary_participant.columns.get_level_values(1)
 )
 df_esm_JCQ_summary_participant[
     "JCQ_subscale"
@@ -265,11 +286,23 @@ df_esm_JCQ_summary_participant[
 )
 
 # %%
-sns.displot(
-    data=df_esm_JCQ_summary_participant, x="mean", hue="JCQ_subscale", binwidth=0.1,
+fig4 = sns.displot(
+    data=df_esm_JCQ_summary_participant,
+    x="mean",
+    hue="JCQ_subscale",
+    binwidth=0.1,
 )
+if save_figs:
+    fig4.figure.savefig("JCQ_mean_participant.png", dpi=300)
 
 # %%
-sns.displot(
-    data=df_esm_JCQ_summary_participant, x="std", hue="JCQ_subscale", binwidth=0.05,
+fig5 = sns.displot(
+    data=df_esm_JCQ_summary_participant,
+    x="std",
+    hue="JCQ_subscale",
+    binwidth=0.05,
 )
+if save_figs:
+    fig5.figure.savefig("JCQ_std_participant.png", dpi=300)
+
+# %%