From d44f776b471ec6f545a1426380d419a7274392f5 Mon Sep 17 00:00:00 2001
From: Szzzesan <98128431+Szzzesan@users.noreply.github.com>
Date: Fri, 23 May 2025 12:15:23 -0400
Subject: [PATCH 1/3] Rie's edits for swapping port assignments and minor fixes
 to the task code.

---
 simple_plots.py                      | 774 +++++++++++++++++++++------
 stand_alone/gui.py                   |  14 +-
 stand_alone/main.py                  |  80 +--
 stand_alone/support_classes.py       | 111 +++-
 stand_alone/user_settings_generic.py |   4 +
 upload_cloud_backup.py               |  40 ++
 6 files changed, 797 insertions(+), 226 deletions(-)
 create mode 100644 upload_cloud_backup.py

diff --git a/simple_plots.py b/simple_plots.py
index de1dbb3..4add81e 100644
--- a/simple_plots.py
+++ b/simple_plots.py
@@ -10,24 +10,31 @@
 import seaborn as sns
 from matplotlib.collections import PatchCollection
 from matplotlib.patches import Rectangle
+import matplotlib.gridspec as gridspec
+
+from upload_cloud_backup import behavior_data_dir
 from user_info import get_user_info
 import shutil
+from datetime import datetime
+import pdb
 
 info_dict = get_user_info()
 initials = info_dict['initials']
 start_date = info_dict['start_date']
+data_dir = behavior_data_dir
 
 
 def get_today_filepaths(days_back=0):
     file_paths = []
-    for root, dirs, filenames in walk(os.path.join(os.getcwd(), 'data')):
-        if len(dirs) == 0 and os.path.basename(root)[:2] == initials:
+    for root, dirs, filenames in walk(data_dir):
+        if len(dirs) == 0 and os.path.basename(root)[:2] in initials:
             mouse = os.path.basename(root)
             for f in filenames:
                 if f == 'desktop.ini':
                     continue
-                file_date = date(int(f[5:9]), int(f[10:12]), int(f[13:15]))
-                dif = date.today() - file_date
+                file_date = datetime.strptime(f[5:-4], '%Y-%m-%d_%H-%M-%S')
+                # file_date = date(int(f[5:9]), int(f[10:12]), int(f[13:15]))
+                dif = datetime.today() - file_date
                 if dif.days <= days_back:
                     # if f[5:15] == time.strftime("%Y-%m-%d"):
                     file_paths.append(os.path.join(mouse, f))
@@ -91,24 +98,39 @@ def gen_data(file_paths, select_mouse=None, return_info=False):
         if select_mouse is not None and mouse not in select_mouse:
             continue
 
-        path = os.path.join(os.getcwd(), 'data', f)
+        path = os.path.join(data_dir, f)
+        meta_data = read_pi_meta(path)
+
         if return_info:
-            data = read_pi_meta(path)
-            # if data['box'] == 'elissapi0':
-            #     session = pd.read_csv(path, na_values=['None'], skiprows=3)
-            #     session_summary(data_reduction(session), mouse)
-            #     ans = input(f'remove broken file? (y/n)\n{path}\n???')
-            #     if ans == 'y':
-            #         file_name = f[6:]
-            #         half_session_path = os.path.join(os.getcwd(), 'data', 'half_sessions', file_name)
-            #         shutil.move(path, half_session_path)
+            data = meta_data
         else:
-            data = pd.read_csv(path, na_values=['None'], skiprows=3)
+            try:
+                data = pd.read_csv(path, na_values=['None'], skiprows=3)
+            except pd.errors.EmptyDataError:
+                print(f'empty file at {path}')
+                continue
             try:
                 data = data_reduction(data)
+
+                #port info as a row
+                port_info_row = [{
+                    'key': meta_data['port1_info']['distribution'],
+                    'port': meta_data['port1_info']['port_num']
+                    # 'phase': last_row_phase,
+                },
+                    {'key': meta_data['port2_info']['distribution'],
+                     'port': meta_data['port2_info']['port_num']
+                     # 'phase': last_row_phase,
+                     }
+                ]
+
+                #Added port info at the end of the dataframe
+                portinfo_rows_df = pd.DataFrame(port_info_row).reindex(columns=data.columns)
+                data = pd.concat([data, portinfo_rows_df], ignore_index=True)
+
             except ValueError:
                 file_name = f[6:]
-                half_session_path = os.path.join(os.getcwd(), 'data', 'half_sessions', file_name)
+                half_session_path = os.path.join(data_dir, 'half_sessions', file_name)
                 if data.session_time.max() < 800:
                     print(f'moving {f} to half sessions, session time: {data.session_time.max():.2f} seconds')
                     shutil.move(path, half_session_path)
@@ -117,6 +139,7 @@ def gen_data(file_paths, select_mouse=None, return_info=False):
                     if ans == 'y':
                         shutil.move(path, half_session_path)
                 continue
+
         if mouse in d.keys():
             d[mouse].append(data)
         else:
@@ -158,9 +181,13 @@ def data_reduction(df, lick_tol=.01, head_tol=.2):
 
 
 def consumption_time(df):
-    bg_end_times = df[(df.key == 'LED') & (df.port == 2) & (df.value == 1)]
-    exp_entries = df[(df.key == 'head') & (df.port == 1) & (df.value == 1)]
+    bgportassignment = df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    expportassignment = df.loc[df['key'] == 'exp_decreasing', 'port'].iloc[-1]
+    bg_end_times = df[(df.key == 'LED') & (df.port == bgportassignment) & (df.value == 1)]
+    exp_entries = df[(df.key == 'head') & (df.port == expportassignment) & (df.value == 1)]
+
     dif = min_dif(bg_end_times.session_time, exp_entries.session_time)
+
     bg_consumption = dif[~np.isnan(dif)]
     if df.task.iloc[10] != 'single_reward':
         consumption_df = pd.DataFrame()
@@ -168,8 +195,8 @@ def consumption_time(df):
         consumption_df['port'] = ['bg'] * len(bg_consumption)
         return consumption_df
 
-    exp_end_times = df[(df.key == 'LED') & (df.port == 1) & (df.value == 1)]
-    bg_entries = df[(df.key == 'head') & (df.port == 2) & (df.value == 1)]
+    exp_end_times = df[(df.key == 'LED') & (df.port == expportassignment) & (df.value == 1)]
+    bg_entries = df[(df.key == 'head') & (df.port == bgportassignment) & (df.value == 1)]
     dif = min_dif(exp_end_times.session_time, bg_entries.session_time)
     exp_consumption = dif[~np.isnan(dif)]
     consumption_df = pd.DataFrame()
@@ -178,12 +205,58 @@ def consumption_time(df):
     return consumption_df
 
 
+def calculate_premature_leave(df, threshold=1.0): #trials with premature leave
+    bgportassignment = df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    premature_leave_trials = 0  # Counter for trials with premature leave
+    blocks = df.phase.dropna().unique()
+    blocks.sort()
+    results = []
+    for block in blocks:
+        block_data = df[df.phase == block]
+        block_total_trials = len(block_data.trial.unique())
+        for trial in block_data.trial.unique():
+            if pd.isna(trial):
+                continue  # Skip invalid trials
+            trial_data = df[df.trial == trial] # Filter data for the current trial
+            bg_on_times = trial_data[(trial_data['key'] == 'trial') &
+                                      (trial_data['value'] == 1)].session_time.values
+            if len(bg_on_times) == 0:
+                continue
+            bg_start = bg_on_times[0]
+            bg_off_times = trial_data[(trial_data['key'] == 'LED') &
+                                      (trial_data['port'] == bgportassignment) &
+                                      (trial_data['value'] == 1)].session_time.values
+            if len(bg_off_times) == 0:  # Check if no BG ON recorded
+                continue
+            bg_end = bg_off_times[0]
+            bg_head_out_times = trial_data[(trial_data['key'] == 'head') & # Extract head-out and head-in times while BG port is active
+                                        (trial_data['port'] == bgportassignment) &
+                                        (trial_data['value'] == 0) &
+                                        (trial_data['session_time'] >= bg_start) &
+                                        (trial_data['session_time'] < bg_end)].session_time.values
+            bg_head_in_times = trial_data[(trial_data['key'] == 'head') &
+                                       (trial_data['port'] == bgportassignment) &
+                                       (trial_data['value'] == 1) &
+                                       (trial_data['session_time'] >= bg_start) &
+                                       (trial_data['session_time'] < bg_end)].session_time.values
+            for head_out in bg_head_out_times: # Check for premature leave: head-out without a valid head-in within threshold
+                if not any((head_in > head_out) and (head_in - head_out) <= threshold for head_in in bg_head_in_times):
+                    premature_leave_trials += 1  # Flag this trial as a premature leave
+                    break  # Stop checking further head-outs in this trial
+        premature_leave_rate = (premature_leave_trials / block_total_trials)
+        results.append({"block": block,"premature leave numbers":premature_leave_trials, "premature leave rate": premature_leave_rate})
+    premature_leave_df = pd.DataFrame(results)
+    return premature_leave_df
+
+
 def block_leave_times(df):
+    bgportassignment = df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    expportassignment = df.loc[df['key'] == 'exp_decreasing', 'port'].iloc[-1]
     reward_trials = df[(df.key == 'reward_initiate')].trial.to_numpy()
     non_reward = ~df.trial.isin(reward_trials)
-    bg_end_times = df[(df.key == 'LED') & (df.port == 2) & (df.value == 1) & non_reward]
-    exp_entries = df[(df.key == 'head') & (df.value == 1) & (df.port == 1) & non_reward]
-    exp_exits = df[(df.key == 'head') & (df.value == 0) & (df.port == 1) & non_reward]
+    bg_end_times = df[(df.key == 'LED') & (df.port == bgportassignment) & (df.value == 1) & non_reward]
+    exp_entries = df[(df.key == 'head') & (df.value == 1) & (df.port == expportassignment) & non_reward]
+    exp_exits = df[(df.key == 'head') & (df.value == 0) & (df.port == expportassignment) & non_reward]
     bg_end_times = bg_end_times[bg_end_times.session_time < exp_entries.session_time.max()]
     ind, dif = min_dif(bg_end_times.session_time, exp_entries.session_time, return_index=True)
     exp_entries = exp_entries.iloc[np.unique(ind)]
@@ -193,8 +266,14 @@ def block_leave_times(df):
     valid_trials = np.intersect1d(valid_trials, bg_end_times.trial.values)
     exp_exits = exp_exits.loc[valid_trials]
     exp_entries = exp_entries.loc[valid_trials]
+
     if len(exp_exits.to_numpy()) != len(exp_entries.to_numpy()):
-        print()
+        print(
+            f"Mismatch in exp_entries and exp_exits. "
+            f"exp_entries: {exp_entries}, exp_exits: {exp_exits}. "
+            "Using clean_entries_exits to resolve."
+        )
+        exp_entries, exp_exits = clean_entries_exits(exp_entries, exp_exits)
     leave_times = exp_exits.to_numpy() - exp_entries.to_numpy()
 
     trial_blocks = bg_end_times[bg_end_times.trial.isin(exp_entries.index.values)].phase.to_numpy()
@@ -208,15 +287,17 @@ def get_entry_exit(df, trial):
     is_trial = df.trial == trial
     start = df.value == 1
     end = df.value == 0
-    port1 = df.port == 1
-    port2 = df.port == 2
+    bgport = df.port == df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    expport = df.port == df.loc[df['key'] == 'exp_decreasing', 'port'].iloc[-1]
+    # port1 = df.port == 1
+    # port2 = df.port == 2
 
     trial_start = df[is_trial & start & (df.key == 'trial')].session_time.values[0]
-    trial_middle = df[is_trial & end & (df.key == 'LED') & port2].session_time.values[0]
+    trial_middle = df[is_trial & end & (df.key == 'LED') & bgport].session_time.values[0] #head in to EXP, bg LED off
     trial_end = df[is_trial & end & (df.key == 'trial')].session_time.values[0]
 
-    bg_entries = df[is_trial & port2 & start & (df.key == 'head')].session_time.to_numpy()
-    bg_exits = df[is_trial & port2 & end & (df.key == 'head')].session_time.to_numpy()
+    bg_entries = df[is_trial & bgport & start & (df.key == 'head')].session_time.to_numpy()
+    bg_exits = df[is_trial & bgport & end & (df.key == 'head')].session_time.to_numpy()
 
     if len(bg_entries) == 0 or len(bg_exits) == 0 or bg_entries[0] > bg_exits[0]:
         bg_entries = np.concatenate([[trial_start], bg_entries])
@@ -225,15 +306,15 @@ def get_entry_exit(df, trial):
     if len(bg_exits) == 0 or bg_entries[-1] > bg_exits[-1]:
         bg_exits = np.concatenate([bg_exits, [trial_middle]])
 
-    exp_entries = df[is_trial & port1 & start & (df.key == 'head') &
+    exp_entries = df[is_trial & expport & start & (df.key == 'head') &
                      (df.session_time > trial_middle)].session_time.to_numpy()
-    exp_exits = df[is_trial & port1 & end & (df.key == 'head') &
+    exp_exits = df[is_trial & expport & end & (df.key == 'head') &
                    (df.session_time > trial_middle)].session_time.to_numpy()
 
     if not (len(exp_entries) == 0 and len(exp_exits) == 0):
-        if len(exp_entries) == 0:
+        if len(exp_entries) == 0: #only exp out
             exp_entries = np.concatenate([[trial_middle], exp_entries])
-        if len(exp_exits) == 0:
+        if len(exp_exits) == 0: #only exp in
             exp_exits = np.concatenate([exp_exits, [trial_end]])
 
         if exp_entries[0] > exp_exits[0]:
@@ -241,12 +322,12 @@ def get_entry_exit(df, trial):
         if exp_entries[-1] > exp_exits[-1]:
             exp_exits = np.concatenate([exp_exits, [trial_end]])
 
-    early_exp_entries = df[is_trial & port1 & start & (df.key == 'head') &
+    early_exp_entries = df[is_trial & expport & start & (df.key == 'head') &
                            (df.session_time < trial_middle)].session_time.to_numpy()
-    early_exp_exits = df[is_trial & port1 & end & (df.key == 'head') &
+    early_exp_exits = df[is_trial & expport & end & (df.key == 'head') &
                          (df.session_time < trial_middle)].session_time.to_numpy()
 
-    if not (len(early_exp_entries) == 0 and len(early_exp_exits) == 0):
+    if not (len(early_exp_entries) == 0 and len(early_exp_exits) == 0): #any early exp in/out
         if len(early_exp_entries) == 0:
             early_exp_entries = np.concatenate([[trial_start], early_exp_entries])
         if len(early_exp_exits) == 0:
@@ -264,114 +345,164 @@ def get_entry_exit(df, trial):
     if len(early_exp_entries) != len(early_exp_exits):
         print()
 
+    if len(exp_entries):
+        if len(exp_exits) != len(exp_entries):
+            print(
+                f"Mismatch in exp_entries and exp_exits in trial {trial}. "
+                f"exp_entries: {exp_entries}, exp_exits: {exp_exits}. "
+                "Using clean_entries_exits to resolve."
+            )
+            exp_entries, exp_exits = clean_entries_exits(exp_entries, exp_exits)
+
+    if len(bg_entries) != len(bg_exits):
+        print(
+            f"Mismatch in bg_entries and bg_exits in trial {trial}. "
+            f"bg_entries: {bg_entries}, bg_exits: {bg_exits}. "
+            "Using clean_entries_exits to resolve."
+        )
+        bg_entries, bg_exits = clean_entries_exits(bg_entries, bg_exits)
+
     return bg_entries, bg_exits, exp_entries, exp_exits, early_exp_entries, early_exp_exits
 
 
+def clean_entries_exits(entries, exits):
+    """
+    Cleans mismatched entries and exits such that each entry is paired with the nearest valid exit.
+    """
+    valid_entries = []
+    valid_exits = []
+    e_idx, x_idx = 0, 0
+
+    while e_idx < len(entries) and x_idx < len(exits):
+        if entries[e_idx] < exits[x_idx]:  # Valid entry-exit pair
+            valid_entries.append(entries[e_idx])
+            valid_exits.append(exits[x_idx])
+            e_idx += 1
+            x_idx += 1  # Move to the next entry and exit
+        else:
+            x_idx += 1  # Skip unmatched exits
+
+    return valid_entries, valid_exits
+
 def percent_engaged(df):
-    travel_time = .5
-    blocks = df.phase.unique()
-    blocks.sort()
-    time_engaged = []
-    block_time = []
-    block_rewards = []
-    for block in blocks:
-        engaged = []
-        all_time = []
-        rewards = []
-        block_trials = df[(df.value == 0) & (df.key == 'trial') & (df.phase == block)].trial
-        for trial in block_trials:
-            bg_entries, bg_exits, exp_entries, exp_exits, _, _ = get_entry_exit(df, trial)
-            is_trial = df.trial == trial
-            start = df.value == 1
-            end = df.value == 0
-            # port1 = df.port == 1
-            # port2 = df.port == 2
-
-            #
-            trial_start = df[is_trial & start & (df.key == 'trial')].session_time.values[0]
-            # trial_middle = df[is_trial & start & (df.key == 'LED') & port2].session_time.values[0]
-            trial_end = df[is_trial & end & (df.key == 'trial')].session_time.values[0]
-            #
-            # bg_entries = df[is_trial & port2 & start & (df.key == 'head')].session_time.to_numpy()
-            # bg_exits = df[is_trial & port2 & end & (df.key == 'head')].session_time.to_numpy()
-            #
-            # if len(bg_entries) == 0 or bg_entries[0] > bg_exits[0]:
-            #     bg_entries = np.concatenate([[trial_start], bg_entries])
-            # if trial_end - bg_entries[-1] < .1:
-            #     bg_entries = bg_entries[:-1]
-            # if len(bg_exits) == 0 or bg_entries[-1] > bg_exits[-1]:
-            #     bg_entries = np.concatenate([bg_exits, [trial_middle]])
-            #
-            # if not (len(bg_entries) == len(bg_exits) and np.all(bg_exits - bg_entries > 0)):
-            #     print('stop')
-            # bg_engaged = sum(bg_exits - bg_entries)
-            #
-            # exp_entries = df[is_trial & port1 & start & (df.key == 'head') &
-            #                  (df.session_time > trial_middle)].session_time.to_numpy()
-            # exp_exits = df[is_trial & port1 & end & (df.key == 'head') &
-            #                (df.session_time > trial_middle)].session_time.to_numpy()
-            #
-            # if len(exp_entries) == 0 and len(exp_exits) == 0:
-            #     exp_engaged = 0
-            # else:
-            #     if len(exp_entries) == 0:
-            #         exp_entries = np.concatenate([[trial_middle], exp_entries])
-            #     if len(exp_exits) == 0:
-            #         exp_exits = np.concatenate([exp_exits, [trial_end]])
-            #
-            #     if exp_entries[0] > exp_exits[0]:
-            #         exp_entries = np.concatenate([[trial_middle], exp_entries])
-            #     if exp_entries[-1] > exp_exits[-1]:
-            #         exp_exits = np.concatenate([exp_exits, [trial_end]])
-            #     exp_engaged = sum(exp_exits - exp_entries)
-            #
-            #     # if not len(exp_entries) == len(exp_exits) and len(exp_entries):
-            #     #     print('stop')
-            #     # if len(exp_entries):
-
-            if len(exp_entries):
-                exp_engaged = sum(exp_exits - exp_entries)
-            else:
-                exp_engaged = 0
-            bg_engaged = sum(bg_exits - bg_entries)
+    try:
+        travel_time = .5
+        blocks = df.phase.dropna().unique()
+        blocks.sort()
+        time_engaged = []
+        block_time = []
+        block_rewards = []
+
+        for block in blocks:
+            engaged = []
+            all_time = []
+            rewards = []
+            block_trials = df[(df.value == 0) & (df.key == 'trial') & (df.phase == block)].trial
+            for trial in block_trials:
+                bg_entries, bg_exits, exp_entries, exp_exits, _, _ = get_entry_exit(df, trial)
+                is_trial = df.trial == trial
+                start = df.value == 1
+                end = df.value == 0
+                # port1 = df.port == 1
+                # port2 = df.port == 2
+
+                #
+                trial_start = df[is_trial & start & (df.key == 'trial')].session_time.values[0]
+                # trial_middle = df[is_trial & start & (df.key == 'LED') & port2].session_time.values[0]
+                trial_end = df[is_trial & end & (df.key == 'trial')].session_time.values[0]
+                #
+                # bg_entries = df[is_trial & port2 & start & (df.key == 'head')].session_time.to_numpy()
+                # bg_exits = df[is_trial & port2 & end & (df.key == 'head')].session_time.to_numpy()
+                #
+                # if len(bg_entries) == 0 or bg_entries[0] > bg_exits[0]:
+                #     bg_entries = np.concatenate([[trial_start], bg_entries])
+                # if trial_end - bg_entries[-1] < .1:
+                #     bg_entries = bg_entries[:-1]
+                # if len(bg_exits) == 0 or bg_entries[-1] > bg_exits[-1]:
+                #     bg_entries = np.concatenate([bg_exits, [trial_middle]])
+                #
+                # if not (len(bg_entries) == len(bg_exits) and np.all(bg_exits - bg_entries > 0)):
+                #     print('stop')
+                # bg_engaged = sum(bg_exits - bg_entries)
+                #
+                # exp_entries = df[is_trial & port1 & start & (df.key == 'head') &
+                #                  (df.session_time > trial_middle)].session_time.to_numpy()
+                # exp_exits = df[is_trial & port1 & end & (df.key == 'head') &
+                #                (df.session_time > trial_middle)].session_time.to_numpy()
+                #
+                # if len(exp_entries) == 0 and len(exp_exits) == 0:
+                #     exp_engaged = 0
+                # else:
+                #     if len(exp_entries) == 0:
+                #         exp_entries = np.concatenate([[trial_middle], exp_entries])
+                #     if len(exp_exits) == 0:
+                #         exp_exits = np.concatenate([exp_exits, [trial_end]])
+                #
+                #     if exp_entries[0] > exp_exits[0]:
+                #         exp_entries = np.concatenate([[trial_middle], exp_entries])
+                #     if exp_entries[-1] > exp_exits[-1]:
+                #         exp_exits = np.concatenate([exp_exits, [trial_end]])
+                #     exp_engaged = sum(exp_exits - exp_entries)
+                #
+                #     # if not len(exp_entries) == len(exp_exits) and len(exp_entries):
+                #     #     print('stop')
+                #     # if len(exp_entries):
+
+                if len(exp_entries):
+                    exp_engaged = sum([exit - entry for entry, exit in zip(exp_entries, exp_exits)])
+                else:
+                    exp_engaged = 0
+
+                bg_engaged = sum([exit - entry for entry, exit in zip(bg_entries, bg_exits)])
 
-            all_time.append(trial_end - trial_start)
-            engaged.append(bg_engaged + exp_engaged)
-            rewards.append(len(df[is_trial & start & (df.key == 'reward')]))
+                all_time.append(trial_end - trial_start)
+                engaged.append(bg_engaged + exp_engaged)
+                rewards.append(len(df[is_trial & start & (df.key == 'reward')]))
 
-        time_engaged.append(sum(engaged) + travel_time * 2 * len(block_trials))
-        block_time.append(sum(all_time))
-        block_rewards.append(sum(rewards))
-    engaged_df = pd.DataFrame()
-    engaged_df['percent engaged'] = np.array(time_engaged) / np.array(block_time)
-    engaged_df['block'] = blocks
-    engaged_df['time engaged'] = time_engaged
-    engaged_df['rewards earned'] = block_rewards
-    engaged_df['reward rate'] = np.array(block_rewards) / np.array(time_engaged)
-    return engaged_df
+            time_engaged.append(sum(engaged) + travel_time * 2 * len(block_trials))
+            block_time.append(sum(all_time))
+            block_rewards.append(sum(rewards))
+        engaged_df = pd.DataFrame()
+        engaged_df['percent engaged'] = np.array(time_engaged) / np.array(block_time)
+        engaged_df['block'] = blocks
+        engaged_df['time engaged'] = time_engaged
+        engaged_df['rewards earned'] = block_rewards
+        engaged_df['reward rate'] = np.array(block_rewards) / np.array(time_engaged)
+
+        return engaged_df
+    except Exception as e:
+        print ("Error in function.")
+        raise
 
 
 def reentry_index(df):
-    is_bg_exit = (df.port == 2) & (df.key == 'head') & (df.value == 0)
-    is_slow_block = df.groupby('trial').phase.agg(pd.Series.mode) == '0.4'
-    is_fast_block = df.groupby('trial').phase.agg(pd.Series.mode) == '0.8'
-    num_ideal_bg_entry_slow = len(np.unique(df.trial.dropna())[is_slow_block])
-    num_bg_entry_slow = len(df.index[is_bg_exit & df.trial.isin(
-        np.unique(df.trial.dropna())[is_slow_block])])
-    num_ideal_bg_entry_fast = len(np.unique(df.trial.dropna())[is_fast_block])
-    num_bg_entry_fast = len(df.index[is_bg_exit & df.trial.isin(
-        np.unique(df.trial.dropna())[is_fast_block])])
-
-    reentry_index_slow = num_bg_entry_slow / num_ideal_bg_entry_slow
-    reentry_index_fast = num_bg_entry_fast / num_ideal_bg_entry_fast
+    bgportassignment = df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    expportassignment = df.loc[df['key'] == 'exp_decreasing', 'port'].iloc[-1]
+    is_bg_exit = (df.port == bgportassignment) & (df.key == 'head') & (df.value == 0)
+    phase_mode = df.groupby('trial').phase.agg(
+        lambda s: s.iloc[-1] if (s.value_counts().get('0.4', 0) == 1 and s.value_counts().get('0.8', 0) == 1)
+        else pd.Series.mode(s).iloc[-1]
+    )
+
+    is_low_block = (phase_mode == '0.4')
+    is_high_block = (phase_mode == '0.8')
+    num_ideal_bg_entry_low = len(np.unique(df.trial.dropna())[is_low_block])  # gets number of low block trials
+    num_bg_entry_low = len(df.index[is_bg_exit & df.trial.isin(
+        np.unique(df.trial.dropna())[is_low_block])])
+    num_ideal_bg_entry_high = len(np.unique(df.trial.dropna())[is_high_block])  # gets number of high block trials
+    num_bg_entry_high = len(df.index[is_bg_exit & df.trial.isin(
+        np.unique(df.trial.dropna())[is_high_block])])
+
+    reentry_index_low = num_bg_entry_low / num_ideal_bg_entry_low if num_ideal_bg_entry_low > 0 else 0
+    reentry_index_high = num_bg_entry_high / num_ideal_bg_entry_high if num_ideal_bg_entry_high > 0 else 0
     reentry_df = pd.DataFrame()
     reentry_df['block'] = ['0.4', '0.8']
-    reentry_df['bg_reentry_index'] = [reentry_index_slow, reentry_index_fast]
+    reentry_df['bg_reentry_index'] = [reentry_index_low, reentry_index_high]
     return reentry_df
 
 
 def add_h_lines(data=None, x=None, y=None, hue=None, ax=None, palette=None, estimator='mean'):
-    days_back = 10
+    days_back = 4
     palette = sns.color_palette(palette)
     for i, hue_key in enumerate(data[hue].unique()):
         df = data[data[hue] == hue_key]
@@ -390,6 +521,7 @@ def merge_old_trials(session):
 
 def simple_plots(select_mouse=None):
     plot_single_mouse_plots = True
+    save_folder = "C:\\Users\\Shichen\\OneDrive - Johns Hopkins\\ShulerLab\\Rie_behavior\\summary_graphs"
     if select_mouse is None:
         dif = date.today() - start_date
         data = gen_data(get_today_filepaths(days_back=dif.days), select_mouse=select_mouse)
@@ -405,73 +537,96 @@ def simple_plots(select_mouse=None):
         consumption = pd.DataFrame()
         block_leaves = pd.DataFrame()
         reentry = pd.DataFrame()
+        premature_leave = pd.DataFrame()
+
         for i, session in enumerate(data[mouse]):
-            if info[mouse][i]['task'] == 'cued_forgo_forced':
-                continue
+            # if info[mouse][i]['task'] == 'multi_reward':
+            #     continue
             try:
                 session = merge_old_trials(session)
 
                 engaged_df = percent_engaged(session)
-                engaged_df['day'] = [i] * len(engaged_df)
+                engaged_df['session'] = [i] * len(engaged_df)
                 engaged = pd.concat([engaged, engaged_df])
 
                 consumption_df = consumption_time(session)
-                consumption_df['day'] = [i] * len(consumption_df)
+                consumption_df['session'] = [i] * len(consumption_df)
                 consumption = pd.concat([consumption, consumption_df])
 
                 block_leaves_df = block_leave_times(session)
-                block_leaves_df['day'] = [i] * len(block_leaves_df)
+                block_leaves_df['session'] = [i] * len(block_leaves_df)
                 block_leaves = pd.concat([block_leaves, block_leaves_df])
 
                 reentry_df = reentry_index(session)
-                reentry_df['day'] = [i] * len(reentry_df)
+                reentry_df['session'] = [i] * len(reentry_df)
                 reentry = pd.concat([reentry, reentry_df])
+
+                premature_leave_df= calculate_premature_leave(session)
+                premature_leave_df['session'] = [i] * len(premature_leave_df)
+                premature_leave = pd.concat([premature_leave,premature_leave_df])
+
             except Exception as e:
-                raise e
+                print(f"Error processing session {i} for mouse {mouse}: {e}")
+                raise
 
         engaged.sort_values('block', inplace=True)
         block_leaves.sort_values('block', inplace=True)
         if plot_single_mouse_plots:
-            fig, axes = plt.subplots(2, 2, figsize=[11, 8], layout="constrained")
-            sns.lineplot(data=block_leaves.reset_index(), x='day', y='leave time', hue='block', ax=axes[0, 0],
+            fig, axes = plt.subplots(3, 2, figsize=[11, 8], layout="constrained")
+            sns.lineplot(data=block_leaves.reset_index(), x='session', y='leave time', hue='block', style = 'block',markers=True, ax=axes[0, 0],
                          palette='Set2')
-            add_h_lines(data=block_leaves.reset_index(), x='day', y='leave time', hue='block', ax=axes[0, 0],
+            add_h_lines(data=block_leaves.reset_index(), x='session', y='leave time', hue='block', ax=axes[0, 0],
                         palette='Set2')
-            sns.lineplot(data=consumption.reset_index(), x='day', y='consumption time', hue='port', ax=axes[0, 1],
+            sns.lineplot(data=consumption.reset_index(), x='session', y='consumption time', hue='port', style = 'port', markers=True, ax=axes[0, 1],
                          palette='Set1', estimator=np.median)
-            add_h_lines(data=consumption.reset_index(), x='day', y='consumption time', hue='port', ax=axes[0, 1],
+            add_h_lines(data=consumption.reset_index(), x='session', y='consumption time', hue='port', ax=axes[0, 1],
                         palette='Set1', estimator='median')
-            sns.lineplot(data=engaged.reset_index(), x='day', y='reward rate', hue='block', ax=axes[1, 0],
+            sns.lineplot(data=engaged.reset_index(), x='session', y='reward rate', hue='block', style = 'block', markers=True,ax=axes[1, 0],
+                         palette='Set2')
+            add_h_lines(data=engaged.reset_index(), x='session', y='reward rate', hue='block', ax=axes[1, 0],
+                        palette='Set2')
+            sns.lineplot(data=engaged.reset_index(), x='session', y='percent engaged', hue='block',style = 'block', markers=True, ax=axes[1, 1],
                          palette='Set2')
-            add_h_lines(data=engaged.reset_index(), x='day', y='reward rate', hue='block', ax=axes[1, 0],
+            add_h_lines(data=engaged.reset_index(), x='session', y='percent engaged', hue='block', ax=axes[1, 1],
                         palette='Set2')
-            sns.lineplot(data=engaged.reset_index(), x='day', y='percent engaged', hue='block', ax=axes[1, 1],
+            sns.lineplot(data=premature_leave.reset_index(), x='session', y='premature leave rate', hue='block', style = 'block',markers=True,ax=axes[2, 0],
                          palette='Set2')
-            add_h_lines(data=engaged.reset_index(), x='day', y='percent engaged', hue='block', ax=axes[1, 1],
+            add_h_lines(data=premature_leave.reset_index(), x='session', y='premature leave rate', hue='block', ax=axes[2, 0],
                         palette='Set2')
 
+            axes[2, 0].axhline(y=0.2, color='red', linestyle='--', linewidth=1.5, label='Threshold = 0.2')
+
+            # Add legend to show the line label
+            axes[2, 0].legend()
+
             axes[0, 0].set_title('Leave Time by Block')
             axes[0, 1].set_title('Consumption Time by Port')
             axes[1, 0].set_title('Reward Rate by Block')
             axes[1, 1].set_title('Percent Time Engaged by Block')
+            axes[2, 0].set_title('Premature leave from BG port by Block')
 
             axes[0, 0].set_ylim([0, 20])
             axes[0, 1].set_ylim([0, 20])
             axes[1, 0].set_ylim([0, .65])
             axes[1, 1].set_ylim([0, 1])
+            axes[2, 0].set_ylim([0, 1])
             plt.suptitle(mouse, fontsize=20)
+            os.makedirs(save_folder, exist_ok=True)
+            # Construct the filename
+            filename = f'{mouse}_session_summary.png'
+            save_path = os.path.join(save_folder, filename)
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+            print(f"Graph saved to: {save_path}")
             plt.show()
 
-        block_leaves_last10_df = block_leaves[(block_leaves.day >= block_leaves.day.max() - 10)].groupby('block')[
+        block_leaves_last10_df = block_leaves[(block_leaves.session >= block_leaves.session.max() - 10)].groupby('block')[
             'leave time'].mean().reset_index()
         block_leaves_last10_df['animal'] = mouse
         block_leaves_last10 = pd.concat([block_leaves_last10, block_leaves_last10_df])
 
     fig, axes = plt.subplots(1, 1)
     sns.boxplot(data=block_leaves_last10.reset_index(), x='block', y='leave time')
-    for mouse in data.keys():
-        plt.plot([-0.1, 0.9], block_leaves_last10[block_leaves_last10.animal == mouse]['leave time'], 'o-',
-                 color='darkgray')
+
     plt.show()
 
 
@@ -492,7 +647,125 @@ def single_session(select_mouse=None, num_back=2):
             session_summary(last_session, mouse, last_info)
 
 
+# def session_summary(data, mouse, info):
+#     base_save_folder= "C:\\Users\\riepo\\Experiment1\\graphs\\each_session"
+#     save_folder = os.path.join(base_save_folder, mouse)
+#     os.makedirs(save_folder, exist_ok=True)
+#     fig, [ax1, ax2] = plt.subplots(1, 2, figsize=[10, 10])
+#     port_palette = sns.color_palette('Set1')
+#     block_palette = sns.color_palette('Set2')
+#     start = data.value == 1
+#     end = data.value == 0
+#     head = data.key == 'head'
+#     lick = data.key == 'lick'
+#     reward = data.key == 'reward'
+#     bgport = data.port == data.loc[data['key'] == 'background', 'port'].iloc[-1]
+#     expport = data.port == data.loc[data['key'] == 'exp_decreasing', 'port'].iloc[-1]
+#     # port1 = data.port == 1
+#     # port2 = data.port == 2
+#     # print(f"port2 is {port2}")
+#     # print(f"bgport is {bgport}")
+#     max_trial = data.trial.max()
+#
+#     bg_rectangles = []
+#     exp_rectangles_in_bg = []
+#     exp_rectangles = []
+#     block1_rectangles = []
+#     block2_rectangles = []
+#     bg_reward_events = []
+#     exp_reward_events = []
+#     bg_lick_events = []
+#     exp_lick_events = []
+#     bg_lengths = []
+#     exp_lengths = []
+#     trial_blocks = data.groupby(['trial'])['phase'].agg(pd.Series.mode)
+#     blocks = data.phase.dropna().unique()
+#     blocks.sort()
+#     for trial in data.trial.unique():
+#         if np.isnan(trial):
+#             continue
+#         is_trial = data.trial == trial
+#         try:
+#             trial_start = data[is_trial & start & (data.key == 'trial')].session_time.values[0]
+#             trial_middle = data[is_trial & end & (data.key == 'LED') & bgport].session_time.values[0]
+#             trial_end = data[is_trial & end & (data.key == 'trial')].session_time.values[0]
+#         except IndexError:
+#             continue
+#
+#         bg_rewards = data[is_trial & start & bgport & reward].session_time.values
+#         exp_rewards = data[is_trial & start & expport & reward].session_time.values
+#         bg_licks = data[is_trial & start & lick & (data.session_time < trial_middle)].session_time.values
+#         exp_licks = data[is_trial & start & lick & (data.session_time > trial_middle)].session_time.values
+#
+#         bg_lengths.append(trial_middle - trial_start)
+#         exp_lengths.append(trial_end - trial_middle)
+#
+#         bg_entries, bg_exits, exp_entries, exp_exits, early_exp_entries, early_exp_exits = get_entry_exit(data, trial)
+#         bg_intervals = list(zip(bg_entries, bg_exits))
+#         exp_intervals = list(zip(exp_entries, exp_exits))
+#         early_exp_intervals = list(zip(early_exp_entries, early_exp_exits))
+#         for [s, e] in bg_intervals:
+#             bg_rectangles.append(Rectangle((s - trial_start, trial), e - s, .7))
+#         for [s, e] in early_exp_intervals:
+#             exp_rectangles_in_bg.append(Rectangle((s - trial_start, trial), e - s, .7))
+#         for [s, e] in exp_intervals:
+#             exp_rectangles.append(Rectangle((s - trial_middle, trial), e - s, .7))
+#         if np.where(blocks == trial_blocks.loc[trial])[0][0] == 0:
+#             block1_rectangles.append(Rectangle((0, trial), 100, 1))
+#         else:
+#             block2_rectangles.append(Rectangle((0, trial), 100, 1))
+#         bg_reward_events.append(bg_rewards - trial_start)
+#         exp_reward_events.append(exp_rewards - trial_middle)
+#         bg_lick_events.append(bg_licks - trial_start)
+#         exp_lick_events.append(exp_licks - trial_middle)
+#
+#     alpha = .5
+#     pc_b1 = PatchCollection(block1_rectangles, facecolors=block_palette[0], alpha=alpha)
+#     pc_b2 = PatchCollection(block2_rectangles, facecolors=block_palette[1], alpha=alpha)
+#     ax1.add_collection(pc_b1)
+#     ax1.add_collection(pc_b2)
+#     pc_b12 = PatchCollection(block1_rectangles, facecolors=block_palette[0], alpha=alpha)
+#     pc_b22 = PatchCollection(block2_rectangles, facecolors=block_palette[1], alpha=alpha)
+#     ax2.add_collection(pc_b12)
+#     ax2.add_collection(pc_b22)
+#
+#     pc_bg = PatchCollection(bg_rectangles, edgecolor=port_palette[0], facecolor='w', alpha=1)
+#     ax1.add_collection(pc_bg)
+#
+#     pc_exp_bg = PatchCollection(exp_rectangles_in_bg, edgecolor=port_palette[1], facecolor='w', alpha=1)
+#     ax1.add_collection(pc_exp_bg)
+#
+#     pc_exp = PatchCollection(exp_rectangles, edgecolor=port_palette[1], facecolor='w', alpha=1)
+#     ax2.add_collection(pc_exp)
+#
+#     offsets = np.array(list(range(len(bg_reward_events)))) + 1.4
+#     ax1.eventplot(bg_reward_events, color='purple', linelengths=.62, lineoffsets=offsets)
+#     offsets = np.array(list(range(len(exp_reward_events)))) + 1.4
+#     ax2.eventplot(exp_reward_events, color='purple', linelengths=.62, lineoffsets=offsets)
+#
+#     light = [.8, .7, .8]
+#     dark = [.2, .2, .2]
+#     offsets = np.array(list(range(len(bg_lick_events)))) + 1.4
+#     ax1.eventplot(bg_lick_events, color=light, linelengths=.25, lineoffsets=offsets)
+#     offsets = np.array(list(range(len(exp_lick_events)))) + 1.4
+#     ax2.eventplot(exp_lick_events, color=light, linelengths=.25, lineoffsets=offsets)
+#
+#     session_summary_axis_settings([ax1, ax2], max_trial)
+#     plt.suptitle(f'{mouse}: {info["date"]} {info["time"]}')
+#
+#     # Construct the filename
+#     filename = f'{mouse}_{info["date"]}_{info["time"]}.png'
+#     save_path = os.path.join(save_folder, filename)
+#
+#     # Save the plot
+#     plt.savefig(save_path, dpi=300, bbox_inches='tight')
+#     print(f"Graph saved to: {save_path}")
+#     plt.show()
+
 def session_summary(data, mouse, info):
+    base_save_folder= save_folder = "C:\\Users\\Shichen\\OneDrive - Johns Hopkins\\ShulerLab\\Rie_behavior\\each_session"
+    save_folder = os.path.join(base_save_folder, mouse)
+    os.makedirs(save_folder, exist_ok=True)
     fig, [ax1, ax2] = plt.subplots(1, 2, figsize=[10, 10])
     port_palette = sns.color_palette('Set1')
     block_palette = sns.color_palette('Set2')
@@ -501,8 +774,8 @@ def session_summary(data, mouse, info):
     head = data.key == 'head'
     lick = data.key == 'lick'
     reward = data.key == 'reward'
-    port1 = data.port == 1
-    port2 = data.port == 2
+    bgport = data.port == data.loc[data['key'] == 'background', 'port'].iloc[-1]
+    expport = data.port == data.loc[data['key'] == 'exp_decreasing', 'port'].iloc[-1]
     max_trial = data.trial.max()
 
     bg_rectangles = []
@@ -517,7 +790,7 @@ def session_summary(data, mouse, info):
     bg_lengths = []
     exp_lengths = []
     trial_blocks = data.groupby(['trial'])['phase'].agg(pd.Series.mode)
-    blocks = data.phase.unique()
+    blocks = data.phase.dropna().unique()
     blocks.sort()
     for trial in data.trial.unique():
         if np.isnan(trial):
@@ -525,13 +798,13 @@ def session_summary(data, mouse, info):
         is_trial = data.trial == trial
         try:
             trial_start = data[is_trial & start & (data.key == 'trial')].session_time.values[0]
-            trial_middle = data[is_trial & end & (data.key == 'LED') & port2].session_time.values[0]
+            trial_middle = data[is_trial & end & (data.key == 'LED') & bgport].session_time.values[0]
             trial_end = data[is_trial & end & (data.key == 'trial')].session_time.values[0]
         except IndexError:
             continue
 
-        bg_rewards = data[is_trial & start & port2 & reward].session_time.values
-        exp_rewards = data[is_trial & start & port1 & reward].session_time.values
+        bg_rewards = data[is_trial & start & bgport & reward].session_time.values
+        exp_rewards = data[is_trial & start & expport & reward].session_time.values
         bg_licks = data[is_trial & start & lick & (data.session_time < trial_middle)].session_time.values
         exp_licks = data[is_trial & start & lick & (data.session_time > trial_middle)].session_time.values
 
@@ -548,10 +821,13 @@ def session_summary(data, mouse, info):
             exp_rectangles_in_bg.append(Rectangle((s - trial_start, trial), e - s, .7))
         for [s, e] in exp_intervals:
             exp_rectangles.append(Rectangle((s - trial_middle, trial), e - s, .7))
-        if np.where(blocks == trial_blocks.loc[trial])[0][0] == 0:
+
+        block_value = trial_blocks.loc[trial]
+        if block_value == "0.4":
             block1_rectangles.append(Rectangle((0, trial), 100, 1))
         else:
             block2_rectangles.append(Rectangle((0, trial), 100, 1))
+
         bg_reward_events.append(bg_rewards - trial_start)
         exp_reward_events.append(exp_rewards - trial_middle)
         bg_lick_events.append(bg_licks - trial_start)
@@ -590,8 +866,15 @@ def session_summary(data, mouse, info):
 
     session_summary_axis_settings([ax1, ax2], max_trial)
     plt.suptitle(f'{mouse}: {info["date"]} {info["time"]}')
-    plt.show()
 
+    # Construct the filename
+    filename = f'{mouse}_{info["date"]}_{info["time"]}.png'
+    save_path = os.path.join(save_folder, filename)
+
+    # Save the plot
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"Graph saved to: {save_path}")
+    plt.show()
 
 def session_summary_axis_settings(axes, max_trial):
     for ax in axes:
@@ -601,17 +884,178 @@ def session_summary_axis_settings(axes, max_trial):
         ax.spines['bottom'].set_visible(True)
         ax.get_yaxis().set_visible(False)
         ax.set_ylim([-1, max_trial + 1])
-        ax.set_xlim([0, 20])
+        ax.set_xlim([0, 30])
         ax.invert_yaxis()
         ax.set_ylabel('Trial')
         ax.set_xlabel('Time (sec)')
 
 
+
+def simple_plots2(select_mouse=None):
+    plot_single_mouse_plots = True
+    save_folder = "C:\\Users\\Shichen\\OneDrive - Johns Hopkins\\ShulerLab\\Rie_behavior\\summary_graphs"
+
+    #insert vertical line between stages or pre/post surgery
+    stage_changes = {
+        # "RK001": {"1->2": "2025-01-22", "2->3": "2025-01-31"},
+        # "RK002": {"1->2": "2025-01-03", "2->3": "2025-03-14"},
+        # "RK003": {"1->2": "2025-01-03", "2->3": "2025-01-22"},
+        # "RK004": {"1->2": "2025-02-24", "2->3": "2025-03-05"},
+        # "RK005": {"1->2": "2025-01-27", "2->3": "2025-01-31"},
+        # "RK006": {"1->2": "2024-12-20", "2->3": "2025-02-26"},
+        # "ES039": {"1->2": "2024-02-14", "2->3": None},
+        # "ES045": {"1->2": "2024-04-24", "2->3": None},
+        # "ES046": {"1->2": "2024-05-02", "2->3": None},
+        # "ES047": {"1->2": "2024-04-24", "2->3": None},
+        "RK001": {"1->2": "2025-04-13", "2->3": None}, #surgery date
+        "RK002": {"1->2": "2025-04-13", "2->3": None},
+        "RK003": {"1->2": "2025-04-13", "2->3": None},
+        "RK004": {"1->2": "2025-03-30", "2->3": None},
+        "RK005": {"1->2": "2025-03-30", "2->3": None},
+        "RK006": {"1->2": "2025-03-30", "2->3": None}
+    }
+
+    if select_mouse is None:
+        dif = date.today() - start_date
+        data = gen_data(get_today_filepaths(days_back=dif.days), select_mouse=select_mouse)
+        info = gen_data(get_today_filepaths(days_back=dif.days), select_mouse=select_mouse, return_info=True)
+    else:
+        data = gen_data(get_today_filepaths(days_back=1000), select_mouse=select_mouse)
+        info = gen_data(get_today_filepaths(days_back=1000), select_mouse=select_mouse, return_info=True)
+    block_leaves_last10 = pd.DataFrame()
+
+    for mouse in data.keys():
+        if select_mouse is not None and mouse not in select_mouse:
+            continue
+
+        # Initialize empty dataframes for aggregated data.
+        engaged = pd.DataFrame()
+        consumption = pd.DataFrame()
+        block_leaves = pd.DataFrame()
+        reentry = pd.DataFrame()
+        premature_leave = pd.DataFrame()
+
+        # Extract session dates
+        session_dates = []
+        for sess_info in info[mouse]:
+            try:
+                dt = datetime.strptime(sess_info['date'], "%Y-%m-%d")
+                session_dates.append(dt)
+            except Exception as e:
+                print(f"Error parsing date in session info: {sess_info.get('date', None)}: {e}")
+
+        for i, session in enumerate(data[mouse]):
+            # if info[mouse][i]['task'] == 'multi_reward':
+            #     continue
+            try:
+                session = merge_old_trials(session)
+
+                engaged_df = percent_engaged(session)
+                engaged_df['session'] = [i] * len(engaged_df)
+                engaged = pd.concat([engaged, engaged_df])
+
+                consumption_df = consumption_time(session)
+                consumption_df['session'] = [i] * len(consumption_df)
+                consumption = pd.concat([consumption, consumption_df])
+
+                block_leaves_df = block_leave_times(session)
+                block_leaves_df['session'] = [i] * len(block_leaves_df)
+                block_leaves = pd.concat([block_leaves, block_leaves_df])
+
+                reentry_df = reentry_index(session)
+                reentry_df['session'] = [i] * len(reentry_df)
+                reentry = pd.concat([reentry, reentry_df])
+
+                premature_leave_df = calculate_premature_leave(session)
+                premature_leave_df['session'] = [i] * len(premature_leave_df)
+                premature_leave = pd.concat([premature_leave, premature_leave_df])
+
+            except Exception as e:
+                print(f"Error processing session {i} for mouse {mouse}: {e}")
+                raise
+
+        engaged.sort_values('block', inplace=True)
+        block_leaves.sort_values('block', inplace=True)
+
+        if plot_single_mouse_plots:
+            fig, axes = plt.subplots(3, 2, figsize=[11, 8], layout="constrained")
+
+            sns.lineplot(data=block_leaves.reset_index(), x='session', y='leave time',
+                         hue='block', style='block', markers=True, ax=axes[0, 0], palette='Set2')
+            add_h_lines(data=block_leaves.reset_index(), x='session', y='leave time',
+                        hue='block', ax=axes[0, 0], palette='Set2')
+
+            sns.lineplot(data=consumption.reset_index(), x='session', y='consumption time',
+                         hue='port', style='port', markers=True, ax=axes[0, 1],
+                         palette='Set1', estimator=np.median)
+            add_h_lines(data=consumption.reset_index(), x='session', y='consumption time',
+                        hue='port', ax=axes[0, 1], palette='Set1', estimator='median')
+
+            sns.lineplot(data=engaged.reset_index(), x='session', y='reward rate',
+                         hue='block', style='block', markers=True, ax=axes[1, 0], palette='Set2')
+            add_h_lines(data=engaged.reset_index(), x='session', y='reward rate',
+                        hue='block', ax=axes[1, 0], palette='Set2')
+
+            sns.lineplot(data=engaged.reset_index(), x='session', y='percent engaged',
+                         hue='block', style='block', markers=True, ax=axes[1, 1], palette='Set2')
+            add_h_lines(data=engaged.reset_index(), x='session', y='percent engaged',
+                        hue='block', ax=axes[1, 1], palette='Set2')
+
+            sns.lineplot(data=premature_leave.reset_index(), x='session', y='premature leave rate',
+                         hue='block', style='block', markers=True, ax=axes[2, 0], palette='Set2')
+            add_h_lines(data=premature_leave.reset_index(), x='session', y='premature leave rate',
+                        hue='block', ax=axes[2, 0], palette='Set2')
+
+            axes[2, 0].axhline(y=0.2, color='red', linestyle='--', linewidth=1.5, label='Threshold = 0.2')
+            axes[2, 0].legend()
+
+            # Add vertical lines for stage changes (blue dashed lines)
+            if mouse in stage_changes:
+                for change, change_date_str in stage_changes[mouse].items():
+                    if change_date_str is not None:
+                        try:
+                            change_date = datetime.strptime(change_date_str, "%Y-%m-%d")
+                            idx = next((j for j, d in enumerate(session_dates) if d >= change_date), None)
+                            if idx is not None and idx > 0:
+                                pos = idx - 0.5
+                                for ax in axes.flatten():
+                                    ax.axvline(x=pos, color='blue', linestyle='--', linewidth=1.5)
+                        except Exception as e:
+                            print(f"Error processing stage change date {change_date_str} for mouse {mouse}: {e}")
+
+            axes[0, 0].set_title('Leave Time by Block')
+            axes[0, 1].set_title('Consumption Time by Port')
+            axes[1, 0].set_title('Reward Rate by Block')
+            axes[1, 1].set_title('Percent Time Engaged by Block')
+            axes[2, 0].set_title('Premature leave from BG port by Block')
+
+            axes[0, 0].set_ylim([0, 20])
+            axes[0, 1].set_ylim([0, 20])
+            axes[1, 0].set_ylim([0, 0.65])
+            axes[1, 1].set_ylim([0, 1])
+            axes[2, 0].set_ylim([0, 1])
+
+            plt.suptitle(f"{mouse}", fontsize=20)
+            os.makedirs(save_folder, exist_ok=True)
+            filename = f'{mouse}_session_summary.png'
+            save_path = os.path.join(save_folder, filename)
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+            print(f"Graph saved to: {save_path}")
+            plt.show()
+
+            block_leaves_last10_df = block_leaves[(block_leaves.session >= block_leaves.session.max() - 10)]\
+                .groupby('block')['leave time'].mean().reset_index()
+            block_leaves_last10_df['animal'] = mouse
+            block_leaves_last10 = pd.concat([block_leaves_last10, block_leaves_last10_df])
+
+    fig, axes = plt.subplots(1, 1)
+    sns.boxplot(data=block_leaves_last10.reset_index(), x='block', y='leave time')
+    plt.show()
+
 if __name__ == '__main__':
-    # mice = ['ES057', 'ES058', 'ES059', 'ES060', 'ES061', 'ES062']
-    # mice = ['ES045', 'ES046', 'ES047', 'ES051', 'ES052', 'ES053', 'ES057', 'ES060', 'ES061', 'ES062']
-    # mice = ['ES058', 'ES059', 'ES045', 'ES047']
-    mice = ['ES057', 'ES046']
-    # mice = ['ES051', 'ES052', 'ES053', 'ES060', 'ES061', 'ES062']
-    simple_plots(mice)
-    # single_session(mice)
+    mice = ['RK007','RK008','RK009','RK010']
+    # mice = ['RK004']
+    # mice = ['RK001', 'RK002','RK003','RK004','RK005','RK006','ES039', 'ES045','ES046','ES047','ES051', 'ES052', 'ES053']
+    simple_plots2(mice) #history day by day summary statistics
+    single_session(mice)  # most recent session when num_back = 0
+
diff --git a/stand_alone/gui.py b/stand_alone/gui.py
index 5127c4e..d354113 100644
--- a/stand_alone/gui.py
+++ b/stand_alone/gui.py
@@ -12,20 +12,20 @@
 from user_settings import get_user_info
 import os
 
-info_dict = get_user_info(box=os.uname()[1])
+info_dict = get_user_info()
 
 # scp -r C:\Users\Elissa\GoogleDrive\Code\Python\behavior_code\stand_alone pi@elissapi0:\home\pi
 # scp C:\Users\Elissa\GoogleDrive\Code\Python\behavior_code\stand_alone\scp_rescue.py pi@elissapi1:\home\pi\behavior
 
 # scp -r "C:\Users\Shichen\OneDrive - Johns Hopkins\ShulerLab\behavior_code\stand_alone" pi@elissapi1:\home\pi\behavior1
 pastel_colors = ['#ffffcc', '#99ccff', '#cc99ff', '#ff99cc', '#ffcc99', '#ffffcc', '#99ffcc', '#ccffff', '#ccccff',
-                 '#ffccff', '#ffcccc', '#D3D3D3']
+                 '#ffccff', '#ffcccc', '#D3D3D3', '#f0a207']
 
 
 class Gui:
     def __init__(self):
         self.root = tk.Tk()
-        self.root.geometry('2000x700')
+        self.root.geometry('960x360')
         self.root.title('BehaviorGui')
 
         with open('durations.pkl', 'rb') as f:
@@ -37,9 +37,11 @@ def __init__(self):
         myFont = font.Font(size=30)
         mouse_rows = len(info_dict['mouse_buttons'])
         self.mouse_assignments = info_dict['mouse_assignments']
+        self.port_swap_assignments = info_dict['port_swap_assignments']
+        self.block_swap_assignments = info_dict['block_swap_assignments']
         tasks = {
             'single_reward': single_reward,
-            'cued_forgo': cued_forgo
+            'multi_reward': multi_reward
         }
         for key in self.mouse_assignments.keys():
             self.mouse_assignments[key] = tasks[self.mouse_assignments[key]]
@@ -105,8 +107,10 @@ def __init__(self):
 
     def run_behavior(self, mouse):
         task = self.mouse_assignments[mouse]
+        swap_port = self.port_swap_assignments[mouse]
+        swap_block = self.block_swap_assignments[mouse]
         print(f"running {task} for {mouse}")
-        main(mouse, task, forgo=False, forced_trials=True)
+        main(mouse, task, swap_port=swap_port, swap_block=swap_block)
 
     def calibrate(self, port):
         print(f'calibrating port {port}')
diff --git a/stand_alone/main.py b/stand_alone/main.py
index a5e0a45..5e523de 100644
--- a/stand_alone/main.py
+++ b/stand_alone/main.py
@@ -1,5 +1,6 @@
 from support_classes import *
-from tasks import *
+# from tasks import *
+from tasks_RK import multireward_task
 from timescapes import *
 import random
 
@@ -36,68 +37,76 @@
 #     session.start([task1])
 
 
-def cued_forgo(session, reward_level, starting_prob, session_length, forgo=False, forced_trials=True):
+def multi_reward(session, reward_level, starting_prob, session_length, swap_port=False, swap_block = False):
     print(reward_level)
     print(starting_prob)
     print(session_length)
-    task_structure = cued_forgo_task
-
-    if forgo:
-        task_name = 'cued_forgo'
-        print('cued forgo')
+    task_structure = multireward_task
+
+    # rates = [.4, .8, .4, .8, .4, .8]
+    # if np.random.random() > .5:
+    #     rates.reverse()
+
+    # rates = [.4, .8] #swap_block random is session by session change, swap block trial is trial by trial change
+    if swap_block == "random":
+        rates = [random.choice([0.4, 0.8])]
+        print (rates)
+    elif swap_block == "trial":
+        rates = [.4, .8, .4, .8, .4, .8]
+        if np.random.random() > .5:
+            rates.reverse()
     else:
-        task_name = 'cued_no_forgo'
-        print('cued bg without forgo option')
-
-    if forced_trials:
-        task_name = task_name + '_forced'
-        print('forced trials included')
+        rates = [0.4] if not swap_block else [0.8] #0.4 if swap_block is false, 0.8 if true
+    #     rates.reverse()
 
-    rates = [.4, .8, .4, .8, .4, .8]
-    if np.random.random() > .5:
-        rates.reverse()
+    exp_port = 1 if not swap_port else 2
+    bg_port = 2 if not swap_port else 1
+    print(f'exp_port = {exp_port}')
+    print(f'bg_port = {bg_port}')
     background_dist = {'distribution': 'background',
                        'rates': rates,
                        'duration': 5,
-                       'port_num': 2}
+                       'port_num': bg_port}
     print(background_dist['rates'])
     exp_dist = {'distribution': exp_decreasing,
                 'cumulative': reward_level,
                 'starting_probability': starting_prob,
-                'port_num': 1}
+                'port_num': exp_port}
     ports = {'exp': Port(exp_dist['port_num'], dist_info=exp_dist),
              'background': Port(background_dist['port_num'], dist_info=background_dist)}
-    task1 = Task(session, name=task_name, structure=task_structure, ports=ports,
-                 maximum=session_length, limit='time', forgo=forgo, forced_trials=forced_trials)
+    task1 = Task(session, name='multi_reward', structure=task_structure, ports=ports,
+                 maximum=session_length, limit='time')
     session.start([task1])
 
 
-def single_reward(session, reward_level, starting_prob, session_length, forgo=False, forced_trials=True):
+def single_reward(session, reward_level, starting_prob, session_length, swap_port=False, swap_block= False):
     reward_level = 0.5994974874371859  # cumulative for an 8 reward version
     starting_prob = 0.1301005025125628
     print(reward_level)
     print(starting_prob)
     print(session_length)
     task_structure = single_reward_task
-
+    
     task_name = 'single_reward'
-
+    
     rates = [.4, .8, .4, .8, .4, .8]
     if np.random.random() > .5:
         rates.reverse()
+    exp_port = 1 if not swap_port else 2
+    bg_port = 2 if not swap_port else 1
     background_dist = {'distribution': 'background',
-                       'rates': rates,
-                       'duration': 5,
-                       'port_num': 2}
+                   'rates': rates,
+                   'duration': 5,
+                   'port_num': bg_port}
     print(background_dist['rates'])
     exp_dist = {'distribution': exp_decreasing,
-                'cumulative': reward_level,
-                'starting_probability': starting_prob,
-                'port_num': 1}
+            'cumulative': reward_level,
+            'starting_probability': starting_prob,
+            'port_num': exp_port}
     ports = {'exp': Port(exp_dist['port_num'], dist_info=exp_dist),
-             'background': Port(background_dist['port_num'], dist_info=background_dist)}
+         'background': Port(background_dist['port_num'], dist_info=background_dist)}
     task1 = Task(session, name=task_name, structure=task_structure, ports=ports,
-                 maximum=session_length, limit='time')
+             maximum=session_length, limit='time')
     session.start([task1])
 
 # def give_up_blocked(session, reward_level, starting_prob, session_length, forgo=True, forced_trials=False):
@@ -122,7 +131,8 @@ def single_reward(session, reward_level, starting_prob, session_length, forgo=Fa
 #     session.start([task1])
 
 
-def main(mouse, to_run, forgo=False, forced_trials=False):
+# def main(mouse, to_run, forgo=False, forced_trials=False, swap_port=False):
+def main(mouse, to_run, swap_port=False, swap_block = False): #swap port swaps BG/EXP, swap_block is the stage
     cumulative = 8
     start_prob = 1
     session_time = 18
@@ -135,9 +145,9 @@ def main(mouse, to_run, forgo=False, forced_trials=False):
 
     try:
         if mouse not in mouse_settings.keys():
-            to_run(session, *mouse_settings['default'], forgo=forgo, forced_trials=forced_trials)  # Run the task
+            to_run(session, *mouse_settings['default'], swap_port=swap_port, swap_block=swap_block)  # Run the task
         else:
-            to_run(session, *mouse_settings[mouse], forgo=forgo, forced_trials=forced_trials)  # Run the task
+            to_run(session, *mouse_settings[mouse], swap_port=swap_port, swap_block=swap_block)  # Run the task
         session.smooth_finish = True
         print('smooth finish')
     except KeyboardInterrupt:  # Catch if the task is stopped via ctrl-C or the stop button
@@ -149,6 +159,6 @@ def main(mouse, to_run, forgo=False, forced_trials=False):
 if __name__ == "__main__":
     # main('testmouse', cued_forgo, forgo=False, forced_trials=True)
     # main('ES024', cued_forgo, forgo=False, forced_trials=True)
-    main('ES030', cued_forgo, forgo=False, forced_trials=True)
+    main('ES030', multi_reward)
     # main('testmouse', give_up_blocked)
 
diff --git a/stand_alone/support_classes.py b/stand_alone/support_classes.py
index b5ad92a..427c4d5 100644
--- a/stand_alone/support_classes.py
+++ b/stand_alone/support_classes.py
@@ -42,7 +42,72 @@ def test(i):
     print('test ' + str(i))
 
 
-def ssh(host, cmd, user, password, timeout=30, bg_run=False):
+# def ssh(host, cmd, user, password, timeout=30, bg_run=False):
+#     """SSH'es to a host using the supplied credentials and executes a command.
+#     Throws an exception if the command doesn't return 0.
+#     bgrun: run command in the background"""
+#
+#     options = '-q -oStrictHostKeyChecking=no -oUserKnownHostsFile=/dev/null -oPubkeyAuthentication=no'
+#     if bg_run:
+#         options += ' -f'
+#
+#     ssh_cmd = 'ssh %s@%s %s \'%s\'' % (user, host, options, cmd)
+#     print(ssh_cmd)
+#     child = pexpect.spawnu(ssh_cmd, timeout=timeout)
+#     child.expect([r'(?i).*password.*:'])
+#     child.sendline(password)
+#     child.expect(pexpect.EOF)
+#     child.close()
+
+
+# def scp(host, filename, destination, user, password, timeout=30, bg_run=False, recursive=False, cmd=False):
+#     """Scp's to a host using the supplied credentials and executes a command.
+#     Throws an exception if the command doesn't return 0.
+#     bgrun: run command in the background"""
+#
+#     options = '-q -oStrictHostKeyChecking=no -oUserKnownHostsFile=/dev/null -oPubkeyAuthentication=no'
+#     if recursive:
+#         options += ' -r'
+#     if bg_run:
+#         options += ' -f'
+#     scp_cmd = 'scp %s %s %s@%s:\'"%s"\'' % (options, filename, user, host, os.path.join(destination, filename))
+#     password_prompt = rf"{user}@{host}'s password: "
+#     print(password_prompt)
+#     child = None
+#     status = 1
+#     print(scp_cmd)
+#     try:
+#         if cmd:
+#             return scp_cmd
+#         child = pexpect.spawnu(scp_cmd, timeout=timeout)  # spawnu for Python 3
+#         child.expect(password_prompt)
+#         child.sendline(password)
+#         child.expect(pexpect.EOF)
+#         child.close()
+#         status = child.exitstatus
+#     except Exception as e:
+#         print('scp didn\'t work the first time so we are trying again with different password')
+#         try:
+#             # scp_cmd2 = 'scp %s %s %s@%s:"%s"' % (options, filename, user, host, os.path.join(destination, filename))
+#             # print(scp_cmd2)
+#             # if cmd:
+#             #     return scp_cmd2
+#             child = pexpect.spawnu(scp_cmd, timeout=timeout)  # spawnu for Python 3
+#             child.expect([r"(?i).*password.*:?(\s*)"])
+#             child.sendline(password)
+#             child.expect(pexpect.EOF)
+#             child.close()
+#             status = child.exitstatus
+#         except Exception as e_2:
+#             print(f"still didnt work: {e_2}")
+#             status = 1
+#     if status == 0:
+#         print("file transfer completed!")
+#     else:
+#         print("transfer failed.")
+#     return status
+
+def ssh(host, cmd, user, password, timeout=60, bg_run=False):
     """SSH'es to a host using the supplied credentials and executes a command.
     Throws an exception if the command doesn't return 0.
     bgrun: run command in the background"""
@@ -54,12 +119,10 @@ def ssh(host, cmd, user, password, timeout=30, bg_run=False):
     ssh_cmd = 'ssh %s@%s %s \'%s\'' % (user, host, options, cmd)
     print(ssh_cmd)
     child = pexpect.spawnu(ssh_cmd, timeout=timeout)
-    child.expect(['[Pp]assword: '])
+    child.expect([r"(?i).*password.*:?(\s*)"])
     child.sendline(password)
     child.expect(pexpect.EOF)
     child.close()
-
-
 def scp(host, filename, destination, user, password, timeout=30, bg_run=False, recursive=False, cmd=False):
     """Scp's to a host using the supplied credentials and executes a command.
     Throws an exception if the command doesn't return 0.
@@ -96,7 +159,6 @@ def scp(host, filename, destination, user, password, timeout=30, bg_run=False, r
             print('worked!')
     return child.exitstatus
 
-
 def sync_stream(self):
     pin_map = {'session': 25,
                'trial': 8,
@@ -142,7 +204,7 @@ def __init__(self, mouse):
         self.data_write_path = "/data/" + self.mouse
         self.datetime = time.strftime("%Y-%m-%d_%H-%M-%S")
         self.filename = "data_" + self.datetime + ".txt"
-
+        self.isstarted = False
         self.halted = False
         self.smooth_finish = False
         GPIO.setmode(GPIO.BCM)
@@ -167,6 +229,7 @@ def __init__(self, mouse):
         GPIO.output(self.camera_pin, GPIO.LOW)
 
     def start(self, task_list):
+        self.isstarted = True
         for val in self.sync_pins.values():
             GPIO.setup(val, GPIO.OUT)
             GPIO.output(val, GPIO.LOW)
@@ -190,10 +253,11 @@ def start(self, task_list):
             perform(task)
 
     def log(self, string):  # Adds session time stamp to beginning of string and logs it
-        session_time = time.time() - self.start_time
-        new_line = str(session_time) + ',' + string + '\n'
-        # print(new_line)
-        self.f.write(new_line)
+        if self.isstarted:
+            session_time = time.time() - self.start_time
+            new_line = str(session_time) + ',' + string + '\n'
+            # print(new_line)
+            self.f.write(new_line)
 
     def end(self):
         self.log('nan,nan,nan,nan,setup,nan,0,session')
@@ -201,16 +265,21 @@ def end(self):
         os.system('sudo chmod o-w ' + self.filename)
         mkdir_command = 'if not exist "%s" mkdir "%s"' % (
             self.ssh_path.replace('/', '\\'), self.ssh_path.replace('/', '\\'))
-        ssh(self.ip, mkdir_command, self.user, self.password)
-        res = scp(self.ip, self.filename, self.data_send_path, self.user, self.password)
-        if not res:
-            print('\nSuccessful file transfer to "%s"\nDeleting local file from pi.' % self.data_send_path)
-            os.remove(self.filename)
-        else:
-            print('connection back to desktop timed out')
-        GPIO.cleanup()
-        os.chdir(os.path.join(os.getcwd(), '..', '..'))
-        print('\nFile closed and clean up complete')
+        try:
+            ssh(self.ip, mkdir_command, self.user, self.password)
+            res = scp(self.ip, self.filename, self.data_send_path, self.user, self.password)
+            if res == 0:
+                print('\nSuccessful file transfer to "%s"\nDeleting local file from pi.' % self.data_send_path)
+                os.remove(self.filename)
+        except Exception as e:
+            res = 1
+            print('scp failed with the following error')
+            print(e)
+        finally:
+            os.chdir(os.path.join(os.getcwd(), '..', '..'))
+            print('file location is back to the project root')
+            GPIO.cleanup()
+            print('\nFile closed and clean up complete')
         if self.halted:
             print('Session stopped early')
         elif self.smooth_finish:
@@ -323,7 +392,7 @@ def lick_status_change(self):
 
 class Task:
     def __init__(self, session, name='blank', structure=None, ports=None, limit='trials',
-                 maximum=None, forgo=True, forced_trials=False):
+                 maximum=None, forgo=False, forced_trials=True):
         print('Starting task: %s' % name)
         self.structure = structure
         self.port_dict = ports
diff --git a/stand_alone/user_settings_generic.py b/stand_alone/user_settings_generic.py
index fb41112..f5ef573 100644
--- a/stand_alone/user_settings_generic.py
+++ b/stand_alone/user_settings_generic.py
@@ -16,6 +16,10 @@ def get_user_info():
             'MOUSE_NAME': 'TASK_NAME',
             'testmouse': 'TASK_NAME',
         },
+        'mouse_swap_assignments': {
+            'MOUSE_NAME': False,
+            'testmouse': False,
+        },
         'mouse_colors': {
             'MOUSE_NAME': 1,
             'testmouse': 1,
diff --git a/upload_cloud_backup.py b/upload_cloud_backup.py
new file mode 100644
index 0000000..f9ca9da
--- /dev/null
+++ b/upload_cloud_backup.py
@@ -0,0 +1,40 @@
+import os
+import shutil
+import time
+from pathlib import Path
+
+# Define local and OneDrive paths
+behavior_data_dir = Path("C:/Users/Shichen/OneDrive - Johns Hopkins/ShulerLab/Rie_behavior/data")
+# onedrive_dir = Path("C:/Users/Valued Customer/OneDrive - Johns Hopkins/behavior_data")
+# onedrive_dir = Path("C:/Users/Rie/OneDrive - Johns Hopkins/behavior_data")
+
+
+def upload_cloud_backup():
+    while True:
+        try:
+            # Walk through the behavior_data directory and identify files and folders
+            for root, dirs, files in os.walk(behavior_data_dir):
+                # Calculate the relative path from the base behavior_data directory
+                relative_path = Path(root).relative_to(behavior_data_dir)
+                # Create the corresponding subdirectory path in the onedrive backup folder
+                backup_subdir = onedrive_dir / relative_path
+                backup_subdir.mkdir(parents=True, exist_ok=True)  # Create directories if needed
+
+                # Copy each file if it's not already in the backup folder
+                for file_name in files:
+                    src_file = Path(root) / file_name
+                    dest_file = backup_subdir / file_name
+
+                    # Only copy if the file doesn't exist in the destination
+                    if not dest_file.exists():
+                        shutil.copy2(src_file, dest_file)
+                        print(f"Copied {src_file} to {dest_file}")
+
+        except Exception as e:
+            print(f"Error during backup: {e}")
+
+        # Wait 24 hours before the next check
+        time.sleep(3600 * 24)
+
+if __name__ == '__main__':
+    upload_cloud_backup()

From fb37f09a509b8951f3a51713f6116fa7d7155289 Mon Sep 17 00:00:00 2001
From: riekane <rkaneko2@jh.edu>
Date: Fri, 23 May 2025 14:16:33 -0400
Subject: [PATCH 2/3] Rie's edits for swapping port assignments and minor fixes
 to the task code.

---
 simple_plots.py                | 774 ++++++++++++++++++++++++++-------
 stand_alone/gui.py             |  14 +-
 stand_alone/main.py            |  80 ++--
 stand_alone/support_classes.py | 111 ++++-
 upload_to_pi.py                |  12 +-
 5 files changed, 762 insertions(+), 229 deletions(-)

diff --git a/simple_plots.py b/simple_plots.py
index de1dbb3..4add81e 100644
--- a/simple_plots.py
+++ b/simple_plots.py
@@ -10,24 +10,31 @@
 import seaborn as sns
 from matplotlib.collections import PatchCollection
 from matplotlib.patches import Rectangle
+import matplotlib.gridspec as gridspec
+
+from upload_cloud_backup import behavior_data_dir
 from user_info import get_user_info
 import shutil
+from datetime import datetime
+import pdb
 
 info_dict = get_user_info()
 initials = info_dict['initials']
 start_date = info_dict['start_date']
+data_dir = behavior_data_dir
 
 
 def get_today_filepaths(days_back=0):
     file_paths = []
-    for root, dirs, filenames in walk(os.path.join(os.getcwd(), 'data')):
-        if len(dirs) == 0 and os.path.basename(root)[:2] == initials:
+    for root, dirs, filenames in walk(data_dir):
+        if len(dirs) == 0 and os.path.basename(root)[:2] in initials:
             mouse = os.path.basename(root)
             for f in filenames:
                 if f == 'desktop.ini':
                     continue
-                file_date = date(int(f[5:9]), int(f[10:12]), int(f[13:15]))
-                dif = date.today() - file_date
+                file_date = datetime.strptime(f[5:-4], '%Y-%m-%d_%H-%M-%S')
+                # file_date = date(int(f[5:9]), int(f[10:12]), int(f[13:15]))
+                dif = datetime.today() - file_date
                 if dif.days <= days_back:
                     # if f[5:15] == time.strftime("%Y-%m-%d"):
                     file_paths.append(os.path.join(mouse, f))
@@ -91,24 +98,39 @@ def gen_data(file_paths, select_mouse=None, return_info=False):
         if select_mouse is not None and mouse not in select_mouse:
             continue
 
-        path = os.path.join(os.getcwd(), 'data', f)
+        path = os.path.join(data_dir, f)
+        meta_data = read_pi_meta(path)
+
         if return_info:
-            data = read_pi_meta(path)
-            # if data['box'] == 'elissapi0':
-            #     session = pd.read_csv(path, na_values=['None'], skiprows=3)
-            #     session_summary(data_reduction(session), mouse)
-            #     ans = input(f'remove broken file? (y/n)\n{path}\n???')
-            #     if ans == 'y':
-            #         file_name = f[6:]
-            #         half_session_path = os.path.join(os.getcwd(), 'data', 'half_sessions', file_name)
-            #         shutil.move(path, half_session_path)
+            data = meta_data
         else:
-            data = pd.read_csv(path, na_values=['None'], skiprows=3)
+            try:
+                data = pd.read_csv(path, na_values=['None'], skiprows=3)
+            except pd.errors.EmptyDataError:
+                print(f'empty file at {path}')
+                continue
             try:
                 data = data_reduction(data)
+
+                #port info as a row
+                port_info_row = [{
+                    'key': meta_data['port1_info']['distribution'],
+                    'port': meta_data['port1_info']['port_num']
+                    # 'phase': last_row_phase,
+                },
+                    {'key': meta_data['port2_info']['distribution'],
+                     'port': meta_data['port2_info']['port_num']
+                     # 'phase': last_row_phase,
+                     }
+                ]
+
+                #Added port info at the end of the dataframe
+                portinfo_rows_df = pd.DataFrame(port_info_row).reindex(columns=data.columns)
+                data = pd.concat([data, portinfo_rows_df], ignore_index=True)
+
             except ValueError:
                 file_name = f[6:]
-                half_session_path = os.path.join(os.getcwd(), 'data', 'half_sessions', file_name)
+                half_session_path = os.path.join(data_dir, 'half_sessions', file_name)
                 if data.session_time.max() < 800:
                     print(f'moving {f} to half sessions, session time: {data.session_time.max():.2f} seconds')
                     shutil.move(path, half_session_path)
@@ -117,6 +139,7 @@ def gen_data(file_paths, select_mouse=None, return_info=False):
                     if ans == 'y':
                         shutil.move(path, half_session_path)
                 continue
+
         if mouse in d.keys():
             d[mouse].append(data)
         else:
@@ -158,9 +181,13 @@ def data_reduction(df, lick_tol=.01, head_tol=.2):
 
 
 def consumption_time(df):
-    bg_end_times = df[(df.key == 'LED') & (df.port == 2) & (df.value == 1)]
-    exp_entries = df[(df.key == 'head') & (df.port == 1) & (df.value == 1)]
+    bgportassignment = df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    expportassignment = df.loc[df['key'] == 'exp_decreasing', 'port'].iloc[-1]
+    bg_end_times = df[(df.key == 'LED') & (df.port == bgportassignment) & (df.value == 1)]
+    exp_entries = df[(df.key == 'head') & (df.port == expportassignment) & (df.value == 1)]
+
     dif = min_dif(bg_end_times.session_time, exp_entries.session_time)
+
     bg_consumption = dif[~np.isnan(dif)]
     if df.task.iloc[10] != 'single_reward':
         consumption_df = pd.DataFrame()
@@ -168,8 +195,8 @@ def consumption_time(df):
         consumption_df['port'] = ['bg'] * len(bg_consumption)
         return consumption_df
 
-    exp_end_times = df[(df.key == 'LED') & (df.port == 1) & (df.value == 1)]
-    bg_entries = df[(df.key == 'head') & (df.port == 2) & (df.value == 1)]
+    exp_end_times = df[(df.key == 'LED') & (df.port == expportassignment) & (df.value == 1)]
+    bg_entries = df[(df.key == 'head') & (df.port == bgportassignment) & (df.value == 1)]
     dif = min_dif(exp_end_times.session_time, bg_entries.session_time)
     exp_consumption = dif[~np.isnan(dif)]
     consumption_df = pd.DataFrame()
@@ -178,12 +205,58 @@ def consumption_time(df):
     return consumption_df
 
 
+def calculate_premature_leave(df, threshold=1.0): #trials with premature leave
+    bgportassignment = df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    premature_leave_trials = 0  # Counter for trials with premature leave
+    blocks = df.phase.dropna().unique()
+    blocks.sort()
+    results = []
+    for block in blocks:
+        block_data = df[df.phase == block]
+        block_total_trials = len(block_data.trial.unique())
+        for trial in block_data.trial.unique():
+            if pd.isna(trial):
+                continue  # Skip invalid trials
+            trial_data = df[df.trial == trial] # Filter data for the current trial
+            bg_on_times = trial_data[(trial_data['key'] == 'trial') &
+                                      (trial_data['value'] == 1)].session_time.values
+            if len(bg_on_times) == 0:
+                continue
+            bg_start = bg_on_times[0]
+            bg_off_times = trial_data[(trial_data['key'] == 'LED') &
+                                      (trial_data['port'] == bgportassignment) &
+                                      (trial_data['value'] == 1)].session_time.values
+            if len(bg_off_times) == 0:  # Check if no BG ON recorded
+                continue
+            bg_end = bg_off_times[0]
+            bg_head_out_times = trial_data[(trial_data['key'] == 'head') & # Extract head-out and head-in times while BG port is active
+                                        (trial_data['port'] == bgportassignment) &
+                                        (trial_data['value'] == 0) &
+                                        (trial_data['session_time'] >= bg_start) &
+                                        (trial_data['session_time'] < bg_end)].session_time.values
+            bg_head_in_times = trial_data[(trial_data['key'] == 'head') &
+                                       (trial_data['port'] == bgportassignment) &
+                                       (trial_data['value'] == 1) &
+                                       (trial_data['session_time'] >= bg_start) &
+                                       (trial_data['session_time'] < bg_end)].session_time.values
+            for head_out in bg_head_out_times: # Check for premature leave: head-out without a valid head-in within threshold
+                if not any((head_in > head_out) and (head_in - head_out) <= threshold for head_in in bg_head_in_times):
+                    premature_leave_trials += 1  # Flag this trial as a premature leave
+                    break  # Stop checking further head-outs in this trial
+        premature_leave_rate = (premature_leave_trials / block_total_trials)
+        results.append({"block": block,"premature leave numbers":premature_leave_trials, "premature leave rate": premature_leave_rate})
+    premature_leave_df = pd.DataFrame(results)
+    return premature_leave_df
+
+
 def block_leave_times(df):
+    bgportassignment = df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    expportassignment = df.loc[df['key'] == 'exp_decreasing', 'port'].iloc[-1]
     reward_trials = df[(df.key == 'reward_initiate')].trial.to_numpy()
     non_reward = ~df.trial.isin(reward_trials)
-    bg_end_times = df[(df.key == 'LED') & (df.port == 2) & (df.value == 1) & non_reward]
-    exp_entries = df[(df.key == 'head') & (df.value == 1) & (df.port == 1) & non_reward]
-    exp_exits = df[(df.key == 'head') & (df.value == 0) & (df.port == 1) & non_reward]
+    bg_end_times = df[(df.key == 'LED') & (df.port == bgportassignment) & (df.value == 1) & non_reward]
+    exp_entries = df[(df.key == 'head') & (df.value == 1) & (df.port == expportassignment) & non_reward]
+    exp_exits = df[(df.key == 'head') & (df.value == 0) & (df.port == expportassignment) & non_reward]
     bg_end_times = bg_end_times[bg_end_times.session_time < exp_entries.session_time.max()]
     ind, dif = min_dif(bg_end_times.session_time, exp_entries.session_time, return_index=True)
     exp_entries = exp_entries.iloc[np.unique(ind)]
@@ -193,8 +266,14 @@ def block_leave_times(df):
     valid_trials = np.intersect1d(valid_trials, bg_end_times.trial.values)
     exp_exits = exp_exits.loc[valid_trials]
     exp_entries = exp_entries.loc[valid_trials]
+
     if len(exp_exits.to_numpy()) != len(exp_entries.to_numpy()):
-        print()
+        print(
+            f"Mismatch in exp_entries and exp_exits. "
+            f"exp_entries: {exp_entries}, exp_exits: {exp_exits}. "
+            "Using clean_entries_exits to resolve."
+        )
+        exp_entries, exp_exits = clean_entries_exits(exp_entries, exp_exits)
     leave_times = exp_exits.to_numpy() - exp_entries.to_numpy()
 
     trial_blocks = bg_end_times[bg_end_times.trial.isin(exp_entries.index.values)].phase.to_numpy()
@@ -208,15 +287,17 @@ def get_entry_exit(df, trial):
     is_trial = df.trial == trial
     start = df.value == 1
     end = df.value == 0
-    port1 = df.port == 1
-    port2 = df.port == 2
+    bgport = df.port == df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    expport = df.port == df.loc[df['key'] == 'exp_decreasing', 'port'].iloc[-1]
+    # port1 = df.port == 1
+    # port2 = df.port == 2
 
     trial_start = df[is_trial & start & (df.key == 'trial')].session_time.values[0]
-    trial_middle = df[is_trial & end & (df.key == 'LED') & port2].session_time.values[0]
+    trial_middle = df[is_trial & end & (df.key == 'LED') & bgport].session_time.values[0] #head in to EXP, bg LED off
     trial_end = df[is_trial & end & (df.key == 'trial')].session_time.values[0]
 
-    bg_entries = df[is_trial & port2 & start & (df.key == 'head')].session_time.to_numpy()
-    bg_exits = df[is_trial & port2 & end & (df.key == 'head')].session_time.to_numpy()
+    bg_entries = df[is_trial & bgport & start & (df.key == 'head')].session_time.to_numpy()
+    bg_exits = df[is_trial & bgport & end & (df.key == 'head')].session_time.to_numpy()
 
     if len(bg_entries) == 0 or len(bg_exits) == 0 or bg_entries[0] > bg_exits[0]:
         bg_entries = np.concatenate([[trial_start], bg_entries])
@@ -225,15 +306,15 @@ def get_entry_exit(df, trial):
     if len(bg_exits) == 0 or bg_entries[-1] > bg_exits[-1]:
         bg_exits = np.concatenate([bg_exits, [trial_middle]])
 
-    exp_entries = df[is_trial & port1 & start & (df.key == 'head') &
+    exp_entries = df[is_trial & expport & start & (df.key == 'head') &
                      (df.session_time > trial_middle)].session_time.to_numpy()
-    exp_exits = df[is_trial & port1 & end & (df.key == 'head') &
+    exp_exits = df[is_trial & expport & end & (df.key == 'head') &
                    (df.session_time > trial_middle)].session_time.to_numpy()
 
     if not (len(exp_entries) == 0 and len(exp_exits) == 0):
-        if len(exp_entries) == 0:
+        if len(exp_entries) == 0: #only exp out
             exp_entries = np.concatenate([[trial_middle], exp_entries])
-        if len(exp_exits) == 0:
+        if len(exp_exits) == 0: #only exp in
             exp_exits = np.concatenate([exp_exits, [trial_end]])
 
         if exp_entries[0] > exp_exits[0]:
@@ -241,12 +322,12 @@ def get_entry_exit(df, trial):
         if exp_entries[-1] > exp_exits[-1]:
             exp_exits = np.concatenate([exp_exits, [trial_end]])
 
-    early_exp_entries = df[is_trial & port1 & start & (df.key == 'head') &
+    early_exp_entries = df[is_trial & expport & start & (df.key == 'head') &
                            (df.session_time < trial_middle)].session_time.to_numpy()
-    early_exp_exits = df[is_trial & port1 & end & (df.key == 'head') &
+    early_exp_exits = df[is_trial & expport & end & (df.key == 'head') &
                          (df.session_time < trial_middle)].session_time.to_numpy()
 
-    if not (len(early_exp_entries) == 0 and len(early_exp_exits) == 0):
+    if not (len(early_exp_entries) == 0 and len(early_exp_exits) == 0): #any early exp in/out
         if len(early_exp_entries) == 0:
             early_exp_entries = np.concatenate([[trial_start], early_exp_entries])
         if len(early_exp_exits) == 0:
@@ -264,114 +345,164 @@ def get_entry_exit(df, trial):
     if len(early_exp_entries) != len(early_exp_exits):
         print()
 
+    if len(exp_entries):
+        if len(exp_exits) != len(exp_entries):
+            print(
+                f"Mismatch in exp_entries and exp_exits in trial {trial}. "
+                f"exp_entries: {exp_entries}, exp_exits: {exp_exits}. "
+                "Using clean_entries_exits to resolve."
+            )
+            exp_entries, exp_exits = clean_entries_exits(exp_entries, exp_exits)
+
+    if len(bg_entries) != len(bg_exits):
+        print(
+            f"Mismatch in bg_entries and bg_exits in trial {trial}. "
+            f"bg_entries: {bg_entries}, bg_exits: {bg_exits}. "
+            "Using clean_entries_exits to resolve."
+        )
+        bg_entries, bg_exits = clean_entries_exits(bg_entries, bg_exits)
+
     return bg_entries, bg_exits, exp_entries, exp_exits, early_exp_entries, early_exp_exits
 
 
+def clean_entries_exits(entries, exits):
+    """
+    Cleans mismatched entries and exits such that each entry is paired with the nearest valid exit.
+    """
+    valid_entries = []
+    valid_exits = []
+    e_idx, x_idx = 0, 0
+
+    while e_idx < len(entries) and x_idx < len(exits):
+        if entries[e_idx] < exits[x_idx]:  # Valid entry-exit pair
+            valid_entries.append(entries[e_idx])
+            valid_exits.append(exits[x_idx])
+            e_idx += 1
+            x_idx += 1  # Move to the next entry and exit
+        else:
+            x_idx += 1  # Skip unmatched exits
+
+    return valid_entries, valid_exits
+
 def percent_engaged(df):
-    travel_time = .5
-    blocks = df.phase.unique()
-    blocks.sort()
-    time_engaged = []
-    block_time = []
-    block_rewards = []
-    for block in blocks:
-        engaged = []
-        all_time = []
-        rewards = []
-        block_trials = df[(df.value == 0) & (df.key == 'trial') & (df.phase == block)].trial
-        for trial in block_trials:
-            bg_entries, bg_exits, exp_entries, exp_exits, _, _ = get_entry_exit(df, trial)
-            is_trial = df.trial == trial
-            start = df.value == 1
-            end = df.value == 0
-            # port1 = df.port == 1
-            # port2 = df.port == 2
-
-            #
-            trial_start = df[is_trial & start & (df.key == 'trial')].session_time.values[0]
-            # trial_middle = df[is_trial & start & (df.key == 'LED') & port2].session_time.values[0]
-            trial_end = df[is_trial & end & (df.key == 'trial')].session_time.values[0]
-            #
-            # bg_entries = df[is_trial & port2 & start & (df.key == 'head')].session_time.to_numpy()
-            # bg_exits = df[is_trial & port2 & end & (df.key == 'head')].session_time.to_numpy()
-            #
-            # if len(bg_entries) == 0 or bg_entries[0] > bg_exits[0]:
-            #     bg_entries = np.concatenate([[trial_start], bg_entries])
-            # if trial_end - bg_entries[-1] < .1:
-            #     bg_entries = bg_entries[:-1]
-            # if len(bg_exits) == 0 or bg_entries[-1] > bg_exits[-1]:
-            #     bg_entries = np.concatenate([bg_exits, [trial_middle]])
-            #
-            # if not (len(bg_entries) == len(bg_exits) and np.all(bg_exits - bg_entries > 0)):
-            #     print('stop')
-            # bg_engaged = sum(bg_exits - bg_entries)
-            #
-            # exp_entries = df[is_trial & port1 & start & (df.key == 'head') &
-            #                  (df.session_time > trial_middle)].session_time.to_numpy()
-            # exp_exits = df[is_trial & port1 & end & (df.key == 'head') &
-            #                (df.session_time > trial_middle)].session_time.to_numpy()
-            #
-            # if len(exp_entries) == 0 and len(exp_exits) == 0:
-            #     exp_engaged = 0
-            # else:
-            #     if len(exp_entries) == 0:
-            #         exp_entries = np.concatenate([[trial_middle], exp_entries])
-            #     if len(exp_exits) == 0:
-            #         exp_exits = np.concatenate([exp_exits, [trial_end]])
-            #
-            #     if exp_entries[0] > exp_exits[0]:
-            #         exp_entries = np.concatenate([[trial_middle], exp_entries])
-            #     if exp_entries[-1] > exp_exits[-1]:
-            #         exp_exits = np.concatenate([exp_exits, [trial_end]])
-            #     exp_engaged = sum(exp_exits - exp_entries)
-            #
-            #     # if not len(exp_entries) == len(exp_exits) and len(exp_entries):
-            #     #     print('stop')
-            #     # if len(exp_entries):
-
-            if len(exp_entries):
-                exp_engaged = sum(exp_exits - exp_entries)
-            else:
-                exp_engaged = 0
-            bg_engaged = sum(bg_exits - bg_entries)
+    try:
+        travel_time = .5
+        blocks = df.phase.dropna().unique()
+        blocks.sort()
+        time_engaged = []
+        block_time = []
+        block_rewards = []
+
+        for block in blocks:
+            engaged = []
+            all_time = []
+            rewards = []
+            block_trials = df[(df.value == 0) & (df.key == 'trial') & (df.phase == block)].trial
+            for trial in block_trials:
+                bg_entries, bg_exits, exp_entries, exp_exits, _, _ = get_entry_exit(df, trial)
+                is_trial = df.trial == trial
+                start = df.value == 1
+                end = df.value == 0
+                # port1 = df.port == 1
+                # port2 = df.port == 2
+
+                #
+                trial_start = df[is_trial & start & (df.key == 'trial')].session_time.values[0]
+                # trial_middle = df[is_trial & start & (df.key == 'LED') & port2].session_time.values[0]
+                trial_end = df[is_trial & end & (df.key == 'trial')].session_time.values[0]
+                #
+                # bg_entries = df[is_trial & port2 & start & (df.key == 'head')].session_time.to_numpy()
+                # bg_exits = df[is_trial & port2 & end & (df.key == 'head')].session_time.to_numpy()
+                #
+                # if len(bg_entries) == 0 or bg_entries[0] > bg_exits[0]:
+                #     bg_entries = np.concatenate([[trial_start], bg_entries])
+                # if trial_end - bg_entries[-1] < .1:
+                #     bg_entries = bg_entries[:-1]
+                # if len(bg_exits) == 0 or bg_entries[-1] > bg_exits[-1]:
+                #     bg_entries = np.concatenate([bg_exits, [trial_middle]])
+                #
+                # if not (len(bg_entries) == len(bg_exits) and np.all(bg_exits - bg_entries > 0)):
+                #     print('stop')
+                # bg_engaged = sum(bg_exits - bg_entries)
+                #
+                # exp_entries = df[is_trial & port1 & start & (df.key == 'head') &
+                #                  (df.session_time > trial_middle)].session_time.to_numpy()
+                # exp_exits = df[is_trial & port1 & end & (df.key == 'head') &
+                #                (df.session_time > trial_middle)].session_time.to_numpy()
+                #
+                # if len(exp_entries) == 0 and len(exp_exits) == 0:
+                #     exp_engaged = 0
+                # else:
+                #     if len(exp_entries) == 0:
+                #         exp_entries = np.concatenate([[trial_middle], exp_entries])
+                #     if len(exp_exits) == 0:
+                #         exp_exits = np.concatenate([exp_exits, [trial_end]])
+                #
+                #     if exp_entries[0] > exp_exits[0]:
+                #         exp_entries = np.concatenate([[trial_middle], exp_entries])
+                #     if exp_entries[-1] > exp_exits[-1]:
+                #         exp_exits = np.concatenate([exp_exits, [trial_end]])
+                #     exp_engaged = sum(exp_exits - exp_entries)
+                #
+                #     # if not len(exp_entries) == len(exp_exits) and len(exp_entries):
+                #     #     print('stop')
+                #     # if len(exp_entries):
+
+                if len(exp_entries):
+                    exp_engaged = sum([exit - entry for entry, exit in zip(exp_entries, exp_exits)])
+                else:
+                    exp_engaged = 0
+
+                bg_engaged = sum([exit - entry for entry, exit in zip(bg_entries, bg_exits)])
 
-            all_time.append(trial_end - trial_start)
-            engaged.append(bg_engaged + exp_engaged)
-            rewards.append(len(df[is_trial & start & (df.key == 'reward')]))
+                all_time.append(trial_end - trial_start)
+                engaged.append(bg_engaged + exp_engaged)
+                rewards.append(len(df[is_trial & start & (df.key == 'reward')]))
 
-        time_engaged.append(sum(engaged) + travel_time * 2 * len(block_trials))
-        block_time.append(sum(all_time))
-        block_rewards.append(sum(rewards))
-    engaged_df = pd.DataFrame()
-    engaged_df['percent engaged'] = np.array(time_engaged) / np.array(block_time)
-    engaged_df['block'] = blocks
-    engaged_df['time engaged'] = time_engaged
-    engaged_df['rewards earned'] = block_rewards
-    engaged_df['reward rate'] = np.array(block_rewards) / np.array(time_engaged)
-    return engaged_df
+            time_engaged.append(sum(engaged) + travel_time * 2 * len(block_trials))
+            block_time.append(sum(all_time))
+            block_rewards.append(sum(rewards))
+        engaged_df = pd.DataFrame()
+        engaged_df['percent engaged'] = np.array(time_engaged) / np.array(block_time)
+        engaged_df['block'] = blocks
+        engaged_df['time engaged'] = time_engaged
+        engaged_df['rewards earned'] = block_rewards
+        engaged_df['reward rate'] = np.array(block_rewards) / np.array(time_engaged)
+
+        return engaged_df
+    except Exception as e:
+        print ("Error in function.")
+        raise
 
 
 def reentry_index(df):
-    is_bg_exit = (df.port == 2) & (df.key == 'head') & (df.value == 0)
-    is_slow_block = df.groupby('trial').phase.agg(pd.Series.mode) == '0.4'
-    is_fast_block = df.groupby('trial').phase.agg(pd.Series.mode) == '0.8'
-    num_ideal_bg_entry_slow = len(np.unique(df.trial.dropna())[is_slow_block])
-    num_bg_entry_slow = len(df.index[is_bg_exit & df.trial.isin(
-        np.unique(df.trial.dropna())[is_slow_block])])
-    num_ideal_bg_entry_fast = len(np.unique(df.trial.dropna())[is_fast_block])
-    num_bg_entry_fast = len(df.index[is_bg_exit & df.trial.isin(
-        np.unique(df.trial.dropna())[is_fast_block])])
-
-    reentry_index_slow = num_bg_entry_slow / num_ideal_bg_entry_slow
-    reentry_index_fast = num_bg_entry_fast / num_ideal_bg_entry_fast
+    bgportassignment = df.loc[df['key'] == 'background', 'port'].iloc[-1]
+    expportassignment = df.loc[df['key'] == 'exp_decreasing', 'port'].iloc[-1]
+    is_bg_exit = (df.port == bgportassignment) & (df.key == 'head') & (df.value == 0)
+    phase_mode = df.groupby('trial').phase.agg(
+        lambda s: s.iloc[-1] if (s.value_counts().get('0.4', 0) == 1 and s.value_counts().get('0.8', 0) == 1)
+        else pd.Series.mode(s).iloc[-1]
+    )
+
+    is_low_block = (phase_mode == '0.4')
+    is_high_block = (phase_mode == '0.8')
+    num_ideal_bg_entry_low = len(np.unique(df.trial.dropna())[is_low_block])  # gets number of low block trials
+    num_bg_entry_low = len(df.index[is_bg_exit & df.trial.isin(
+        np.unique(df.trial.dropna())[is_low_block])])
+    num_ideal_bg_entry_high = len(np.unique(df.trial.dropna())[is_high_block])  # gets number of high block trials
+    num_bg_entry_high = len(df.index[is_bg_exit & df.trial.isin(
+        np.unique(df.trial.dropna())[is_high_block])])
+
+    reentry_index_low = num_bg_entry_low / num_ideal_bg_entry_low if num_ideal_bg_entry_low > 0 else 0
+    reentry_index_high = num_bg_entry_high / num_ideal_bg_entry_high if num_ideal_bg_entry_high > 0 else 0
     reentry_df = pd.DataFrame()
     reentry_df['block'] = ['0.4', '0.8']
-    reentry_df['bg_reentry_index'] = [reentry_index_slow, reentry_index_fast]
+    reentry_df['bg_reentry_index'] = [reentry_index_low, reentry_index_high]
     return reentry_df
 
 
 def add_h_lines(data=None, x=None, y=None, hue=None, ax=None, palette=None, estimator='mean'):
-    days_back = 10
+    days_back = 4
     palette = sns.color_palette(palette)
     for i, hue_key in enumerate(data[hue].unique()):
         df = data[data[hue] == hue_key]
@@ -390,6 +521,7 @@ def merge_old_trials(session):
 
 def simple_plots(select_mouse=None):
     plot_single_mouse_plots = True
+    save_folder = "C:\\Users\\Shichen\\OneDrive - Johns Hopkins\\ShulerLab\\Rie_behavior\\summary_graphs"
     if select_mouse is None:
         dif = date.today() - start_date
         data = gen_data(get_today_filepaths(days_back=dif.days), select_mouse=select_mouse)
@@ -405,73 +537,96 @@ def simple_plots(select_mouse=None):
         consumption = pd.DataFrame()
         block_leaves = pd.DataFrame()
         reentry = pd.DataFrame()
+        premature_leave = pd.DataFrame()
+
         for i, session in enumerate(data[mouse]):
-            if info[mouse][i]['task'] == 'cued_forgo_forced':
-                continue
+            # if info[mouse][i]['task'] == 'multi_reward':
+            #     continue
             try:
                 session = merge_old_trials(session)
 
                 engaged_df = percent_engaged(session)
-                engaged_df['day'] = [i] * len(engaged_df)
+                engaged_df['session'] = [i] * len(engaged_df)
                 engaged = pd.concat([engaged, engaged_df])
 
                 consumption_df = consumption_time(session)
-                consumption_df['day'] = [i] * len(consumption_df)
+                consumption_df['session'] = [i] * len(consumption_df)
                 consumption = pd.concat([consumption, consumption_df])
 
                 block_leaves_df = block_leave_times(session)
-                block_leaves_df['day'] = [i] * len(block_leaves_df)
+                block_leaves_df['session'] = [i] * len(block_leaves_df)
                 block_leaves = pd.concat([block_leaves, block_leaves_df])
 
                 reentry_df = reentry_index(session)
-                reentry_df['day'] = [i] * len(reentry_df)
+                reentry_df['session'] = [i] * len(reentry_df)
                 reentry = pd.concat([reentry, reentry_df])
+
+                premature_leave_df= calculate_premature_leave(session)
+                premature_leave_df['session'] = [i] * len(premature_leave_df)
+                premature_leave = pd.concat([premature_leave,premature_leave_df])
+
             except Exception as e:
-                raise e
+                print(f"Error processing session {i} for mouse {mouse}: {e}")
+                raise
 
         engaged.sort_values('block', inplace=True)
         block_leaves.sort_values('block', inplace=True)
         if plot_single_mouse_plots:
-            fig, axes = plt.subplots(2, 2, figsize=[11, 8], layout="constrained")
-            sns.lineplot(data=block_leaves.reset_index(), x='day', y='leave time', hue='block', ax=axes[0, 0],
+            fig, axes = plt.subplots(3, 2, figsize=[11, 8], layout="constrained")
+            sns.lineplot(data=block_leaves.reset_index(), x='session', y='leave time', hue='block', style = 'block',markers=True, ax=axes[0, 0],
                          palette='Set2')
-            add_h_lines(data=block_leaves.reset_index(), x='day', y='leave time', hue='block', ax=axes[0, 0],
+            add_h_lines(data=block_leaves.reset_index(), x='session', y='leave time', hue='block', ax=axes[0, 0],
                         palette='Set2')
-            sns.lineplot(data=consumption.reset_index(), x='day', y='consumption time', hue='port', ax=axes[0, 1],
+            sns.lineplot(data=consumption.reset_index(), x='session', y='consumption time', hue='port', style = 'port', markers=True, ax=axes[0, 1],
                          palette='Set1', estimator=np.median)
-            add_h_lines(data=consumption.reset_index(), x='day', y='consumption time', hue='port', ax=axes[0, 1],
+            add_h_lines(data=consumption.reset_index(), x='session', y='consumption time', hue='port', ax=axes[0, 1],
                         palette='Set1', estimator='median')
-            sns.lineplot(data=engaged.reset_index(), x='day', y='reward rate', hue='block', ax=axes[1, 0],
+            sns.lineplot(data=engaged.reset_index(), x='session', y='reward rate', hue='block', style = 'block', markers=True,ax=axes[1, 0],
+                         palette='Set2')
+            add_h_lines(data=engaged.reset_index(), x='session', y='reward rate', hue='block', ax=axes[1, 0],
+                        palette='Set2')
+            sns.lineplot(data=engaged.reset_index(), x='session', y='percent engaged', hue='block',style = 'block', markers=True, ax=axes[1, 1],
                          palette='Set2')
-            add_h_lines(data=engaged.reset_index(), x='day', y='reward rate', hue='block', ax=axes[1, 0],
+            add_h_lines(data=engaged.reset_index(), x='session', y='percent engaged', hue='block', ax=axes[1, 1],
                         palette='Set2')
-            sns.lineplot(data=engaged.reset_index(), x='day', y='percent engaged', hue='block', ax=axes[1, 1],
+            sns.lineplot(data=premature_leave.reset_index(), x='session', y='premature leave rate', hue='block', style = 'block',markers=True,ax=axes[2, 0],
                          palette='Set2')
-            add_h_lines(data=engaged.reset_index(), x='day', y='percent engaged', hue='block', ax=axes[1, 1],
+            add_h_lines(data=premature_leave.reset_index(), x='session', y='premature leave rate', hue='block', ax=axes[2, 0],
                         palette='Set2')
 
+            axes[2, 0].axhline(y=0.2, color='red', linestyle='--', linewidth=1.5, label='Threshold = 0.2')
+
+            # Add legend to show the line label
+            axes[2, 0].legend()
+
             axes[0, 0].set_title('Leave Time by Block')
             axes[0, 1].set_title('Consumption Time by Port')
             axes[1, 0].set_title('Reward Rate by Block')
             axes[1, 1].set_title('Percent Time Engaged by Block')
+            axes[2, 0].set_title('Premature leave from BG port by Block')
 
             axes[0, 0].set_ylim([0, 20])
             axes[0, 1].set_ylim([0, 20])
             axes[1, 0].set_ylim([0, .65])
             axes[1, 1].set_ylim([0, 1])
+            axes[2, 0].set_ylim([0, 1])
             plt.suptitle(mouse, fontsize=20)
+            os.makedirs(save_folder, exist_ok=True)
+            # Construct the filename
+            filename = f'{mouse}_session_summary.png'
+            save_path = os.path.join(save_folder, filename)
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+            print(f"Graph saved to: {save_path}")
             plt.show()
 
-        block_leaves_last10_df = block_leaves[(block_leaves.day >= block_leaves.day.max() - 10)].groupby('block')[
+        block_leaves_last10_df = block_leaves[(block_leaves.session >= block_leaves.session.max() - 10)].groupby('block')[
             'leave time'].mean().reset_index()
         block_leaves_last10_df['animal'] = mouse
         block_leaves_last10 = pd.concat([block_leaves_last10, block_leaves_last10_df])
 
     fig, axes = plt.subplots(1, 1)
     sns.boxplot(data=block_leaves_last10.reset_index(), x='block', y='leave time')
-    for mouse in data.keys():
-        plt.plot([-0.1, 0.9], block_leaves_last10[block_leaves_last10.animal == mouse]['leave time'], 'o-',
-                 color='darkgray')
+
     plt.show()
 
 
@@ -492,7 +647,125 @@ def single_session(select_mouse=None, num_back=2):
             session_summary(last_session, mouse, last_info)
 
 
+# def session_summary(data, mouse, info):
+#     base_save_folder= "C:\\Users\\riepo\\Experiment1\\graphs\\each_session"
+#     save_folder = os.path.join(base_save_folder, mouse)
+#     os.makedirs(save_folder, exist_ok=True)
+#     fig, [ax1, ax2] = plt.subplots(1, 2, figsize=[10, 10])
+#     port_palette = sns.color_palette('Set1')
+#     block_palette = sns.color_palette('Set2')
+#     start = data.value == 1
+#     end = data.value == 0
+#     head = data.key == 'head'
+#     lick = data.key == 'lick'
+#     reward = data.key == 'reward'
+#     bgport = data.port == data.loc[data['key'] == 'background', 'port'].iloc[-1]
+#     expport = data.port == data.loc[data['key'] == 'exp_decreasing', 'port'].iloc[-1]
+#     # port1 = data.port == 1
+#     # port2 = data.port == 2
+#     # print(f"port2 is {port2}")
+#     # print(f"bgport is {bgport}")
+#     max_trial = data.trial.max()
+#
+#     bg_rectangles = []
+#     exp_rectangles_in_bg = []
+#     exp_rectangles = []
+#     block1_rectangles = []
+#     block2_rectangles = []
+#     bg_reward_events = []
+#     exp_reward_events = []
+#     bg_lick_events = []
+#     exp_lick_events = []
+#     bg_lengths = []
+#     exp_lengths = []
+#     trial_blocks = data.groupby(['trial'])['phase'].agg(pd.Series.mode)
+#     blocks = data.phase.dropna().unique()
+#     blocks.sort()
+#     for trial in data.trial.unique():
+#         if np.isnan(trial):
+#             continue
+#         is_trial = data.trial == trial
+#         try:
+#             trial_start = data[is_trial & start & (data.key == 'trial')].session_time.values[0]
+#             trial_middle = data[is_trial & end & (data.key == 'LED') & bgport].session_time.values[0]
+#             trial_end = data[is_trial & end & (data.key == 'trial')].session_time.values[0]
+#         except IndexError:
+#             continue
+#
+#         bg_rewards = data[is_trial & start & bgport & reward].session_time.values
+#         exp_rewards = data[is_trial & start & expport & reward].session_time.values
+#         bg_licks = data[is_trial & start & lick & (data.session_time < trial_middle)].session_time.values
+#         exp_licks = data[is_trial & start & lick & (data.session_time > trial_middle)].session_time.values
+#
+#         bg_lengths.append(trial_middle - trial_start)
+#         exp_lengths.append(trial_end - trial_middle)
+#
+#         bg_entries, bg_exits, exp_entries, exp_exits, early_exp_entries, early_exp_exits = get_entry_exit(data, trial)
+#         bg_intervals = list(zip(bg_entries, bg_exits))
+#         exp_intervals = list(zip(exp_entries, exp_exits))
+#         early_exp_intervals = list(zip(early_exp_entries, early_exp_exits))
+#         for [s, e] in bg_intervals:
+#             bg_rectangles.append(Rectangle((s - trial_start, trial), e - s, .7))
+#         for [s, e] in early_exp_intervals:
+#             exp_rectangles_in_bg.append(Rectangle((s - trial_start, trial), e - s, .7))
+#         for [s, e] in exp_intervals:
+#             exp_rectangles.append(Rectangle((s - trial_middle, trial), e - s, .7))
+#         if np.where(blocks == trial_blocks.loc[trial])[0][0] == 0:
+#             block1_rectangles.append(Rectangle((0, trial), 100, 1))
+#         else:
+#             block2_rectangles.append(Rectangle((0, trial), 100, 1))
+#         bg_reward_events.append(bg_rewards - trial_start)
+#         exp_reward_events.append(exp_rewards - trial_middle)
+#         bg_lick_events.append(bg_licks - trial_start)
+#         exp_lick_events.append(exp_licks - trial_middle)
+#
+#     alpha = .5
+#     pc_b1 = PatchCollection(block1_rectangles, facecolors=block_palette[0], alpha=alpha)
+#     pc_b2 = PatchCollection(block2_rectangles, facecolors=block_palette[1], alpha=alpha)
+#     ax1.add_collection(pc_b1)
+#     ax1.add_collection(pc_b2)
+#     pc_b12 = PatchCollection(block1_rectangles, facecolors=block_palette[0], alpha=alpha)
+#     pc_b22 = PatchCollection(block2_rectangles, facecolors=block_palette[1], alpha=alpha)
+#     ax2.add_collection(pc_b12)
+#     ax2.add_collection(pc_b22)
+#
+#     pc_bg = PatchCollection(bg_rectangles, edgecolor=port_palette[0], facecolor='w', alpha=1)
+#     ax1.add_collection(pc_bg)
+#
+#     pc_exp_bg = PatchCollection(exp_rectangles_in_bg, edgecolor=port_palette[1], facecolor='w', alpha=1)
+#     ax1.add_collection(pc_exp_bg)
+#
+#     pc_exp = PatchCollection(exp_rectangles, edgecolor=port_palette[1], facecolor='w', alpha=1)
+#     ax2.add_collection(pc_exp)
+#
+#     offsets = np.array(list(range(len(bg_reward_events)))) + 1.4
+#     ax1.eventplot(bg_reward_events, color='purple', linelengths=.62, lineoffsets=offsets)
+#     offsets = np.array(list(range(len(exp_reward_events)))) + 1.4
+#     ax2.eventplot(exp_reward_events, color='purple', linelengths=.62, lineoffsets=offsets)
+#
+#     light = [.8, .7, .8]
+#     dark = [.2, .2, .2]
+#     offsets = np.array(list(range(len(bg_lick_events)))) + 1.4
+#     ax1.eventplot(bg_lick_events, color=light, linelengths=.25, lineoffsets=offsets)
+#     offsets = np.array(list(range(len(exp_lick_events)))) + 1.4
+#     ax2.eventplot(exp_lick_events, color=light, linelengths=.25, lineoffsets=offsets)
+#
+#     session_summary_axis_settings([ax1, ax2], max_trial)
+#     plt.suptitle(f'{mouse}: {info["date"]} {info["time"]}')
+#
+#     # Construct the filename
+#     filename = f'{mouse}_{info["date"]}_{info["time"]}.png'
+#     save_path = os.path.join(save_folder, filename)
+#
+#     # Save the plot
+#     plt.savefig(save_path, dpi=300, bbox_inches='tight')
+#     print(f"Graph saved to: {save_path}")
+#     plt.show()
+
 def session_summary(data, mouse, info):
+    base_save_folder= save_folder = "C:\\Users\\Shichen\\OneDrive - Johns Hopkins\\ShulerLab\\Rie_behavior\\each_session"
+    save_folder = os.path.join(base_save_folder, mouse)
+    os.makedirs(save_folder, exist_ok=True)
     fig, [ax1, ax2] = plt.subplots(1, 2, figsize=[10, 10])
     port_palette = sns.color_palette('Set1')
     block_palette = sns.color_palette('Set2')
@@ -501,8 +774,8 @@ def session_summary(data, mouse, info):
     head = data.key == 'head'
     lick = data.key == 'lick'
     reward = data.key == 'reward'
-    port1 = data.port == 1
-    port2 = data.port == 2
+    bgport = data.port == data.loc[data['key'] == 'background', 'port'].iloc[-1]
+    expport = data.port == data.loc[data['key'] == 'exp_decreasing', 'port'].iloc[-1]
     max_trial = data.trial.max()
 
     bg_rectangles = []
@@ -517,7 +790,7 @@ def session_summary(data, mouse, info):
     bg_lengths = []
     exp_lengths = []
     trial_blocks = data.groupby(['trial'])['phase'].agg(pd.Series.mode)
-    blocks = data.phase.unique()
+    blocks = data.phase.dropna().unique()
     blocks.sort()
     for trial in data.trial.unique():
         if np.isnan(trial):
@@ -525,13 +798,13 @@ def session_summary(data, mouse, info):
         is_trial = data.trial == trial
         try:
             trial_start = data[is_trial & start & (data.key == 'trial')].session_time.values[0]
-            trial_middle = data[is_trial & end & (data.key == 'LED') & port2].session_time.values[0]
+            trial_middle = data[is_trial & end & (data.key == 'LED') & bgport].session_time.values[0]
             trial_end = data[is_trial & end & (data.key == 'trial')].session_time.values[0]
         except IndexError:
             continue
 
-        bg_rewards = data[is_trial & start & port2 & reward].session_time.values
-        exp_rewards = data[is_trial & start & port1 & reward].session_time.values
+        bg_rewards = data[is_trial & start & bgport & reward].session_time.values
+        exp_rewards = data[is_trial & start & expport & reward].session_time.values
         bg_licks = data[is_trial & start & lick & (data.session_time < trial_middle)].session_time.values
         exp_licks = data[is_trial & start & lick & (data.session_time > trial_middle)].session_time.values
 
@@ -548,10 +821,13 @@ def session_summary(data, mouse, info):
             exp_rectangles_in_bg.append(Rectangle((s - trial_start, trial), e - s, .7))
         for [s, e] in exp_intervals:
             exp_rectangles.append(Rectangle((s - trial_middle, trial), e - s, .7))
-        if np.where(blocks == trial_blocks.loc[trial])[0][0] == 0:
+
+        block_value = trial_blocks.loc[trial]
+        if block_value == "0.4":
             block1_rectangles.append(Rectangle((0, trial), 100, 1))
         else:
             block2_rectangles.append(Rectangle((0, trial), 100, 1))
+
         bg_reward_events.append(bg_rewards - trial_start)
         exp_reward_events.append(exp_rewards - trial_middle)
         bg_lick_events.append(bg_licks - trial_start)
@@ -590,8 +866,15 @@ def session_summary(data, mouse, info):
 
     session_summary_axis_settings([ax1, ax2], max_trial)
     plt.suptitle(f'{mouse}: {info["date"]} {info["time"]}')
-    plt.show()
 
+    # Construct the filename
+    filename = f'{mouse}_{info["date"]}_{info["time"]}.png'
+    save_path = os.path.join(save_folder, filename)
+
+    # Save the plot
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"Graph saved to: {save_path}")
+    plt.show()
 
 def session_summary_axis_settings(axes, max_trial):
     for ax in axes:
@@ -601,17 +884,178 @@ def session_summary_axis_settings(axes, max_trial):
         ax.spines['bottom'].set_visible(True)
         ax.get_yaxis().set_visible(False)
         ax.set_ylim([-1, max_trial + 1])
-        ax.set_xlim([0, 20])
+        ax.set_xlim([0, 30])
         ax.invert_yaxis()
         ax.set_ylabel('Trial')
         ax.set_xlabel('Time (sec)')
 
 
+
+def simple_plots2(select_mouse=None):
+    plot_single_mouse_plots = True
+    save_folder = "C:\\Users\\Shichen\\OneDrive - Johns Hopkins\\ShulerLab\\Rie_behavior\\summary_graphs"
+
+    #insert vertical line between stages or pre/post surgery
+    stage_changes = {
+        # "RK001": {"1->2": "2025-01-22", "2->3": "2025-01-31"},
+        # "RK002": {"1->2": "2025-01-03", "2->3": "2025-03-14"},
+        # "RK003": {"1->2": "2025-01-03", "2->3": "2025-01-22"},
+        # "RK004": {"1->2": "2025-02-24", "2->3": "2025-03-05"},
+        # "RK005": {"1->2": "2025-01-27", "2->3": "2025-01-31"},
+        # "RK006": {"1->2": "2024-12-20", "2->3": "2025-02-26"},
+        # "ES039": {"1->2": "2024-02-14", "2->3": None},
+        # "ES045": {"1->2": "2024-04-24", "2->3": None},
+        # "ES046": {"1->2": "2024-05-02", "2->3": None},
+        # "ES047": {"1->2": "2024-04-24", "2->3": None},
+        "RK001": {"1->2": "2025-04-13", "2->3": None}, #surgery date
+        "RK002": {"1->2": "2025-04-13", "2->3": None},
+        "RK003": {"1->2": "2025-04-13", "2->3": None},
+        "RK004": {"1->2": "2025-03-30", "2->3": None},
+        "RK005": {"1->2": "2025-03-30", "2->3": None},
+        "RK006": {"1->2": "2025-03-30", "2->3": None}
+    }
+
+    if select_mouse is None:
+        dif = date.today() - start_date
+        data = gen_data(get_today_filepaths(days_back=dif.days), select_mouse=select_mouse)
+        info = gen_data(get_today_filepaths(days_back=dif.days), select_mouse=select_mouse, return_info=True)
+    else:
+        data = gen_data(get_today_filepaths(days_back=1000), select_mouse=select_mouse)
+        info = gen_data(get_today_filepaths(days_back=1000), select_mouse=select_mouse, return_info=True)
+    block_leaves_last10 = pd.DataFrame()
+
+    for mouse in data.keys():
+        if select_mouse is not None and mouse not in select_mouse:
+            continue
+
+        # Initialize empty dataframes for aggregated data.
+        engaged = pd.DataFrame()
+        consumption = pd.DataFrame()
+        block_leaves = pd.DataFrame()
+        reentry = pd.DataFrame()
+        premature_leave = pd.DataFrame()
+
+        # Extract session dates
+        session_dates = []
+        for sess_info in info[mouse]:
+            try:
+                dt = datetime.strptime(sess_info['date'], "%Y-%m-%d")
+                session_dates.append(dt)
+            except Exception as e:
+                print(f"Error parsing date in session info: {sess_info.get('date', None)}: {e}")
+
+        for i, session in enumerate(data[mouse]):
+            # if info[mouse][i]['task'] == 'multi_reward':
+            #     continue
+            try:
+                session = merge_old_trials(session)
+
+                engaged_df = percent_engaged(session)
+                engaged_df['session'] = [i] * len(engaged_df)
+                engaged = pd.concat([engaged, engaged_df])
+
+                consumption_df = consumption_time(session)
+                consumption_df['session'] = [i] * len(consumption_df)
+                consumption = pd.concat([consumption, consumption_df])
+
+                block_leaves_df = block_leave_times(session)
+                block_leaves_df['session'] = [i] * len(block_leaves_df)
+                block_leaves = pd.concat([block_leaves, block_leaves_df])
+
+                reentry_df = reentry_index(session)
+                reentry_df['session'] = [i] * len(reentry_df)
+                reentry = pd.concat([reentry, reentry_df])
+
+                premature_leave_df = calculate_premature_leave(session)
+                premature_leave_df['session'] = [i] * len(premature_leave_df)
+                premature_leave = pd.concat([premature_leave, premature_leave_df])
+
+            except Exception as e:
+                print(f"Error processing session {i} for mouse {mouse}: {e}")
+                raise
+
+        engaged.sort_values('block', inplace=True)
+        block_leaves.sort_values('block', inplace=True)
+
+        if plot_single_mouse_plots:
+            fig, axes = plt.subplots(3, 2, figsize=[11, 8], layout="constrained")
+
+            sns.lineplot(data=block_leaves.reset_index(), x='session', y='leave time',
+                         hue='block', style='block', markers=True, ax=axes[0, 0], palette='Set2')
+            add_h_lines(data=block_leaves.reset_index(), x='session', y='leave time',
+                        hue='block', ax=axes[0, 0], palette='Set2')
+
+            sns.lineplot(data=consumption.reset_index(), x='session', y='consumption time',
+                         hue='port', style='port', markers=True, ax=axes[0, 1],
+                         palette='Set1', estimator=np.median)
+            add_h_lines(data=consumption.reset_index(), x='session', y='consumption time',
+                        hue='port', ax=axes[0, 1], palette='Set1', estimator='median')
+
+            sns.lineplot(data=engaged.reset_index(), x='session', y='reward rate',
+                         hue='block', style='block', markers=True, ax=axes[1, 0], palette='Set2')
+            add_h_lines(data=engaged.reset_index(), x='session', y='reward rate',
+                        hue='block', ax=axes[1, 0], palette='Set2')
+
+            sns.lineplot(data=engaged.reset_index(), x='session', y='percent engaged',
+                         hue='block', style='block', markers=True, ax=axes[1, 1], palette='Set2')
+            add_h_lines(data=engaged.reset_index(), x='session', y='percent engaged',
+                        hue='block', ax=axes[1, 1], palette='Set2')
+
+            sns.lineplot(data=premature_leave.reset_index(), x='session', y='premature leave rate',
+                         hue='block', style='block', markers=True, ax=axes[2, 0], palette='Set2')
+            add_h_lines(data=premature_leave.reset_index(), x='session', y='premature leave rate',
+                        hue='block', ax=axes[2, 0], palette='Set2')
+
+            axes[2, 0].axhline(y=0.2, color='red', linestyle='--', linewidth=1.5, label='Threshold = 0.2')
+            axes[2, 0].legend()
+
+            # Add vertical lines for stage changes (blue dashed lines)
+            if mouse in stage_changes:
+                for change, change_date_str in stage_changes[mouse].items():
+                    if change_date_str is not None:
+                        try:
+                            change_date = datetime.strptime(change_date_str, "%Y-%m-%d")
+                            idx = next((j for j, d in enumerate(session_dates) if d >= change_date), None)
+                            if idx is not None and idx > 0:
+                                pos = idx - 0.5
+                                for ax in axes.flatten():
+                                    ax.axvline(x=pos, color='blue', linestyle='--', linewidth=1.5)
+                        except Exception as e:
+                            print(f"Error processing stage change date {change_date_str} for mouse {mouse}: {e}")
+
+            axes[0, 0].set_title('Leave Time by Block')
+            axes[0, 1].set_title('Consumption Time by Port')
+            axes[1, 0].set_title('Reward Rate by Block')
+            axes[1, 1].set_title('Percent Time Engaged by Block')
+            axes[2, 0].set_title('Premature leave from BG port by Block')
+
+            axes[0, 0].set_ylim([0, 20])
+            axes[0, 1].set_ylim([0, 20])
+            axes[1, 0].set_ylim([0, 0.65])
+            axes[1, 1].set_ylim([0, 1])
+            axes[2, 0].set_ylim([0, 1])
+
+            plt.suptitle(f"{mouse}", fontsize=20)
+            os.makedirs(save_folder, exist_ok=True)
+            filename = f'{mouse}_session_summary.png'
+            save_path = os.path.join(save_folder, filename)
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+            print(f"Graph saved to: {save_path}")
+            plt.show()
+
+            block_leaves_last10_df = block_leaves[(block_leaves.session >= block_leaves.session.max() - 10)]\
+                .groupby('block')['leave time'].mean().reset_index()
+            block_leaves_last10_df['animal'] = mouse
+            block_leaves_last10 = pd.concat([block_leaves_last10, block_leaves_last10_df])
+
+    fig, axes = plt.subplots(1, 1)
+    sns.boxplot(data=block_leaves_last10.reset_index(), x='block', y='leave time')
+    plt.show()
+
 if __name__ == '__main__':
-    # mice = ['ES057', 'ES058', 'ES059', 'ES060', 'ES061', 'ES062']
-    # mice = ['ES045', 'ES046', 'ES047', 'ES051', 'ES052', 'ES053', 'ES057', 'ES060', 'ES061', 'ES062']
-    # mice = ['ES058', 'ES059', 'ES045', 'ES047']
-    mice = ['ES057', 'ES046']
-    # mice = ['ES051', 'ES052', 'ES053', 'ES060', 'ES061', 'ES062']
-    simple_plots(mice)
-    # single_session(mice)
+    mice = ['RK007','RK008','RK009','RK010']
+    # mice = ['RK004']
+    # mice = ['RK001', 'RK002','RK003','RK004','RK005','RK006','ES039', 'ES045','ES046','ES047','ES051', 'ES052', 'ES053']
+    simple_plots2(mice) #history day by day summary statistics
+    single_session(mice)  # most recent session when num_back = 0
+
diff --git a/stand_alone/gui.py b/stand_alone/gui.py
index 5127c4e..d354113 100644
--- a/stand_alone/gui.py
+++ b/stand_alone/gui.py
@@ -12,20 +12,20 @@
 from user_settings import get_user_info
 import os
 
-info_dict = get_user_info(box=os.uname()[1])
+info_dict = get_user_info()
 
 # scp -r C:\Users\Elissa\GoogleDrive\Code\Python\behavior_code\stand_alone pi@elissapi0:\home\pi
 # scp C:\Users\Elissa\GoogleDrive\Code\Python\behavior_code\stand_alone\scp_rescue.py pi@elissapi1:\home\pi\behavior
 
 # scp -r "C:\Users\Shichen\OneDrive - Johns Hopkins\ShulerLab\behavior_code\stand_alone" pi@elissapi1:\home\pi\behavior1
 pastel_colors = ['#ffffcc', '#99ccff', '#cc99ff', '#ff99cc', '#ffcc99', '#ffffcc', '#99ffcc', '#ccffff', '#ccccff',
-                 '#ffccff', '#ffcccc', '#D3D3D3']
+                 '#ffccff', '#ffcccc', '#D3D3D3', '#f0a207']
 
 
 class Gui:
     def __init__(self):
         self.root = tk.Tk()
-        self.root.geometry('2000x700')
+        self.root.geometry('960x360')
         self.root.title('BehaviorGui')
 
         with open('durations.pkl', 'rb') as f:
@@ -37,9 +37,11 @@ def __init__(self):
         myFont = font.Font(size=30)
         mouse_rows = len(info_dict['mouse_buttons'])
         self.mouse_assignments = info_dict['mouse_assignments']
+        self.port_swap_assignments = info_dict['port_swap_assignments']
+        self.block_swap_assignments = info_dict['block_swap_assignments']
         tasks = {
             'single_reward': single_reward,
-            'cued_forgo': cued_forgo
+            'multi_reward': multi_reward
         }
         for key in self.mouse_assignments.keys():
             self.mouse_assignments[key] = tasks[self.mouse_assignments[key]]
@@ -105,8 +107,10 @@ def __init__(self):
 
     def run_behavior(self, mouse):
         task = self.mouse_assignments[mouse]
+        swap_port = self.port_swap_assignments[mouse]
+        swap_block = self.block_swap_assignments[mouse]
         print(f"running {task} for {mouse}")
-        main(mouse, task, forgo=False, forced_trials=True)
+        main(mouse, task, swap_port=swap_port, swap_block=swap_block)
 
     def calibrate(self, port):
         print(f'calibrating port {port}')
diff --git a/stand_alone/main.py b/stand_alone/main.py
index a5e0a45..5e523de 100644
--- a/stand_alone/main.py
+++ b/stand_alone/main.py
@@ -1,5 +1,6 @@
 from support_classes import *
-from tasks import *
+# from tasks import *
+from tasks_RK import multireward_task
 from timescapes import *
 import random
 
@@ -36,68 +37,76 @@
 #     session.start([task1])
 
 
-def cued_forgo(session, reward_level, starting_prob, session_length, forgo=False, forced_trials=True):
+def multi_reward(session, reward_level, starting_prob, session_length, swap_port=False, swap_block = False):
     print(reward_level)
     print(starting_prob)
     print(session_length)
-    task_structure = cued_forgo_task
-
-    if forgo:
-        task_name = 'cued_forgo'
-        print('cued forgo')
+    task_structure = multireward_task
+
+    # rates = [.4, .8, .4, .8, .4, .8]
+    # if np.random.random() > .5:
+    #     rates.reverse()
+
+    # rates = [.4, .8] #swap_block random is session by session change, swap block trial is trial by trial change
+    if swap_block == "random":
+        rates = [random.choice([0.4, 0.8])]
+        print (rates)
+    elif swap_block == "trial":
+        rates = [.4, .8, .4, .8, .4, .8]
+        if np.random.random() > .5:
+            rates.reverse()
     else:
-        task_name = 'cued_no_forgo'
-        print('cued bg without forgo option')
-
-    if forced_trials:
-        task_name = task_name + '_forced'
-        print('forced trials included')
+        rates = [0.4] if not swap_block else [0.8] #0.4 if swap_block is false, 0.8 if true
+    #     rates.reverse()
 
-    rates = [.4, .8, .4, .8, .4, .8]
-    if np.random.random() > .5:
-        rates.reverse()
+    exp_port = 1 if not swap_port else 2
+    bg_port = 2 if not swap_port else 1
+    print(f'exp_port = {exp_port}')
+    print(f'bg_port = {bg_port}')
     background_dist = {'distribution': 'background',
                        'rates': rates,
                        'duration': 5,
-                       'port_num': 2}
+                       'port_num': bg_port}
     print(background_dist['rates'])
     exp_dist = {'distribution': exp_decreasing,
                 'cumulative': reward_level,
                 'starting_probability': starting_prob,
-                'port_num': 1}
+                'port_num': exp_port}
     ports = {'exp': Port(exp_dist['port_num'], dist_info=exp_dist),
              'background': Port(background_dist['port_num'], dist_info=background_dist)}
-    task1 = Task(session, name=task_name, structure=task_structure, ports=ports,
-                 maximum=session_length, limit='time', forgo=forgo, forced_trials=forced_trials)
+    task1 = Task(session, name='multi_reward', structure=task_structure, ports=ports,
+                 maximum=session_length, limit='time')
     session.start([task1])
 
 
-def single_reward(session, reward_level, starting_prob, session_length, forgo=False, forced_trials=True):
+def single_reward(session, reward_level, starting_prob, session_length, swap_port=False, swap_block= False):
     reward_level = 0.5994974874371859  # cumulative for an 8 reward version
     starting_prob = 0.1301005025125628
     print(reward_level)
     print(starting_prob)
     print(session_length)
     task_structure = single_reward_task
-
+    
     task_name = 'single_reward'
-
+    
     rates = [.4, .8, .4, .8, .4, .8]
     if np.random.random() > .5:
         rates.reverse()
+    exp_port = 1 if not swap_port else 2
+    bg_port = 2 if not swap_port else 1
     background_dist = {'distribution': 'background',
-                       'rates': rates,
-                       'duration': 5,
-                       'port_num': 2}
+                   'rates': rates,
+                   'duration': 5,
+                   'port_num': bg_port}
     print(background_dist['rates'])
     exp_dist = {'distribution': exp_decreasing,
-                'cumulative': reward_level,
-                'starting_probability': starting_prob,
-                'port_num': 1}
+            'cumulative': reward_level,
+            'starting_probability': starting_prob,
+            'port_num': exp_port}
     ports = {'exp': Port(exp_dist['port_num'], dist_info=exp_dist),
-             'background': Port(background_dist['port_num'], dist_info=background_dist)}
+         'background': Port(background_dist['port_num'], dist_info=background_dist)}
     task1 = Task(session, name=task_name, structure=task_structure, ports=ports,
-                 maximum=session_length, limit='time')
+             maximum=session_length, limit='time')
     session.start([task1])
 
 # def give_up_blocked(session, reward_level, starting_prob, session_length, forgo=True, forced_trials=False):
@@ -122,7 +131,8 @@ def single_reward(session, reward_level, starting_prob, session_length, forgo=Fa
 #     session.start([task1])
 
 
-def main(mouse, to_run, forgo=False, forced_trials=False):
+# def main(mouse, to_run, forgo=False, forced_trials=False, swap_port=False):
+def main(mouse, to_run, swap_port=False, swap_block = False): #swap port swaps BG/EXP, swap_block is the stage
     cumulative = 8
     start_prob = 1
     session_time = 18
@@ -135,9 +145,9 @@ def main(mouse, to_run, forgo=False, forced_trials=False):
 
     try:
         if mouse not in mouse_settings.keys():
-            to_run(session, *mouse_settings['default'], forgo=forgo, forced_trials=forced_trials)  # Run the task
+            to_run(session, *mouse_settings['default'], swap_port=swap_port, swap_block=swap_block)  # Run the task
         else:
-            to_run(session, *mouse_settings[mouse], forgo=forgo, forced_trials=forced_trials)  # Run the task
+            to_run(session, *mouse_settings[mouse], swap_port=swap_port, swap_block=swap_block)  # Run the task
         session.smooth_finish = True
         print('smooth finish')
     except KeyboardInterrupt:  # Catch if the task is stopped via ctrl-C or the stop button
@@ -149,6 +159,6 @@ def main(mouse, to_run, forgo=False, forced_trials=False):
 if __name__ == "__main__":
     # main('testmouse', cued_forgo, forgo=False, forced_trials=True)
     # main('ES024', cued_forgo, forgo=False, forced_trials=True)
-    main('ES030', cued_forgo, forgo=False, forced_trials=True)
+    main('ES030', multi_reward)
     # main('testmouse', give_up_blocked)
 
diff --git a/stand_alone/support_classes.py b/stand_alone/support_classes.py
index b5ad92a..427c4d5 100644
--- a/stand_alone/support_classes.py
+++ b/stand_alone/support_classes.py
@@ -42,7 +42,72 @@ def test(i):
     print('test ' + str(i))
 
 
-def ssh(host, cmd, user, password, timeout=30, bg_run=False):
+# def ssh(host, cmd, user, password, timeout=30, bg_run=False):
+#     """SSH'es to a host using the supplied credentials and executes a command.
+#     Throws an exception if the command doesn't return 0.
+#     bgrun: run command in the background"""
+#
+#     options = '-q -oStrictHostKeyChecking=no -oUserKnownHostsFile=/dev/null -oPubkeyAuthentication=no'
+#     if bg_run:
+#         options += ' -f'
+#
+#     ssh_cmd = 'ssh %s@%s %s \'%s\'' % (user, host, options, cmd)
+#     print(ssh_cmd)
+#     child = pexpect.spawnu(ssh_cmd, timeout=timeout)
+#     child.expect([r'(?i).*password.*:'])
+#     child.sendline(password)
+#     child.expect(pexpect.EOF)
+#     child.close()
+
+
+# def scp(host, filename, destination, user, password, timeout=30, bg_run=False, recursive=False, cmd=False):
+#     """Scp's to a host using the supplied credentials and executes a command.
+#     Throws an exception if the command doesn't return 0.
+#     bgrun: run command in the background"""
+#
+#     options = '-q -oStrictHostKeyChecking=no -oUserKnownHostsFile=/dev/null -oPubkeyAuthentication=no'
+#     if recursive:
+#         options += ' -r'
+#     if bg_run:
+#         options += ' -f'
+#     scp_cmd = 'scp %s %s %s@%s:\'"%s"\'' % (options, filename, user, host, os.path.join(destination, filename))
+#     password_prompt = rf"{user}@{host}'s password: "
+#     print(password_prompt)
+#     child = None
+#     status = 1
+#     print(scp_cmd)
+#     try:
+#         if cmd:
+#             return scp_cmd
+#         child = pexpect.spawnu(scp_cmd, timeout=timeout)  # spawnu for Python 3
+#         child.expect(password_prompt)
+#         child.sendline(password)
+#         child.expect(pexpect.EOF)
+#         child.close()
+#         status = child.exitstatus
+#     except Exception as e:
+#         print('scp didn\'t work the first time so we are trying again with different password')
+#         try:
+#             # scp_cmd2 = 'scp %s %s %s@%s:"%s"' % (options, filename, user, host, os.path.join(destination, filename))
+#             # print(scp_cmd2)
+#             # if cmd:
+#             #     return scp_cmd2
+#             child = pexpect.spawnu(scp_cmd, timeout=timeout)  # spawnu for Python 3
+#             child.expect([r"(?i).*password.*:?(\s*)"])
+#             child.sendline(password)
+#             child.expect(pexpect.EOF)
+#             child.close()
+#             status = child.exitstatus
+#         except Exception as e_2:
+#             print(f"still didnt work: {e_2}")
+#             status = 1
+#     if status == 0:
+#         print("file transfer completed!")
+#     else:
+#         print("transfer failed.")
+#     return status
+
+def ssh(host, cmd, user, password, timeout=60, bg_run=False):
     """SSH'es to a host using the supplied credentials and executes a command.
     Throws an exception if the command doesn't return 0.
     bgrun: run command in the background"""
@@ -54,12 +119,10 @@ def ssh(host, cmd, user, password, timeout=30, bg_run=False):
     ssh_cmd = 'ssh %s@%s %s \'%s\'' % (user, host, options, cmd)
     print(ssh_cmd)
     child = pexpect.spawnu(ssh_cmd, timeout=timeout)
-    child.expect(['[Pp]assword: '])
+    child.expect([r"(?i).*password.*:?(\s*)"])
     child.sendline(password)
     child.expect(pexpect.EOF)
     child.close()
-
-
 def scp(host, filename, destination, user, password, timeout=30, bg_run=False, recursive=False, cmd=False):
     """Scp's to a host using the supplied credentials and executes a command.
     Throws an exception if the command doesn't return 0.
@@ -96,7 +159,6 @@ def scp(host, filename, destination, user, password, timeout=30, bg_run=False, r
             print('worked!')
     return child.exitstatus
 
-
 def sync_stream(self):
     pin_map = {'session': 25,
                'trial': 8,
@@ -142,7 +204,7 @@ def __init__(self, mouse):
         self.data_write_path = "/data/" + self.mouse
         self.datetime = time.strftime("%Y-%m-%d_%H-%M-%S")
         self.filename = "data_" + self.datetime + ".txt"
-
+        self.isstarted = False
         self.halted = False
         self.smooth_finish = False
         GPIO.setmode(GPIO.BCM)
@@ -167,6 +229,7 @@ def __init__(self, mouse):
         GPIO.output(self.camera_pin, GPIO.LOW)
 
     def start(self, task_list):
+        self.isstarted = True
         for val in self.sync_pins.values():
             GPIO.setup(val, GPIO.OUT)
             GPIO.output(val, GPIO.LOW)
@@ -190,10 +253,11 @@ def start(self, task_list):
             perform(task)
 
     def log(self, string):  # Adds session time stamp to beginning of string and logs it
-        session_time = time.time() - self.start_time
-        new_line = str(session_time) + ',' + string + '\n'
-        # print(new_line)
-        self.f.write(new_line)
+        if self.isstarted:
+            session_time = time.time() - self.start_time
+            new_line = str(session_time) + ',' + string + '\n'
+            # print(new_line)
+            self.f.write(new_line)
 
     def end(self):
         self.log('nan,nan,nan,nan,setup,nan,0,session')
@@ -201,16 +265,21 @@ def end(self):
         os.system('sudo chmod o-w ' + self.filename)
         mkdir_command = 'if not exist "%s" mkdir "%s"' % (
             self.ssh_path.replace('/', '\\'), self.ssh_path.replace('/', '\\'))
-        ssh(self.ip, mkdir_command, self.user, self.password)
-        res = scp(self.ip, self.filename, self.data_send_path, self.user, self.password)
-        if not res:
-            print('\nSuccessful file transfer to "%s"\nDeleting local file from pi.' % self.data_send_path)
-            os.remove(self.filename)
-        else:
-            print('connection back to desktop timed out')
-        GPIO.cleanup()
-        os.chdir(os.path.join(os.getcwd(), '..', '..'))
-        print('\nFile closed and clean up complete')
+        try:
+            ssh(self.ip, mkdir_command, self.user, self.password)
+            res = scp(self.ip, self.filename, self.data_send_path, self.user, self.password)
+            if res == 0:
+                print('\nSuccessful file transfer to "%s"\nDeleting local file from pi.' % self.data_send_path)
+                os.remove(self.filename)
+        except Exception as e:
+            res = 1
+            print('scp failed with the following error')
+            print(e)
+        finally:
+            os.chdir(os.path.join(os.getcwd(), '..', '..'))
+            print('file location is back to the project root')
+            GPIO.cleanup()
+            print('\nFile closed and clean up complete')
         if self.halted:
             print('Session stopped early')
         elif self.smooth_finish:
@@ -323,7 +392,7 @@ def lick_status_change(self):
 
 class Task:
     def __init__(self, session, name='blank', structure=None, ports=None, limit='trials',
-                 maximum=None, forgo=True, forced_trials=False):
+                 maximum=None, forgo=False, forced_trials=True):
         print('Starting task: %s' % name)
         self.structure = structure
         self.port_dict = ports
diff --git a/upload_to_pi.py b/upload_to_pi.py
index 3a98b8d..654333d 100644
--- a/upload_to_pi.py
+++ b/upload_to_pi.py
@@ -7,7 +7,7 @@
 def upload_to_pi(pi_host_name, durations=False):
     local_path = os.path.join(os.getcwd(), "stand_alone")
     pi_user_name = 'pi'
-    remote_path = '/home/pi/behavior'
+    remote_path = '/home/pi/rie_behavior'
     password = 'shuler'
     # command = f'scp -r {os.path.join(local_path, "stand_alone")} {pi_user_name}@{pi_host_name}:{remote_path}'
     ssh = paramiko.SSHClient()
@@ -51,6 +51,7 @@ def ping_host(pi_host_name):
     try:
         ssh.connect(pi_host_name, username=pi_user_name, password=password, timeout=1)
         ssh.close()
+        print(f'{pi_host_name} ping host succeeded')
         return True
     except Exception as e:
         ssh.close()
@@ -61,6 +62,11 @@ def ping_host(pi_host_name):
 
 if __name__ == '__main__':
     info_dict = get_user_info()
+    # for name in ['shichenpi2']:
     for name in info_dict['pi_names']:
-        upload_to_pi(name, durations=False)
-        reset_time(name)
+        ping_host(name)
+        try:
+            upload_to_pi(name, durations=False)
+            reset_time(name)
+        except Exception as e:
+            print(f'{name}: {e}')

From 4bf9b2e5960d19732f3659661760ba03c6e3073e Mon Sep 17 00:00:00 2001
From: riekane <rkaneko2@jh.edu>
Date: Fri, 23 May 2025 14:21:38 -0400
Subject: [PATCH 3/3] Rie specific changes

---
 stand_alone/user_settings.py         | 78 ++++++++++++++++++++++++++++
 stand_alone/user_settings_generic.py |  4 ++
 upload_cloud_backup.py               | 40 ++++++++++++++
 user_info.py                         |  9 ++--
 4 files changed, 127 insertions(+), 4 deletions(-)
 create mode 100644 stand_alone/user_settings.py
 create mode 100644 upload_cloud_backup.py

diff --git a/stand_alone/user_settings.py b/stand_alone/user_settings.py
new file mode 100644
index 0000000..88ede77
--- /dev/null
+++ b/stand_alone/user_settings.py
@@ -0,0 +1,78 @@
+import os
+
+"""
+Copy this file to a new one with the name 'user_settings.py' in the same folder location as this one. Replace all BOLD_CASE
+variables with your own. GOOD LUCK!
+"""
+
+
+def get_user_info():
+    info_dict = {
+        'initials': 'RK',
+        # 'mouse_buttons': [['RK001', 'RK002', 'RK003'],
+        #                   ['RK004', 'RK005', 'RK006'],
+        #                   ['testmouse1', 'testmouse2', 'testmouse3'],
+        #                   ],
+        'mouse_buttons': [['RK007', 'RK008', 'RK009'],
+                          ['RK010', 'testmouse', 'testmouse'],
+                          ['testmouse1', 'testmouse2', 'testmouse3'],
+                          ],
+        'button_colors': [[12, 12, 12],
+                          [12, 12, 12],
+                          [2, 2, 2],
+                          [2, 3, 3]
+                          ],
+        'mouse_assignments': {
+            # 'RK001': 'multi_reward',
+            # 'RK002': 'multi_reward',
+            # 'RK003': 'multi_reward',
+            # 'RK004': 'multi_reward',
+            # 'RK005': 'multi_reward',
+            # 'RK006': 'multi_reward',
+            'RK007': 'multi_reward',
+            'RK008': 'multi_reward',
+            'RK009': 'multi_reward',
+            'RK010': 'multi_reward',
+            'testmouse1': 'multi_reward'
+        },
+        'port_swap_assignments': {
+            'RK007': True,
+            'RK008': False,
+            'RK009': True,
+            'RK010': False,
+            'testmouse1': False
+        },
+        'block_swap_assignments': { #0.4 if swap_block is false, 0.8 if true.
+            'RK007': 'trial',
+            'RK008': 'trial',
+            'RK009': 'trial',
+            'RK010': 'trial',
+            'testmouse1': 'trial'
+        },
+        'mouse_colors': {
+            # 'RK001': 1,
+            # 'RK002': 1,
+            # 'RK003': 1,
+            # 'RK004': 2,
+            # 'RK005': 2,
+            # 'RK006': 2,
+            # 'testmouse1': 3,
+            # 'testmouse2': 3,
+            # 'testmouse3': 3,
+            'RK007': 1,
+            'RK008': 1,
+            'RK009': 1,
+            'RK010': 1,
+        },
+        # 'desktop_ip': '192.168.0.150',
+        # 'desktop_user': 'rie',
+        # 'desktop_password': 'shuler',
+        # 'desktop_user_root': os.path.join('C:/', 'Users', 'Rie'),
+        # 'desktop_save_path': os.path.join('behavior_data'),
+        'desktop_ip': '10.203.164.9',
+        'desktop_user': 'Shichen',
+        'desktop_password': 'shuler_914WBSB',
+        'desktop_user_root': os.path.join('C:/', 'Users', 'Shichen'),
+        'desktop_save_path': os.path.join('OneDrive - Johns Hopkins', 'ShulerLab', 'Rie_behavior', 'data'),
+    }
+    return info_dict
diff --git a/stand_alone/user_settings_generic.py b/stand_alone/user_settings_generic.py
index fb41112..f5ef573 100644
--- a/stand_alone/user_settings_generic.py
+++ b/stand_alone/user_settings_generic.py
@@ -16,6 +16,10 @@ def get_user_info():
             'MOUSE_NAME': 'TASK_NAME',
             'testmouse': 'TASK_NAME',
         },
+        'mouse_swap_assignments': {
+            'MOUSE_NAME': False,
+            'testmouse': False,
+        },
         'mouse_colors': {
             'MOUSE_NAME': 1,
             'testmouse': 1,
diff --git a/upload_cloud_backup.py b/upload_cloud_backup.py
new file mode 100644
index 0000000..f9ca9da
--- /dev/null
+++ b/upload_cloud_backup.py
@@ -0,0 +1,40 @@
+import os
+import shutil
+import time
+from pathlib import Path
+
+# Define local and OneDrive paths
+behavior_data_dir = Path("C:/Users/Shichen/OneDrive - Johns Hopkins/ShulerLab/Rie_behavior/data")
+# onedrive_dir = Path("C:/Users/Valued Customer/OneDrive - Johns Hopkins/behavior_data")
+# onedrive_dir = Path("C:/Users/Rie/OneDrive - Johns Hopkins/behavior_data")
+
+
+def upload_cloud_backup():
+    while True:
+        try:
+            # Walk through the behavior_data directory and identify files and folders
+            for root, dirs, files in os.walk(behavior_data_dir):
+                # Calculate the relative path from the base behavior_data directory
+                relative_path = Path(root).relative_to(behavior_data_dir)
+                # Create the corresponding subdirectory path in the onedrive backup folder
+                backup_subdir = onedrive_dir / relative_path
+                backup_subdir.mkdir(parents=True, exist_ok=True)  # Create directories if needed
+
+                # Copy each file if it's not already in the backup folder
+                for file_name in files:
+                    src_file = Path(root) / file_name
+                    dest_file = backup_subdir / file_name
+
+                    # Only copy if the file doesn't exist in the destination
+                    if not dest_file.exists():
+                        shutil.copy2(src_file, dest_file)
+                        print(f"Copied {src_file} to {dest_file}")
+
+        except Exception as e:
+            print(f"Error during backup: {e}")
+
+        # Wait 24 hours before the next check
+        time.sleep(3600 * 24)
+
+if __name__ == '__main__':
+    upload_cloud_backup()
diff --git a/user_info.py b/user_info.py
index e11feb9..a393c3e 100644
--- a/user_info.py
+++ b/user_info.py
@@ -3,9 +3,10 @@
 
 def get_user_info():
     info_dict = {
-        'initials': 'ES',
-        'pi_names': ['elissapi0', 'elissapi2', 'elissapi4'],
-        # 'pi_names': ['elissapi1'],
-        'start_date': date(2023, 10, 11)  # For the current cohort, for the sake of simple_plots.py
+        'initials': ['RK','ES'],
+        # 'pi_names': ['elissapi0', 'elissapi4'],
+        # 'pi_names': ['elissapi1', 'shichenpi2', 'shichenpi3','elissapi0', 'elissapi4'],
+        'pi_names': ['elissapi1','shichenpi2'  , 'shichenpi3'],
+        'start_date': date(2024, 12, 1)  # For the current cohort, for the sake of simple_plots.py
     }
     return info_dict