{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "03bee28b",
   "source": "# MLflow Manual Logging Tutorial",
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "id": "fc6d0b82",
   "metadata": {},
   "source": [
    "## 1. Environment Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25705aee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Core MLflow imports\n",
    "import mlflow\n",
    "import mlflow.pytorch\n",
    "\n",
    "from datamint import Api\n",
    "from datamint.mlflow import set_project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "198168f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "api = Api()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfa7d7c3",
   "metadata": {},
   "source": [
    "## 2. Project Configuration\n",
    "\n",
    "Before starting any MLflow operations, we must set the active Datamint project. This ensures all experiments are properly associated with your project.\n",
    "\n",
    "**IMPORTANT**: Replace `'BoneSeg'` with your actual project name!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "623a0cae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace this with your actual project name\n",
    "PROJECT_NAME = 'FracAtlas'\n",
    "project = set_project(PROJECT_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64afce13",
   "metadata": {},
   "source": [
    "## 3. Logging Metrics and Model to MLflow\n",
    "\n",
    "Now we'll log the pre-computed metrics and the trained model to MLflow. This demonstrates how to log results after training is complete.\n",
    "\n",
    "### Key MLflow Functions Used:\n",
    "- `mlflow.start_run()`: Start a new MLflow run, representing the past training session\n",
    "- `mlflow.log_params()`: Log hyperparameters\n",
    "- `mlflow.log_metric()`: Log a metric value at a specific step (epoch)\n",
    "- `mlflow.pytorch.log_model()`: Save the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfb910d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assume these are your pre-computed metrics from training\n",
    "# Replace these with your actual training results\n",
    "# You can also log per-epoch metrics if you saved them\n",
    "# For example, if you saved metrics for each epoch during training:\n",
    "epoch_metrics = [\n",
    "    {\"train/loss\": 0.653, \"train/dice_score\": 0.1, \"val/loss\": 0.621, \"val/dice_score\": 0.09},  # Epoch 0\n",
    "    {\"train/loss\": 0.512, \"train/dice_score\": 0.2, \"val/loss\": 0.534, \"val/dice_score\": 0.2},  # Epoch 1\n",
    "    {\"train/loss\": 0.423, \"train/dice_score\": 0.25, \"val/loss\": 0.467, \"val/dice_score\": 0.22},  # Epoch 2\n",
    "    {\"train/loss\": 0.356, \"train/dice_score\": 0.3, \"val/loss\": 0.398, \"val/dice_score\": 0.33},  # Epoch 3\n",
    "    {\"train/loss\": 0.298, \"train/dice_score\": 0.33, \"val/loss\": 0.345, \"val/dice_score\": 0.32},  # Epoch 4\n",
    "    {\"train/loss\": 0.267, \"train/dice_score\": 0.36, \"val/loss\": 0.328, \"val/dice_score\": 0.35},  # Epoch 5\n",
    "    {\"train/loss\": 0.251, \"train/dice_score\": 0.5, \"val/loss\": 0.319, \"val/dice_score\": 0.44},  # Epoch 6\n",
    "    {\"train/loss\": 0.242, \"train/dice_score\": 0.6, \"val/loss\": 0.315, \"val/dice_score\": 0.5},  # Epoch 7\n",
    "    {\"train/loss\": 0.237, \"train/dice_score\": 0.62, \"val/loss\": 0.313, \"val/dice_score\": 0.55},  # Epoch 8\n",
    "    {\"train/loss\": 0.234, \"train/dice_score\": 0.75, \"val/loss\": 0.312, \"val/dice_score\": 0.64},  # Epoch 9\n",
    "]\n",
    "\n",
    "best_metrics = {\n",
    "    \"val/best/loss\": 0.312,\n",
    "    \"val/best/dice_score\": 0.64,\n",
    "}\n",
    "\n",
    "mlflow.set_experiment('nnUnet_training1')\n",
    "# Start MLflow run\n",
    "with mlflow.start_run(run_name=\"log_metrics_run\") as run:\n",
    "    print(f\"🚀 MLflow Run ID: {run.info.run_id}\")\n",
    "    print(f\"Experiment ID: {run.info.experiment_id}\")\n",
    "    print()\n",
    "\n",
    "    # Log hyperparameters\n",
    "    mlflow.log_params({\n",
    "        \"learning_rate\": 1e-4,\n",
    "        \"batch_size\": 16,\n",
    "        \"num_epochs\": 10,\n",
    "        \"model_architecture\": \"nnUnet1\",\n",
    "    })\n",
    "    print(\"✅ Hyperparameters logged to MLflow\")\n",
    "    print()\n",
    "\n",
    "    # Log per-epoch metrics (if available)\n",
    "    print(\"Logging per-epoch metrics...\")\n",
    "    for epoch, metrics in enumerate(epoch_metrics):\n",
    "        mlflow.log_metrics({\n",
    "            \"train/loss\": metrics[\"train/loss\"],\n",
    "            \"train/accuracy\": metrics[\"train/dice_score\"],\n",
    "            \"val/loss\": metrics[\"val/loss\"],\n",
    "            \"val/accuracy\": metrics[\"val/dice_score\"],\n",
    "        }, step=epoch)\n",
    "    print(f\"✅ Logged metrics for {len(epoch_metrics)} epochs\")\n",
    "    print()\n",
    "\n",
    "    # Log final/best metrics\n",
    "    print(\"Logging final metrics...\")\n",
    "    mlflow.log_metrics(best_metrics)\n",
    "    print(\"✅ Final metrics logged\")\n",
    "    print()\n",
    "\n",
    "print(\"=\" * 60)\n",
    "print(\"Logging complete!\")\n",
    "print(f\"View results at: {api.projects.get_by_name(PROJECT_NAME).url}\")\n",
    "print(\"=\" * 60)"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}