{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "61dd4252",
   "metadata": {},
   "source": [
    "# Custom Gravity Model Lab\n",
    "Students can input zones, productions, attractions, and travel costs. Fixed beta velues."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6edd4191",
   "metadata": {},
   "source": [
    "## Enter Zone Count and Initial Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63d64c34",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display, clear_output\n",
    "\n",
    "# Function to safely collect matrix input\n",
    "def parse_matrix_input(text, expected_rows):\n",
    "    try:\n",
    "        rows = text.strip().split(\"\\n\")\n",
    "        if len(rows) != expected_rows:\n",
    "            raise ValueError(\"Row count does not match zone count.\")\n",
    "        matrix = [list(map(float, row.strip().split())) for row in rows]\n",
    "        return np.array(matrix)\n",
    "    except:\n",
    "        print(\" Invalid matrix input format.\")\n",
    "        return None\n",
    "\n",
    "zone_count = widgets.BoundedIntText(value=3, min=2, max=10, step=1, description='Number of Zones:')\n",
    "display(zone_count)\n",
    "\n",
    "# Production input\n",
    "prod_input = widgets.Textarea(\n",
    "    value=\"400 350 250\",\n",
    "    placeholder=\"Enter productions separated by space\",\n",
    "    description='Productions:',\n",
    "    layout=widgets.Layout(width='70%', height='50px')\n",
    ")\n",
    "display(prod_input)\n",
    "\n",
    "# Attraction input\n",
    "attr_input = widgets.Textarea(\n",
    "    value=\"300 200 500\",\n",
    "    placeholder=\"Enter attractions separated by space\",\n",
    "    description='Attractions:',\n",
    "    layout=widgets.Layout(width='70%', height='50px')\n",
    ")\n",
    "display(attr_input)\n",
    "\n",
    "# Skim matrix input\n",
    "skim_input = widgets.Textarea(\n",
    "    value=\"5 10 18\\n13 5 15\\n20 16 6\",\n",
    "    placeholder=\"Enter travel times as space-separated rows (one row per line)\",\n",
    "    description='Cost Matrix:',\n",
    "    layout=widgets.Layout(width='70%', height='100px')\n",
    ")\n",
    "display(skim_input)\n",
    "\n",
    "start_button = widgets.Button(description=\"✅ Run Model\", button_style='success')\n",
    "display(start_button)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a87d194b",
   "metadata": {},
   "source": [
    "## Gravity Model Calculation and Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f55ec2a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_gravity_model(_):\n",
    "    n = zone_count.value\n",
    "    try:\n",
    "        P = np.array(list(map(float, prod_input.value.strip().split())))\n",
    "        A = np.array(list(map(float, attr_input.value.strip().split())))\n",
    "        C = parse_matrix_input(skim_input.value, n)\n",
    "        if len(P) != n or len(A) != n or C is None or C.shape != (n, n):\n",
    "            print(\" Make sure the number of zones matches all input dimensions.\")\n",
    "            return\n",
    "    except:\n",
    "        print(\" Error parsing inputs. Check all values.\")\n",
    "        return\n",
    "\n",
    "    beta = -0.035 # FIXED BETA VALUE\n",
    "    F = np.exp(beta * C)\n",
    "    Tij = np.zeros((n, n))\n",
    "    row_errors = []\n",
    "    col_errors = []\n",
    "\n",
    "    # Initial estimate\n",
    "    for i in range(n):\n",
    "        denom = np.sum(F[i, :] * A)\n",
    "        for j in range(n):\n",
    "            Tij[i, j] = P[i] * (F[i, j] * A[j]) / denom\n",
    "\n",
    "    # Furness balancing\n",
    "    tolerance = 1e-4\n",
    "    max_iter = 100\n",
    "    for _ in range(max_iter):\n",
    "        row_errors.append(np.abs(Tij.sum(axis=1) - P).sum())\n",
    "        col_errors.append(np.abs(Tij.sum(axis=0) - A).sum())\n",
    "\n",
    "        col_sums = Tij.sum(axis=0)\n",
    "        for j in range(n):\n",
    "            if col_sums[j] != 0:\n",
    "                Tij[:, j] *= A[j] / col_sums[j]\n",
    "\n",
    "        row_sums = Tij.sum(axis=1)\n",
    "        for i in range(n):\n",
    "            if row_sums[i] != 0:\n",
    "                Tij[i, :] *= P[i] / row_sums[i]\n",
    "\n",
    "        if np.allclose(Tij.sum(axis=0), A, rtol=tolerance) and np.allclose(Tij.sum(axis=1), P, rtol=tolerance):\n",
    "            break\n",
    "\n",
    "    Tij = Tij.round(1)\n",
    "    for i in range(n):\n",
    "        Tij[i, -1] += P[i] - Tij[i, :].sum()\n",
    "    for j in range(n):\n",
    "        Tij[-1, j] += A[j] - Tij[:, j].sum()\n",
    "\n",
    "    df = pd.DataFrame(\n",
    "        Tij.round(1),\n",
    "        index=[f\"From Zone {i+1}\" for i in range(n)],\n",
    "        columns=[f\"To Zone {j+1}\" for j in range(n)]\n",
    "    )\n",
    "    df[\"Row Total\"] = df.sum(axis=1)\n",
    "    df.loc[\"Column Total\"] = df.sum(axis=0)\n",
    "\n",
    "    clear_output(wait=True)\n",
    "    print(\"✅ Model Complete!\")\n",
    "    display(df)\n",
    "\n",
    "    avg_trip_length = (Tij * C).sum() / Tij.sum()\n",
    "    print(f\"Estimated Average Trip Length: {avg_trip_length:.2f}\")\n",
    "\n",
    "    # Plot: Convergence\n",
    "    plt.figure(figsize=(6, 4))\n",
    "    plt.plot(row_errors, label='Row Error')\n",
    "    plt.plot(col_errors, label='Column Error')\n",
    "    plt.title(\"Convergence Plot\")\n",
    "    plt.xlabel(\"Iteration\")\n",
    "    plt.ylabel(\"Error\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    # Plot: Heatmap\n",
    "    plt.figure(figsize=(6, 5))\n",
    "    sns.heatmap(Tij, annot=True, fmt=\".1f\", cmap=\"YlGnBu\")\n",
    "    plt.title(\"Trip Distribution Heatmap\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    # Plot: Friction factor curve\n",
    "    cost_range = np.linspace(0, 30, 100)\n",
    "    F_curve = np.exp(beta * cost_range)\n",
    "    plt.figure(figsize=(6, 4))\n",
    "    plt.plot(cost_range, F_curve, color='purple')\n",
    "    plt.title(\"Friction Factor vs Travel Time\")\n",
    "    plt.xlabel(\"Travel Time (minutes)\")\n",
    "    plt.ylabel(\"F(c)\")\n",
    "    plt.grid(True)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "start_button.on_click(run_gravity_model)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}