{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8f11af0-cf74-47de-a453-67e16400a775",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPARE NO-ATR VS ATR (60D05crosses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ff0e2a6d-d50d-4b80-9b08-d0f4a12de73a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ethoscopy as etho\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from scipy.stats import ttest_ind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7d4afe2a-bb1b-4a72-a463-b2aee6f70417",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing No-ATR Data...\n",
      "Processing ATR Data...\n"
     ]
    }
   ],
   "source": [
    "# 1. LOAD SAVED PICKLE FILES\n",
    "df_no_atr_raw = pd.read_pickle('/home/rdingjin/UASxGal4_2026-05-05/phase6.5.pkl')\n",
    "df_atr_raw = pd.read_pickle('/home/rdingjin/UASxGal4_ATR_2026-05-10/phase7(60D)1.pkl')\n",
    "\n",
    "# 2. HELPER FUNCTION\n",
    "def calculate_delta_sleep(df, base_start, base_end, stim_start, stim_end):\n",
    " \n",
    "    # --- BASELINE ---\n",
    "    base_df = df.t_filter(start_time=base_start, end_time=base_end).copy()\n",
    "    sleep_hours_base = (base_df.groupby('id')['asleep'].sum() * 10) / 3600\n",
    "    base_overall = (sleep_hours_base / ((base_end - base_start) / 24)).reset_index(name='Base_Sleep_24h')\n",
    "    \n",
    "    base_df['zt_hour'] = (base_df['t'] / 3600) % 24\n",
    "    base_df['Phase'] = np.where(base_df['zt_hour'] < 12, 'Day', 'Night')\n",
    "    base_phase = base_df.groupby(['id', 'Phase']).agg(\n",
    "        total_sleep_epochs=('asleep', 'sum'), total_tracking_epochs=('t', 'count')\n",
    "    ).reset_index()\n",
    "    \n",
    "    base_phase['Sleep_Hours'] = (base_phase['total_sleep_epochs'] * 10) / 3600\n",
    "    base_phase['Tracking_Hours'] = (base_phase['total_tracking_epochs'] * 10) / 3600\n",
    "    base_phase['Base_Sleep_12h'] = (base_phase['Sleep_Hours'] / base_phase['Tracking_Hours']) * 12\n",
    "    base_pivot = base_phase.pivot(index='id', columns='Phase', values='Base_Sleep_12h').reset_index()\n",
    "    base_pivot.columns = ['id', 'Base_Day_12h', 'Base_Night_12h']\n",
    "    baseline_stats = pd.merge(base_overall, base_pivot, on='id')\n",
    "    \n",
    "    # --- STIMULUS ---\n",
    "    stim_df = df.t_filter(start_time=stim_start, end_time=stim_end).copy()\n",
    "    sleep_hours_stim = (stim_df.groupby('id')['asleep'].sum() * 10) / 3600\n",
    "    stim_overall = (sleep_hours_stim / ((stim_end - stim_start) / 24)).reset_index(name='Stim_Sleep_24h')\n",
    "    \n",
    "    stim_df['zt_hour'] = (stim_df['t'] / 3600) % 24\n",
    "    stim_df['Phase'] = np.where(stim_df['zt_hour'] < 12, 'Day', 'Night')\n",
    "    stim_phase = stim_df.groupby(['id', 'Phase']).agg(\n",
    "        total_sleep_epochs=('asleep', 'sum'), total_tracking_epochs=('t', 'count')\n",
    "    ).reset_index()\n",
    "    \n",
    "    stim_phase['Sleep_Hours'] = (stim_phase['total_sleep_epochs'] * 10) / 3600\n",
    "    stim_phase['Tracking_Hours'] = (stim_phase['total_tracking_epochs'] * 10) / 3600\n",
    "    stim_phase['Stim_Sleep_12h'] = (stim_phase['Sleep_Hours'] / stim_phase['Tracking_Hours']) * 12\n",
    "    stim_pivot = stim_phase.pivot(index='id', columns='Phase', values='Stim_Sleep_12h').reset_index()\n",
    "    stim_pivot.columns = ['id', 'Stim_Day_12h', 'Stim_Night_12h']\n",
    "    stim_stats = pd.merge(stim_overall, stim_pivot, on='id')\n",
    "    \n",
    "    # --- MERGE & CALCULATE DELTA ---\n",
    "    combined = pd.merge(baseline_stats, stim_stats, on='id')\n",
    "    unique_meta = df.meta['machine_name'][~df.meta.index.duplicated(keep='first')]\n",
    "    combined['Machine'] = combined['id'].map(unique_meta)\n",
    "    combined['Tube'] = combined['id'].apply(lambda x: str(x).split('|')[-1])\n",
    "    combined['Fly_Label'] = combined['Machine'] + \" - Fly \" + combined['Tube']\n",
    "    combined['Delta_24h'] = combined['Stim_Sleep_24h'] - combined['Base_Sleep_24h']\n",
    "    combined['Delta_Day'] = combined['Stim_Day_12h'] - combined['Base_Day_12h']\n",
    "    combined['Delta_Night'] = combined['Stim_Night_12h'] - combined['Base_Night_12h']\n",
    "    \n",
    "    return combined[['Fly_Label', 'Delta_24h', 'Delta_Day', 'Delta_Night']].dropna()\n",
    "\n",
    "# 3. PROCESS GROUPS\n",
    "print(\"Processing No-ATR Data...\")\n",
    "delta_no_atr = calculate_delta_sleep(df_no_atr_raw, base_start=0, base_end=48, stim_start=48, stim_end=72)\n",
    "delta_no_atr['Condition'] = 'No-ATR'\n",
    "\n",
    "print(\"Processing ATR Data...\")\n",
    "delta_atr = calculate_delta_sleep(df_atr_raw, base_start=0, base_end=48, stim_start=48, stim_end=72)\n",
    "delta_atr['Condition'] = 'ATR'\n",
    "\n",
    "master_delta = pd.concat([delta_no_atr, delta_atr], ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "157772e1-0537-4cd6-8f0f-620f87118272",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4. MELT DATA FOR PLOTTING AND STATS\n",
    "melted_delta = master_delta.melt(\n",
    "    id_vars=['Fly_Label', 'Condition'],\n",
    "    value_vars=['Delta_24h', 'Delta_Day', 'Delta_Night'],\n",
    "    var_name='Phase', \n",
    "    value_name='Sleep_Change'\n",
    ")\n",
    "# Clean up phase names\n",
    "melted_delta['Phase'] = melted_delta['Phase'].map({'Delta_24h': '24h Total', 'Delta_Day': 'Day', 'Delta_Night': 'Night'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "89badf41-b028-47d3-931a-d314696047ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "==================================================\n",
      "--- SCIPY T-TESTS (No-ATR vs ATR Delta) ---\n",
      "==================================================\n",
      "24h Total Phase Comparison:\n",
      "  No-ATR Δ Mean: -0.62 hours\n",
      "  ATR Δ Mean:    -3.82 hours\n",
      "  P-value:       7.2160e-04\n",
      "\n",
      "Day Phase Comparison:\n",
      "  No-ATR Δ Mean: -0.75 hours\n",
      "  ATR Δ Mean:    -1.37 hours\n",
      "  P-value:       2.9477e-01\n",
      "\n",
      "Night Phase Comparison:\n",
      "  No-ATR Δ Mean: -0.61 hours\n",
      "  ATR Δ Mean:    -4.25 hours\n",
      "  P-value:       2.3451e-08\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 5. STATISTICAL COMPARISON \n",
    "print(\"\\n\" + \"=\"*50)\n",
    "print(\"--- SCIPY T-TESTS (No-ATR vs ATR Delta) ---\")\n",
    "print(\"=\"*50)\n",
    "\n",
    "p_values = {}\n",
    "phases = ['24h Total', 'Day', 'Night']\n",
    "\n",
    "for phase in phases:\n",
    "    group_no_atr = melted_delta[(melted_delta['Phase'] == phase) & (melted_delta['Condition'] == 'No-ATR')]['Sleep_Change']\n",
    "    group_atr = melted_delta[(melted_delta['Phase'] == phase) & (melted_delta['Condition'] == 'ATR')]['Sleep_Change']\n",
    "    \n",
    "    # Run Independent T-Test\n",
    "    t_stat, p_val = ttest_ind(group_no_atr, group_atr)\n",
    "    p_values[phase] = p_val\n",
    "    \n",
    "    print(f\"{phase} Phase Comparison:\")\n",
    "    print(f\"  No-ATR Δ Mean: {group_no_atr.mean():.2f} hours\")\n",
    "    print(f\"  ATR Δ Mean:    {group_atr.mean():.2f} hours\")\n",
    "    print(f\"  P-value:       {p_val:.4e}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "121c90cd-3188-4fbf-86a9-d9e2dcbd9a20",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_50770/148058994.py:14: FutureWarning: \n",
      "\n",
      "Setting a gradient palette using color= is deprecated and will be removed in v0.14.0. Set `palette='dark:black'` for the same effect.\n",
      "\n",
      "  sns.stripplot(\n"
     ]
    }
   ],
   "source": [
    "# 6. FINAL COMPARISON PLOT\n",
    "plt.figure(figsize=(8, 6))\n",
    "sns.set_theme(style=\"ticks\")\n",
    "\n",
    "# Plot Boxplot\n",
    "ax = sns.boxplot(\n",
    "    data=melted_delta, x='Phase', y='Sleep_Change', hue='Condition', \n",
    "    palette={'No-ATR': '#cccccc', 'ATR': '#f1c232'}, order=phases,\n",
    "    width=0.6, gap=0.2, boxprops={'edgecolor': 'black', 'alpha': 0.7},\n",
    "    medianprops={'color': 'black', 'linewidth': 1.5}, showfliers=False\n",
    ")\n",
    "\n",
    "# Plot internal points\n",
    "sns.stripplot(\n",
    "    data=melted_delta, x='Phase', y='Sleep_Change', hue='Condition', \n",
    "    order=phases, dodge=True,\n",
    "    color='black', alpha=0.5, size=4, jitter=0.05, legend=False\n",
    ")\n",
    "\n",
    "# Add Median Labels\n",
    "medians = melted_delta.groupby(['Phase', 'Condition'], sort=False)['Sleep_Change'].median()\n",
    "offsets = {'No-ATR': -0.4, 'ATR': 0.4}\n",
    "\n",
    "for i, phase in enumerate(phases):\n",
    "    for condition in ['No-ATR', 'ATR']:\n",
    "        val = medians[(phase, condition)]\n",
    "        ax.text(i + offsets[condition], val, f\"{val:.1f}h\", \n",
    "                ha='center', va='bottom', fontsize=9)\n",
    "\n",
    "# Add Significance Brackets\n",
    "y_max = melted_delta['Sleep_Change'].max()\n",
    "y_min = melted_delta['Sleep_Change'].min()\n",
    "y_range = y_max - y_min\n",
    "y_line = y_max + (y_range * 0.05)\n",
    "h = y_range * 0.02\n",
    "\n",
    "for i, phase in enumerate(phases):\n",
    "    p = p_values[phase]\n",
    "    \n",
    "    # Define the label\n",
    "    if p < 0.0001: star = '****'\n",
    "    elif p < 0.001: star = '***'\n",
    "    elif p < 0.01: star = '**'\n",
    "    elif p < 0.05: star = '*'\n",
    "    else: star = 'ns' # Define 'ns' for non-significant\n",
    "    \n",
    "    # Draw bracket and text for ALL phases (remove the 'if p < 0.05' wrapper)\n",
    "    x1, x2 = i - 0.2, i + 0.2\n",
    "    plt.plot([x1, x1, x2, x2], [y_line, y_line+h, y_line+h, y_line], lw=1.5, c='k')\n",
    "    plt.text((x1+x2)*.5, y_line+h, star, ha='center', va='bottom', fontsize=12)\n",
    "\n",
    "plt.axhline(0, color='black', linestyle='--', linewidth=1.5, alpha=0.7)\n",
    "sns.despine()\n",
    "plt.ylabel('Change in average time spent asleep (Δ hours)', fontsize=12)\n",
    "plt.xlabel('')\n",
    "plt.xlim(-0.5, 2.5) \n",
    "\n",
    "plt.legend(title='', frameon=False, loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "plt.ylim(y_min - (y_range * 0.1), y_max + (y_range * 0.25))\n",
    "plt.tight_layout()\n",
    "\n",
    "plt.savefig(\"ATR_vs_NoATR_Delta-2.png\", dpi=600, bbox_inches='tight')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}