{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4be303ee",
   "metadata": {},
   "source": [
    "## Notebook set up\n",
    "\n",
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f194e0fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from sklearn.datasets import make_friedman1\n",
    "from sklearn.linear_model import LinearRegression\n",
    "from sklearn.tree import DecisionTreeRegressor\n",
    "from sklearn.model_selection import train_test_split, cross_val_score\n",
    "from sklearn.metrics import mean_squared_error, root_mean_squared_error, r2_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5db5d65b",
   "metadata": {},
   "source": [
    "### Create dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0d6bb02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the Friedman1 dataset\n",
    "X, y = make_friedman1(n_samples=5000, n_features=5, random_state=315)\n",
    "\n",
    "# Convert to DataFrame\n",
    "X_df = pd.DataFrame(X, columns=[f'feature_{i}' for i in range(5)])\n",
    "y_series = pd.Series(y, name='label')\n",
    "\n",
    "# Combine into a single DataFrame for easy exploration\n",
    "df = X_df.copy()\n",
    "df['label'] = y_series\n",
    "\n",
    "# Display the first few rows\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87753c18",
   "metadata": {},
   "source": [
    "## Task 1: Train and evaluate a linear regression model\n",
    "\n",
    "**Tasks**:\n",
    "\n",
    "1. Split the data into training and testing sets using an 80-20 split. Use `random_state=315`.\n",
    "\n",
    "2. Train a `LinearRegression` model on the training data.\n",
    "\n",
    "3. Calculate and print the following metrics:\n",
    "   - Training RMSE\n",
    "   - Testing RMSE\n",
    "\n",
    "4. Create a scatter plot of true vs. predicted values for the test set:\n",
    "   - Add a diagonal reference line (y=x) to show perfect predictions\n",
    "   - Label the axes appropriately\n",
    "   - Add a title\n",
    "\n",
    "**Hints**:\n",
    "\n",
    "- Use `train_test_split()` with `test_size=0.2`\n",
    "  - Example: `X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=315)`\n",
    "\n",
    "- To calculate metrics:\n",
    "  - First make predictions: `y_pred = model.predict(X_test)`\n",
    "  - Then calculate: `rmse = root_mean_squared_error(y_test, y_pred)`\n",
    "\n",
    "- To add a reference line to a plot:\n",
    "  - `plt.plot([min, max], [min, max], 'k--', alpha=0.3)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c17f022",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52e629e6",
   "metadata": {},
   "source": [
    "## Task 2: Train and evaluate a decision tree model\n",
    "\n",
    "**Tasks**:\n",
    "\n",
    "1. Train the decision tree model (provided below) on the training data from Task 1.\n",
    "\n",
    "2. Calculate and print the following metrics for the decision tree:\n",
    "   - Training RMSE\n",
    "   - Testing RMSE\n",
    "\n",
    "3. Create a scatter plot of true vs. predicted values for the test set:\n",
    "   - Add a diagonal reference line (y=x) to show perfect predictions\n",
    "   - Label the axes appropriately\n",
    "   - Add a title\n",
    "\n",
    "4. Compare the decision tree metrics to the linear regression metrics from Task 1:\n",
    "   - Which model has lower testing RMSE?\n",
    "   - How do the two models differ in their predictions?\n",
    "\n",
    "**Hints**:\n",
    "\n",
    "- Use the same `X_train`, `X_test`, `y_train`, `y_test` from Task 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caa58443",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize and train the decision tree model\n",
    "dt_model = DecisionTreeRegressor(max_depth=8, min_samples_split=30, min_samples_leaf=15, random_state=315)\n",
    "dt_model.fit(X_train, y_train)\n",
    "\n",
    "# Your code here to calculate metrics and create visualizations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07e0a701",
   "metadata": {},
   "source": [
    "## Task 3: Investigate why the models perform differently\n",
    "\n",
    "**Tasks**:\n",
    "\n",
    "1. Create visualizations to understand the relationship between features and label:\n",
    "   - For features 0 and 1: Create a 2D scatter plot colored by the label value (use a colormap)\n",
    "   - For features 2, 3, and 4: Create individual scatter plots vs. label\n",
    "\n",
    "2. Based on your plots:\n",
    "   - Identify which relationships are linear\n",
    "   - Identify which relationships are non-linear\n",
    "   - Explain how this affects each model's performance\n",
    "\n",
    "3. (Optional) Try to improve the linear regression model by adding polynomial features for the non-linear relationships. Does this improve performance?\n",
    "\n",
    "**Hints**:\n",
    "\n",
    "- For a 2D scatter plot with color mapping:\n",
    "  - `plt.scatter(X[:, 0], X[:, 1], c=y, cmap='viridis')`\n",
    "  - `plt.colorbar(label='Label')`\n",
    "\n",
    "- The decision tree can capture non-linear relationships by splitting the feature space, while linear regression assumes linear relationships"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74694d23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edef9a87",
   "metadata": {},
   "source": [
    "## Reflection\n",
    "\n",
    "Based on your analysis, answer the following questions:\n",
    "\n",
    "1. **Model performance**: Which model performed better and why?\n",
    "\n",
    "2. **Linear assumptions**: What happens when you apply linear regression to non-linear data?\n",
    "\n",
    "3. **Model complexity**: What are the trade-offs between simpler models (linear regression) and more complex models (decision trees)?\n",
    "\n",
    "4. **Real-world implications**: In what situations would you prefer:\n",
    "   - A linear regression model?\n",
    "   - A decision tree model?\n",
    "   - Consider factors like interpretability, performance, and data characteristics."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}