{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2d616210",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DATA\n",
    "\n",
    "import pandas as pd\n",
    "import stan\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import arviz as az\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "\n",
    "# stan problems\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c5cdedfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://raw.githubusercontent.com/stedy/Machine-Learning-with-R-datasets/master/insurance.csv\"\n",
    "data = pd.read_csv(url)\n",
    "\n",
    "# god is dead, force a normal\n",
    "data['smoker'] = data['smoker'].map({'yes': 1, 'no': 0})\n",
    "\n",
    "X = data[['bmi', 'age', 'children', 'smoker']].values\n",
    "y = data['charges'].values\n",
    "\n",
    "\n",
    "# Standardize predictors (mean=0, sd=1)\n",
    "X_mean = X.mean(axis=0)\n",
    "X_std = X.std(axis=0)\n",
    "y_mean = y.mean()\n",
    "y_std = y.std()\n",
    "y = (y - y_mean) / y_std\n",
    "X_std = X.std(axis=0)\n",
    "continuous_vars = data[['bmi', 'age', 'children']]\n",
    "continuous_vars_standardized = (continuous_vars - continuous_vars.mean()) / continuous_vars.std()\n",
    "\n",
    "X_standardized = np.hstack([continuous_vars_standardized, data[['smoker']].values])\n",
    "\n",
    "# Split into train/test\n",
    "X_train, X_test, y_train, y_test = train_test_split(X_standardized, y, test_size=0.2, random_state=42)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1a832a9e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Building: found in cache, done."
     ]
    }
   ],
   "source": [
    "stan_data = {\n",
    "  'N': X_train.shape[0],\n",
    "  'M': X_train.shape[1],\n",
    "  'X': X_train,\n",
    "  'y': y_train,\n",
    "  'N_test': X_test.shape[0],\n",
    "  'X_test': X_test\n",
    "}\n",
    "\n",
    "with open('multiple_regression.stan', 'r') as f:\n",
    "    stan_code = f.read()\n",
    "    \n",
    "posterior = stan.build(stan_code, data=stan_data, random_seed=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "def536a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling:   0%\n",
      "Sampling:  25% (3000/12000)\n",
      "Sampling:  50% (6000/12000)\n",
      "Sampling:  75% (9000/12000)\n",
      "Sampling: 100% (12000/12000)\n",
      "Sampling: 100% (12000/12000), done.\n",
      "Messages received during sampling:\n",
      "  Gradient evaluation took 0.00019 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 1.9 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Informational Message: The current Metropolis proposal is about to be rejected because of the following issue:\n",
      "  Exception: normal_lpdf: Scale parameter is 0, but must be positive! (in '/tmp/httpstan_gquumxii/model_rezrxovk.stan', line 21, column 2 to column 38)\n",
      "  If this warning occurs sporadically, such as for highly constrained variable types like covariance matrices, then the sampler is fine,\n",
      "  but if this warning occurs often then your model may be either severely ill-conditioned or misspecified.\n",
      "  Gradient evaluation took 0.000184 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 1.84 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Gradient evaluation took 0.000168 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 1.68 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Gradient evaluation took 0.000129 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 1.29 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Informational Message: The current Metropolis proposal is about to be rejected because of the following issue:\n",
      "  Exception: normal_lpdf: Scale parameter is 0, but must be positive! (in '/tmp/httpstan_gquumxii/model_rezrxovk.stan', line 21, column 2 to column 38)\n",
      "  If this warning occurs sporadically, such as for highly constrained variable types like covariance matrices, then the sampler is fine,\n",
      "  but if this warning occurs often then your model may be either severely ill-conditioned or misspecified.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "          mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  \\\n",
      "alpha   -0.396  0.017  -0.430   -0.364        0.0      0.0    9346.0   \n",
      "beta[0]  0.165  0.016   0.136    0.195        0.0      0.0   10163.0   \n",
      "beta[1]  0.298  0.016   0.269    0.328        0.0      0.0   10522.0   \n",
      "beta[2]  0.042  0.015   0.014    0.072        0.0      0.0   10397.0   \n",
      "beta[3]  1.951  0.039   1.875    2.020        0.0      0.0    9719.0   \n",
      "sigma    0.507  0.011   0.486    0.527        0.0      0.0    9713.0   \n",
      "\n",
      "         ess_tail  r_hat  \n",
      "alpha      6429.0    1.0  \n",
      "beta[0]    6037.0    1.0  \n",
      "beta[1]    5893.0    1.0  \n",
      "beta[2]    5565.0    1.0  \n",
      "beta[3]    6633.0    1.0  \n",
      "sigma      6303.0    1.0  \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Posterior Distributions of Beta Coefficients')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 21608 characters]",
      "text/plain": [
       "<Figure size 600x440 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "fit = posterior.sample(num_chains=4, num_samples=2000)\n",
    "y_pred_samples = fit['y_pred']\n",
    "\n",
    "# shape --> [num_draws, N_test]\n",
    "# Adjust y_test to match y_pred_samples shape (268, 8000) and compute RMSE per draw <-- Chatgpt did this because everything kep breaking, help\n",
    "rmse_samples = np.sqrt(np.mean((y_pred_samples - y_test[:, None])**2, axis=0))\n",
    "rmse_mean = rmse_samples.mean()\n",
    "rmse_ci = np.percentile(rmse_samples, [2.5, 97.5])\n",
    "\n",
    "idata = az.from_pystan(posterior=fit)\n",
    "print(az.summary(idata, var_names=['alpha', 'beta', 'sigma']))\n",
    "\n",
    "az.plot_forest(idata, var_names=['beta'], combined=True, colors=\"purple\")\n",
    "plt.axvline(0, linestyle='--', color='gray')\n",
    "plt.title(\"Posterior Distributions of Beta Coefficients\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
