{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Probability Theory and Probabilistic Models\n",
    "\n",
    "In this lab, you will apply Bayesian reasoning and Monte Carlo simulation to behavioral assessment data. The primary exercise involves implementing Bayesian updating for functional assessment: given data from a functional analysis, you will quantify how evidence accumulates in favor of one behavioral function over others.\n",
    "\n",
    "## Background\n",
    "\n",
    "A standard functional analysis (FA; Iwata et al., 1982/1994) arranges four conditions -- attention, escape, tangible, and play (control) -- to identify the maintaining variable for problem behavior. Traditionally, clinicians rely on visual analysis to determine which condition produces elevated responding. A Bayesian approach offers a principled, quantitative alternative: we start with prior beliefs about the function (e.g., uniform across all four) and update those beliefs after each session based on the observed rates.\n",
    "\n",
    "In the second part of the lab, you will use Monte Carlo simulation to estimate the sampling distribution and confidence intervals for a behavioral parameter, illustrating how simulation-based methods complement analytic approaches."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 1: Import Libraries\n",
    "\n",
    "Import the libraries you will need: `pandas`, `numpy`, `matplotlib.pyplot`, and `scipy.stats`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.stats as stats\n",
    "\n",
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 2: Load and Inspect the FA Data\n",
    "\n",
    "Load `functional_analysis_data.csv` into a DataFrame. Examine the structure of the data.\n",
    "\n",
    "- How many sessions are there for each condition?\n",
    "- What are the mean rates per minute for each condition?\n",
    "- Based on visual inspection of the summary statistics alone, which function appears most likely?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"functional_analysis_data.csv\")\n",
    "print(\"Shape:\", df.shape)\n",
    "print(df.head())\n",
    "\n",
    "print(\"\\nSessions per condition:\")\n",
    "print(df['condition'].value_counts())\n",
    "\n",
    "print(\"\\nMean rate per minute by condition:\")\n",
    "print(df.groupby('condition')['rate_per_min'].agg(['mean', 'std', 'count']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 3: Traditional Visual Analysis\n",
    "\n",
    "Create a standard FA graph: plot rate per minute (y-axis) as a function of session number (x-axis), with different markers and colors for each condition. Connect data points within each condition with lines.\n",
    "\n",
    "This is the traditional approach to interpreting FA data. Note which condition shows consistently elevated responding."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "conditions = ['attention', 'escape', 'tangible', 'play']\n",
    "markers = {'attention': 'o', 'escape': 's', 'tangible': '^', 'play': 'D'}\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(11, 6))\n",
    "for cond in conditions:\n",
    "    sub = df[df['condition'] == cond].sort_values('session')\n",
    "    ax.plot(sub['session'], sub['rate_per_min'],\n",
    "            marker=markers[cond], label=cond)\n",
    "ax.set_xlabel(\"Session\")\n",
    "ax.set_ylabel(\"Rate per minute\")\n",
    "ax.set_title(\"Functional Analysis: Rate by Session and Condition\")\n",
    "ax.legend(title=\"Condition\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 4: Set Up the Prior Distribution\n",
    "\n",
    "Define a uniform prior probability distribution over four possible behavioral functions: `attention`, `escape`, `tangible`, and `automatic`.\n",
    "\n",
    "Store this as a dictionary mapping each function name to its prior probability. Since we are starting with no prior information, each function should have equal probability.\n",
    "\n",
    "Create a bar plot showing the prior distribution."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "functions = ['attention', 'escape', 'tangible', 'automatic']\n",
    "prior = {f: 1.0 / len(functions) for f in functions}\n",
    "print(\"Prior:\", prior)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 4))\n",
    "ax.bar(prior.keys(), prior.values(), color='gray')\n",
    "ax.set_ylim(0, 1)\n",
    "ax.set_ylabel(\"P(function)\")\n",
    "ax.set_title(\"Prior Distribution (uniform)\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 5: Define the Likelihood Function\n",
    "\n",
    "To perform Bayesian updating, you need a likelihood function: $P(\\text{data} | \\text{function})$. Here is one reasonable approach:\n",
    "\n",
    "For a given session, if the observed condition is the one matching the hypothesized function, the likelihood is proportional to the observed rate (higher rates = stronger evidence for that function). If the observed condition is the **play** (control) condition, a high rate provides weak evidence against all specific functions. If the observed condition does not match the hypothesized function, the likelihood should be lower.\n",
    "\n",
    "Implement a likelihood function with the following logic:\n",
    "- If the session condition matches the hypothesized function, likelihood = `rate + 0.1` (add a small constant to avoid zeros)\n",
    "- If the session condition is play, likelihood = `1.0` (uninformative)\n",
    "- If the session condition does not match and is not play, likelihood = `1.0 / (rate + 0.1)` (high rates in other conditions are weak evidence against)\n",
    "\n",
    "**Note:** This is a simplified likelihood for pedagogical purposes. Research applications would use more principled statistical models."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def likelihood(function, condition, rate):\n",
    "    \"\"\"P(data | function) for a single session under the simplified model.\"\"\"\n",
    "    if condition == function:\n",
    "        return rate + 0.1\n",
    "    elif condition == 'play':\n",
    "        return 1.0\n",
    "    else:\n",
    "        return 1.0 / (rate + 0.1)\n",
    "\n",
    "# Quick sanity check: an attention session with a high rate should most favor 'attention'\n",
    "print({f: round(likelihood(f, 'attention', 8.4), 3) for f in functions})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 6: Implement Bayesian Updating\n",
    "\n",
    "Write a function that takes the current prior, a single session's data (condition and rate), and returns the updated posterior using Bayes' theorem:\n",
    "\n",
    "$$P(\\text{function} | \\text{data}) = \\frac{P(\\text{data} | \\text{function}) \\cdot P(\\text{function})}{P(\\text{data})}$$\n",
    "\n",
    "where $P(\\text{data}) = \\sum_{f} P(\\text{data} | f) \\cdot P(f)$ is the normalizing constant.\n",
    "\n",
    "Then, loop through all sessions in chronological order, updating the posterior after each session. Store the posterior after each update so you can plot how beliefs evolve over time.\n",
    "\n",
    "**Hint:** The posterior from session $n$ becomes the prior for session $n+1$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def bayesian_update(current, condition, rate):\n",
    "    \"\"\"Return the posterior dict after observing one session.\"\"\"\n",
    "    unnormalized = {f: likelihood(f, condition, rate) * current[f] for f in current}\n",
    "    evidence = sum(unnormalized.values())\n",
    "    return {f: unnormalized[f] / evidence for f in current}\n",
    "\n",
    "# Loop through sessions in chronological order, carrying the posterior forward\n",
    "posterior = dict(prior)\n",
    "history = [dict(session=0, **posterior)]\n",
    "for _, row in df.sort_values('session').iterrows():\n",
    "    posterior = bayesian_update(posterior, row['condition'], row['rate_per_min'])\n",
    "    history.append(dict(session=int(row['session']), **posterior))\n",
    "\n",
    "history_df = pd.DataFrame(history)\n",
    "print(history_df.tail())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 7: Plot the Evolution of Posterior Beliefs\n",
    "\n",
    "Create a figure showing how the posterior probability for each function changes across sessions. The x-axis should be session number, and the y-axis should be posterior probability (0 to 1). Plot a separate line for each function.\n",
    "\n",
    "At what point does the Bayesian analysis converge on the correct function? How does this compare to how many sessions a visual analyst might need?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(11, 6))\n",
    "for f in functions:\n",
    "    ax.plot(history_df['session'], history_df[f], marker='o', markersize=3, label=f)\n",
    "ax.set_xlabel(\"Session\")\n",
    "ax.set_ylabel(\"Posterior probability\")\n",
    "ax.set_ylim(0, 1)\n",
    "ax.set_title(\"Evolution of Posterior Beliefs Across Sessions\")\n",
    "ax.legend(title=\"Function\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 8: Final Posterior and Comparison\n",
    "\n",
    "Display the final posterior distribution as a bar plot. Place it side-by-side with the prior distribution so the shift in beliefs is clear.\n",
    "\n",
    "Report the final posterior probabilities for each function. What is the Bayes factor comparing the most likely function to the next most likely? (The Bayes factor is the ratio of their posterior odds.)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4), sharey=True)\n",
    "ax1.bar(prior.keys(), prior.values(), color='gray')\n",
    "ax1.set_title(\"Prior\"); ax1.set_ylabel(\"Probability\"); ax1.set_ylim(0, 1)\n",
    "ax2.bar(posterior.keys(), posterior.values(), color='steelblue')\n",
    "ax2.set_title(\"Final Posterior\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"Final posterior:\")\n",
    "for f, p in sorted(posterior.items(), key=lambda kv: -kv[1]):\n",
    "    print(f\"  {f}: {p:.4f}\")\n",
    "\n",
    "ranked = sorted(posterior.values(), reverse=True)\n",
    "bayes_factor = ranked[0] / ranked[1]\n",
    "print(f\"\\nBayes factor (top vs next most likely): {bayes_factor:.1f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 9: Monte Carlo Simulation for Confidence Intervals\n",
    "\n",
    "Now shift to a different application of probabilistic reasoning. Suppose you want to estimate the mean rate of problem behavior in the attention condition and construct a confidence interval, but you want to use simulation rather than analytic formulas.\n",
    "\n",
    "Implement a bootstrap Monte Carlo procedure:\n",
    "\n",
    "1. Extract the observed rates from the attention condition.\n",
    "2. Set the number of bootstrap samples to 10,000.\n",
    "3. For each bootstrap iteration, resample the attention-condition rates **with replacement** (same sample size as original) and compute the mean.\n",
    "4. Store all 10,000 bootstrap means.\n",
    "5. Compute the 95% confidence interval using the 2.5th and 97.5th percentiles of the bootstrap distribution.\n",
    "6. Plot a histogram of the bootstrap means with vertical lines marking the confidence interval bounds and the observed sample mean."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "attention_rates = df.loc[df['condition'] == 'attention', 'rate_per_min'].values\n",
    "n = len(attention_rates)\n",
    "n_boot = 10000\n",
    "\n",
    "boot_means = np.array([\n",
    "    np.random.choice(attention_rates, size=n, replace=True).mean()\n",
    "    for _ in range(n_boot)\n",
    "])\n",
    "\n",
    "ci_low, ci_high = np.percentile(boot_means, [2.5, 97.5])\n",
    "obs_mean = attention_rates.mean()\n",
    "print(f\"Observed mean: {obs_mean:.2f}\")\n",
    "print(f\"95% bootstrap CI: [{ci_low:.2f}, {ci_high:.2f}]\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "ax.hist(boot_means, bins=40, color='steelblue', alpha=0.7)\n",
    "ax.axvline(obs_mean, color='black', lw=2, label=f\"observed mean = {obs_mean:.2f}\")\n",
    "ax.axvline(ci_low, color='red', ls='--', label=f\"2.5% = {ci_low:.2f}\")\n",
    "ax.axvline(ci_high, color='red', ls='--', label=f\"97.5% = {ci_high:.2f}\")\n",
    "ax.set_xlabel(\"Bootstrap mean rate (attention)\")\n",
    "ax.set_ylabel(\"Frequency\")\n",
    "ax.set_title(\"Bootstrap Distribution of the Attention-Condition Mean\")\n",
    "ax.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 10: Interpretation and Discussion\n",
    "\n",
    "In a markdown cell below, address the following questions:\n",
    "\n",
    "1. What are the advantages of Bayesian updating over traditional visual analysis for FA interpretation? What are the disadvantages or limitations?\n",
    "2. How sensitive was the Bayesian analysis to your choice of likelihood function? What would happen if you used a different likelihood specification?\n",
    "3. How does the bootstrap confidence interval for the attention-condition mean compare to what you would get from a parametric approach (e.g., assuming normality)? You may compute both and compare.\n",
    "4. In what applied scenarios might a probabilistic approach to functional assessment be particularly valuable (e.g., when visual analysis is ambiguous, when data are limited)?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}