{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "04ab3a24",
   "metadata": {},
   "source": [
    "# First Steps\n",
    "\n",
    "This page walks you through your **first complete `braincell` simulation**: a\n",
    "classic Hodgkin–Huxley neuron that we stimulate with a constant current and\n",
    "watch spike. It takes only a few lines and introduces the core workflow you\n",
    "will reuse everywhere.\n",
    "\n",
    "```{tip}\n",
    "Every physical quantity in `braincell` **must** carry an explicit unit from\n",
    "[`brainunit`](https://github.com/chaobrain/brainunit). A bare number like\n",
    "`10` is rejected — write `10 * u.nA` instead. See {doc}`../concepts/units`.\n",
    "```\n",
    "\n",
    "## 1. Define the neuron\n",
    "\n",
    "A single-compartment neuron is a subclass of\n",
    "{class}`braincell.SingleCompartment`. Inside `__init__` you attach **ion\n",
    "species** and the **channels** that carry their currents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "304bb37e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainstate\n",
    "import brainunit as u\n",
    "import braincell\n",
    "\n",
    "\n",
    "class HH(braincell.SingleCompartment):\n",
    "    def __init__(self, size, solver='exp_euler'):\n",
    "        super().__init__(size, solver=solver)\n",
    "\n",
    "        # Sodium: fixed reversal potential + the HH sodium channel\n",
    "        self.na = braincell.ion.SodiumFixed(size, E=50. * u.mV)\n",
    "        self.na.add(INa=braincell.channel.Na_HH1952(size))\n",
    "\n",
    "        # Potassium: fixed reversal potential + the HH potassium channel\n",
    "        self.k = braincell.ion.PotassiumFixed(size, E=-77. * u.mV)\n",
    "        self.k.add(IK=braincell.channel.K_HH1952(size))\n",
    "\n",
    "        # Passive leak\n",
    "        self.IL = braincell.channel.IL(size, E=-54.387 * u.mV,\n",
    "                                       g_max=0.03 * (u.mS / u.cm ** 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7a8e9d9",
   "metadata": {},
   "source": [
    "A few things to notice:\n",
    "\n",
    "- `size` is the number of independent neurons — `braincell` is vectorized, so\n",
    "  `HH(100)` simulates 100 neurons at once.\n",
    "- Channels are *added to* an ion (`self.na.add(...)`) because their current\n",
    "  depends on that ion's reversal potential. The keyword (`INa`, `IK`) names the\n",
    "  channel on the cell.\n",
    "- `solver` selects the numerical integrator by name. See\n",
    "  {doc}`../integration/index` for the full list.\n",
    "\n",
    "## 2. Initialize state\n",
    "\n",
    "`braincell` models are stateful. Allocate and initialize the state (membrane\n",
    "voltage, gating variables, …) before simulating:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92850dc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "hh = HH(1, solver='ind_exp_euler')\n",
    "hh.init_state()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee44e972",
   "metadata": {},
   "source": [
    "## 3. Step the simulation\n",
    "\n",
    "Time stepping is driven by `brainstate`. Each step injects a stimulus current\n",
    "and advances the model by one time step `dt`; we collect the membrane voltage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c87b71e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def step(t):\n",
    "    with brainstate.environ.context(t=t):\n",
    "        hh.update(10. * u.nA / u.cm ** 2)   # injected current density\n",
    "    return hh.V.value\n",
    "\n",
    "\n",
    "with brainstate.environ.context(dt=0.1 * u.ms):\n",
    "    times = u.math.arange(0. * u.ms, 100. * u.ms, brainstate.environ.get_dt())\n",
    "    voltages = brainstate.transform.for_loop(step, times)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c27c96e8",
   "metadata": {},
   "source": [
    "`brainstate.transform.for_loop` compiles the loop with JAX, so the whole 100 ms\n",
    "run executes as a single fused kernel.\n",
    "\n",
    "## 4. Plot the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "473586f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(times.to_decimal(u.ms), u.math.squeeze(voltages).to_decimal(u.mV))\n",
    "plt.xlabel('Time (ms)')\n",
    "plt.ylabel('Membrane potential (mV)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81ec7626",
   "metadata": {},
   "source": [
    "You should see a train of action potentials riding on the injected current.\n",
    "\n",
    "## What you just did\n",
    "\n",
    "```text\n",
    "  define cell  ──▶  init_state()  ──▶  for_loop(step, times)  ──▶  plot V\n",
    "  (ions +                            (advance dt each step,\n",
    "   channels)                          collect a trace)\n",
    "```\n",
    "\n",
    "This is the canonical single-compartment loop. The same `update`/`for_loop`\n",
    "pattern reappears in every single-compartment example.\n",
    "\n",
    "## Next steps\n",
    "\n",
    "- **Understand each piece** — read {doc}`../concepts/cells`,\n",
    "  {doc}`../concepts/ions_channels`, and {doc}`../concepts/units`.\n",
    "- **Explore more single-compartment models** — the\n",
    "  {doc}`../tutorials/index` tutorials cover calcium dynamics,\n",
    "  f–I curves, spike-frequency adaptation, and full thalamic models.\n",
    "- **Add geometry** — when an isopotential point is not enough, move to\n",
    "  morphological cells in {doc}`../tutorials/index`.\n",
    "- **Choose a solver** — for stiff channel kinetics, see\n",
    "  {doc}`../integration/index`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}