{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "925409c8",
   "metadata": {},
   "source": [
    "# Using local data paths with `zea`\n",
    "\n",
    "Most zea examples use Hugging Face links for convenience, but you can also work with local datasets by configuring a `users.yaml` file that points to your data root. This notebook shows how to set up local paths and load data from your own storage."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8eb21ff",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tue-bmd/zea/blob/main/docs/source/notebooks/data/zea_local_data.ipynb)\n",
    " \n",
    "[![View on GitHub](https://img.shields.io/badge/GitHub-View%20Source-blue?logo=github)](https://github.com/tue-bmd/zea/blob/main/docs/source/notebooks/data/zea_local_data.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f6e1f554",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install zea"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "613908ac",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "config_picmus_rf = \"hf://zeahub/configs/config_picmus_rf.yaml\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c5f1be2",
   "metadata": {},
   "source": [
    "## Setting up your `users.yaml`\n",
    "\n",
    "Many codebases and projects are littered with hardcoded absolute paths, which can make it difficult to share code or run it on different machines. To avoid this, zea makes use of a `users.yaml` file to define local data paths. The idea is that users can specify a local data root, and zea will use this to resolve paths dynamically, relative to the user's data root.\n",
    "\n",
    "Create a `users.yaml` file in your project directory. This file tells zea where your local data is stored. Example content:\n",
    "\n",
    "```yaml\n",
    "data_root: /home/your_username/data\n",
    "```\n",
    "\n",
    "Replace `/home/your_username/data` with the actual path to your data directory. You can add an `output` key too, for where zea writes results, but it is optional and can be left out entirely.\n",
    "\n",
    "> **Tip:** You can auto-generate this file by running:\n",
    "> ```\n",
    "> zea datapaths\n",
    "> ```\n",
    "> and following the prompts. It asks where your data lives, and writes the answers to `users.yaml` for you."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dba49a0",
   "metadata": {},
   "source": [
    "So that the rest of this notebook runs anywhere, the next cell writes such a file for you, pointing at a folder it creates alongside this notebook. If you already have a `users.yaml` here, it is left alone and used as is.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "61c4756d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wrote users.yaml, pointing at /content/zea-data\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import yaml\n",
    "\n",
    "users_yaml = Path(\"users.yaml\")\n",
    "data_root = Path(\"zea-data\").resolve()\n",
    "\n",
    "if users_yaml.exists():\n",
    "    print(f\"Using the {users_yaml} that is already here.\")\n",
    "else:\n",
    "    # zea only warns about paths that do not exist, so create the folder as well.\n",
    "    data_root.mkdir(parents=True, exist_ok=True)\n",
    "    users_yaml.write_text(yaml.safe_dump({\"data_root\": str(data_root)}), encoding=\"utf-8\")\n",
    "    print(f\"Wrote {users_yaml}, pointing at {data_root}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60e92c98",
   "metadata": {},
   "source": [
    "## Using Local Data Paths\n",
    "\n",
    "Once your `users.yaml` is set up, you can load data from your local data root. Here's a minimal example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5225fcd8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: Using backend 'jax'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔔 Hi root! You are using data from /content/zea-data\n"
     ]
    }
   ],
   "source": [
    "from zea import set_data_paths\n",
    "\n",
    "user = set_data_paths(\"users.yaml\")\n",
    "\n",
    "data_root = user.data_root\n",
    "username = user.username\n",
    "\n",
    "print(f\"🔔 Hi {username}! You are using data from {data_root}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f5f4a1c",
   "metadata": {},
   "source": [
    "## Resolving Dataset Paths\n",
    "\n",
    "With a data root in place you no longer need absolute paths in your code or configs. `zea` resolves paths with `format_data_path`, which is also what a zea config uses for its `data.path` entry:\n",
    "\n",
    "- a **relative** path is resolved against your `data_root`,\n",
    "- an **absolute** path is used as is,\n",
    "- a **Hugging Face** path (`hf://...`) is passed through untouched.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3c779bd0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/content/zea-data/camus/val/patient0401_4CH_half_sequence.hdf5\n",
      "/mnt/other_disk/patient0401_4CH_half_sequence.hdf5\n",
      "hf://zeahub/camus-sample/val/patient0401_4CH_half_sequence.hdf5\n"
     ]
    }
   ],
   "source": [
    "from zea.datapaths import format_data_path\n",
    "\n",
    "# Relative paths are resolved against the user's data_root\n",
    "print(format_data_path(\"camus/val/patient0401_4CH_half_sequence.hdf5\", user))\n",
    "\n",
    "# Absolute paths and Hugging Face paths are returned unchanged, no user needed\n",
    "print(format_data_path(\"/mnt/other_disk/patient0401_4CH_half_sequence.hdf5\"))\n",
    "print(format_data_path(\"hf://zeahub/camus-sample/val/patient0401_4CH_half_sequence.hdf5\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebec6ada",
   "metadata": {},
   "source": [
    "## Advanced Data Path Configuration\n",
    "\n",
    "In the above example, we use the most simple configuration in `users.yaml`, with just a `data_root` key. However, there are many more advanced options you can configure using `users.yaml`. For example, you can specify multiple data roots, for different users and machines. Additionally, you can define a path for local and remote data (if you use for instance a remote storage). Let's have a look at a more advanced example."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0727e2fd",
   "metadata": {},
   "source": [
    "### Example: Complex `users.yaml` Layout\n",
    "\n",
    "For collaborative projects or when working across multiple machines and operating systems, you can use a more structured `users.yaml` file. Here is an example:\n",
    "\n",
    "```yaml\n",
    "alice:\n",
    "  workstation1:\n",
    "    system: linux\n",
    "    data_root:\n",
    "      local: /mnt/data/alice\n",
    "      remote: /mnt/remote/alice\n",
    "    output: /mnt/data/alice/output\n",
    "  laptop:\n",
    "    system: windows\n",
    "    data_root: D:/data/alice\n",
    "    output: D:/data/alice/output\n",
    "bob:\n",
    "  server:\n",
    "    system: linux\n",
    "    data_root:\n",
    "      local: /mnt/data/bob\n",
    "      remote: /mnt/remote/bob\n",
    "  system: linux\n",
    "  data_root: /mnt/data/bob\n",
    "  output: /mnt/data/bob/output\n",
    "# Default fallback if no user/machine matches\n",
    "data_root: /mnt/shared/data\n",
    "output: /mnt/shared/output\n",
    "```\n",
    "\n",
    "- Each user can have different machines, each with their own `system` and `data_root`.\n",
    "- `data_root` can be a string or a dictionary with `local` and `remote` keys. Pick between the two with `set_data_paths(..., local=True/False)`.\n",
    "- A machine specific `data_root` takes precedence over a user specific one, which in turn takes precedence over the userless `data_root` at the bottom.\n",
    "- If nothing matches, zea warns and falls back to a default path for your operating system, so you know your `users.yaml` still needs an entry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7fa9c2d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Remote data root: /content/zea-data\n",
      "Local data root: /content/zea-data\n"
     ]
    }
   ],
   "source": [
    "# Example: Select remote data root (if defined in users.yaml)\n",
    "user_remote = set_data_paths(\"users.yaml\", local=False)\n",
    "print(\"Remote data root:\", user_remote.data_root)\n",
    "user_local = set_data_paths(\"users.yaml\", local=True)\n",
    "print(\"Local data root:\", user_local.data_root)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f9f007",
   "metadata": {},
   "source": [
    "## Full Environment Setup with `setup`\n",
    "\n",
    "For convenience, zea provides a `setup` function that configures everything in one step: config, data paths, and device (GPU/CPU).\n",
    "\n",
    "- This will prompt for missing user profiles if needed, set up data paths, and initialize the device.\n",
    "- Use this in your main scripts for reproducible and portable setups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a2de03cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from zea.internal.setup_zea import setup\n",
    "\n",
    "# config_path: path to your config YAML file\n",
    "# user_config: path to your users.yaml file\n",
    "config = setup(config_path=config_picmus_rf, user_config=\"users.yaml\")\n",
    "\n",
    "data_root = config.data.user.data_root\n",
    "device = config.device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c967677b",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "- Use `users.yaml` to manage local/remote data roots for different users and systems.\n",
    "- Use `set_data_paths` to resolve your data root dynamically.\n",
    "- Use `format_data_path` to turn a relative dataset path into an absolute one.\n",
    "- For advanced setups, structure `users.yaml` with users, hostnames, and local/remote keys.\n",
    "- Use `setup` for a one-liner to initialize config, data paths, and device."
   ]
  }
 ],
 "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.11"
  },
  "zea": {
   "tags": [
    "data loading"
   ],
   "level": "beginner"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}