{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validation Framework Tutorial\n",
    "\n",
    "This notebook demonstrates the new rich parameter validation framework in epimodels.\n",
    "\n",
    "## Overview\n",
    "\n",
    "The validation framework provides:\n",
    "1. **Declarative Parameter Specifications**: Define parameters with types, bounds, constraints\n",
    "2. **Constraint Language**: Express relationships between parameters\n",
    "3. **Severity Levels**: Error vs warning constraints\n",
    "4. **Backward Compatibility**: Existing models work unchanged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:18:09.729956407Z",
     "start_time": "2026-03-11T18:18:09.227356059Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:18.640330Z",
     "iopub.status.busy": "2026-03-11T19:33:18.639889Z",
     "iopub.status.idle": "2026-03-11T19:33:19.824433Z",
     "shell.execute_reply": "2026-03-11T19:33:19.823438Z"
    }
   },
   "outputs": [],
   "source": [
    "from epimodels import BaseModel, ValidationError\n",
    "from epimodels.validation import (\n",
    "    ParameterSpec,\n",
    "    VariableSpec,\n",
    "    ModelConstraint,\n",
    "    DomainType,\n",
    "    evaluate_constraint,\n",
    ")\n",
    "from collections import OrderedDict\n",
    "import warnings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Parameter Specifications\n",
    "\n",
    "Define parameters with rich metadata including bounds, constraints, units, and documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:18:14.401306456Z",
     "start_time": "2026-03-11T18:18:14.236833048Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.828709Z",
     "iopub.status.busy": "2026-03-11T19:33:19.828340Z",
     "iopub.status.idle": "2026-03-11T19:33:19.833987Z",
     "shell.execute_reply": "2026-03-11T19:33:19.832981Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter: beta\n",
      "Symbol: $\\beta$\n",
      "Description: Transmission rate (contact rate × probability of transmission)\n",
      "Bounds: (0, None)\n",
      "Constraints: ['value > 0']\n",
      "Units: 1/day\n"
     ]
    }
   ],
   "source": [
    "# Create a parameter specification\n",
    "beta_spec = ParameterSpec(\n",
    "    name=\"beta\",\n",
    "    symbol=r\"$\\beta$\",\n",
    "    description=\"Transmission rate (contact rate × probability of transmission)\",\n",
    "    domain_type=DomainType.CONTINUOUS,\n",
    "    bounds=(0, None),  # Must be positive\n",
    "    dtype=float,\n",
    "    required=True,\n",
    "    constraints=[\"value > 0\"],\n",
    "    units=\"1/day\",\n",
    "    typical_range=(0.1, 1.0)\n",
    ")\n",
    "\n",
    "print(f\"Parameter: {beta_spec.name}\")\n",
    "print(f\"Symbol: {beta_spec.symbol}\")\n",
    "print(f\"Description: {beta_spec.description}\")\n",
    "print(f\"Bounds: {beta_spec.bounds}\")\n",
    "print(f\"Constraints: {beta_spec.constraints}\")\n",
    "print(f\"Units: {beta_spec.units}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Creating a Model with Rich Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:18:32.776053911Z",
     "start_time": "2026-03-11T18:18:32.631664844Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.902698Z",
     "iopub.status.busy": "2026-03-11T19:33:19.902285Z",
     "iopub.status.idle": "2026-03-11T19:33:19.916309Z",
     "shell.execute_reply": "2026-03-11T19:33:19.915062Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model type: SIR\n",
      "Parameters: ['beta', 'gamma']\n",
      "Variables: ['S', 'I', 'R']\n",
      "Constraints: 1\n"
     ]
    }
   ],
   "source": [
    "class SIRWithValidation(BaseModel):\n",
    "    \"\"\"SIR model with rich parameter validation.\"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.model_type = \"SIR\"\n",
    "        \n",
    "        # Define parameters with rich specifications\n",
    "        self.define_parameter(ParameterSpec(\n",
    "            name=\"beta\",\n",
    "            symbol=r\"$\\beta$\",\n",
    "            description=\"Transmission rate\",\n",
    "            bounds=(0, None),\n",
    "            constraints=[\"value > 0\"],\n",
    "            units=\"1/day\",\n",
    "            typical_range=(0.1, 1.0)\n",
    "        ))\n",
    "        \n",
    "        self.define_parameter(ParameterSpec(\n",
    "            name=\"gamma\",\n",
    "            symbol=r\"$\\gamma$\",\n",
    "            description=\"Recovery rate (1 / average infectious period)\",\n",
    "            bounds=(0, None),\n",
    "            constraints=[\"value > 0\"],\n",
    "            units=\"1/day\",\n",
    "            typical_range=(0.05, 0.5)\n",
    "        ))\n",
    "        \n",
    "        # Define state variables\n",
    "        self.define_variable(VariableSpec(\n",
    "            name=\"S\",\n",
    "            symbol=\"S\",\n",
    "            description=\"Susceptible individuals\",\n",
    "            non_negative=True,\n",
    "            units=\"individuals\"\n",
    "        ))\n",
    "        \n",
    "        self.define_variable(VariableSpec(\n",
    "            name=\"I\",\n",
    "            symbol=\"I\",\n",
    "            description=\"Infectious individuals\",\n",
    "            non_negative=True,\n",
    "            units=\"individuals\"\n",
    "        ))\n",
    "        \n",
    "        self.define_variable(VariableSpec(\n",
    "            name=\"R\",\n",
    "            symbol=\"R\",\n",
    "            description=\"Recovered/Removed individuals\",\n",
    "            non_negative=True,\n",
    "            units=\"individuals\"\n",
    "        ))\n",
    "        \n",
    "        # Add model-level constraints\n",
    "        self.add_constraint(ModelConstraint(\n",
    "            expression=\"beta / gamma > 1\",\n",
    "            description=\"R0 > 1 required for epidemic spread\",\n",
    "            severity=\"warning\",\n",
    "            name=\"R0_epidemic\"\n",
    "        ))\n",
    "\n",
    "# Create model instance\n",
    "model = SIRWithValidation()\n",
    "print(f\"Model type: {model.model_type}\")\n",
    "print(f\"Parameters: {list(model.parameter_specs.keys())}\")\n",
    "print(f\"Variables: {list(model.variable_specs.keys())}\")\n",
    "print(f\"Constraints: {len(model.model_constraints)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Parameter Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:18:45.379412614Z",
     "start_time": "2026-03-11T18:18:45.261457580Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.918618Z",
     "iopub.status.busy": "2026-03-11T19:33:19.918394Z",
     "iopub.status.idle": "2026-03-11T19:33:19.923358Z",
     "shell.execute_reply": "2026-03-11T19:33:19.922340Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test 1: Valid parameters\n",
      "  ✓ Validation passed\n"
     ]
    }
   ],
   "source": [
    "# Test 1: Valid parameters\n",
    "print(\"Test 1: Valid parameters\")\n",
    "try:\n",
    "    model.validate_parameters({'beta': 0.3, 'gamma': 0.1})\n",
    "    print(\"  ✓ Validation passed\")\n",
    "except ValidationError as e:\n",
    "    print(f\"  ✗ Validation failed: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:18:46.901130820Z",
     "start_time": "2026-03-11T18:18:46.597941127Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.925858Z",
     "iopub.status.busy": "2026-03-11T19:33:19.925580Z",
     "iopub.status.idle": "2026-03-11T19:33:19.931603Z",
     "shell.execute_reply": "2026-03-11T19:33:19.930295Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test 2: Missing required parameter\n",
      "  ✓ Caught error: ValidationError\n",
      "    Message: Missing required parameter: gamma...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/fccoelho/Documentos/Projects_software/epimodels/epimodels/__init__.py:123: UserWarning: Constraint violated: R0 > 1 required for epidemic spread (Failed to evaluate expression: Unknown variable: gamma)\n",
      "  warnings.warn(warning_msg, UserWarning)\n"
     ]
    }
   ],
   "source": [
    "# Test 2: Missing required parameter\n",
    "print(\"Test 2: Missing required parameter\")\n",
    "try:\n",
    "    model.validate_parameters({'beta': 0.3})\n",
    "    print(\"  ✗ Should have raised error\")\n",
    "except ValidationError as e:\n",
    "    print(f\"  ✓ Caught error: {type(e).__name__}\")\n",
    "    print(f\"    Message: {str(e)[:60]}...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:18:56.304724358Z",
     "start_time": "2026-03-11T18:18:56.110241563Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.934360Z",
     "iopub.status.busy": "2026-03-11T19:33:19.933994Z",
     "iopub.status.idle": "2026-03-11T19:33:19.940359Z",
     "shell.execute_reply": "2026-03-11T19:33:19.939249Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test 3: Parameter out of bounds (negative beta)\n",
      "  ✓ Caught error: ValidationError\n",
      "    Message: Parameter 'beta' value -0.3 is below minimum bound 0\n",
      "Paramet...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/fccoelho/Documentos/Projects_software/epimodels/epimodels/__init__.py:123: UserWarning: Constraint violated: R0 > 1 required for epidemic spread\n",
      "  warnings.warn(warning_msg, UserWarning)\n"
     ]
    }
   ],
   "source": [
    "# Test 3: Parameter out of bounds\n",
    "print(\"Test 3: Parameter out of bounds (negative beta)\")\n",
    "try:\n",
    "    model.validate_parameters({'beta': -0.3, 'gamma': 0.1})\n",
    "    print(\"  ✗ Should have raised error\")\n",
    "except ValidationError as e:\n",
    "    print(f\"  ✓ Caught error: {type(e).__name__}\")\n",
    "    print(f\"    Message: {str(e)[:60]}...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:18:58.906767725Z",
     "start_time": "2026-03-11T18:18:58.814684517Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.943412Z",
     "iopub.status.busy": "2026-03-11T19:33:19.943050Z",
     "iopub.status.idle": "2026-03-11T19:33:19.950319Z",
     "shell.execute_reply": "2026-03-11T19:33:19.949395Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test 4: Warning constraint (R0 < 1)\n",
      "  ✓ Warning raised: Constraint violated: R0 > 1 required for epidemic spread...\n"
     ]
    }
   ],
   "source": [
    "# Test 4: Warning constraint (R0 < 1)\n",
    "print(\"Test 4: Warning constraint (R0 < 1)\")\n",
    "with warnings.catch_warnings(record=True) as w:\n",
    "    warnings.simplefilter(\"always\")\n",
    "    try:\n",
    "        model.validate_parameters({'beta': 0.1, 'gamma': 0.3})\n",
    "        if w:\n",
    "            print(f\"  ✓ Warning raised: {str(w[0].message)[:60]}...\")\n",
    "        else:\n",
    "            print(\"  ✗ No warning raised\")\n",
    "    except ValidationError as e:\n",
    "        print(f\"  ✗ Validation failed: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:19:00.983688407Z",
     "start_time": "2026-03-11T18:19:00.776045815Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.952632Z",
     "iopub.status.busy": "2026-03-11T19:33:19.952410Z",
     "iopub.status.idle": "2026-03-11T19:33:19.958181Z",
     "shell.execute_reply": "2026-03-11T19:33:19.956847Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test 5: Constraint violation (gamma = 0)\n",
      "  ✓ Caught error: ValidationError\n",
      "    Message: Parameter 'gamma' value 0.0 violates constraint: value > 0...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/fccoelho/Documentos/Projects_software/epimodels/epimodels/__init__.py:123: UserWarning: Constraint violated: R0 > 1 required for epidemic spread (Failed to evaluate expression: float division by zero)\n",
      "  warnings.warn(warning_msg, UserWarning)\n"
     ]
    }
   ],
   "source": [
    "# Test 5: Constraint violation (gamma = 0)\n",
    "print(\"Test 5: Constraint violation (gamma = 0)\")\n",
    "try:\n",
    "    model.validate_parameters({'beta': 0.3, 'gamma': 0})\n",
    "    print(\"  ✗ Should have raised error\")\n",
    "except ValidationError as e:\n",
    "    print(f\"  ✓ Caught error: {type(e).__name__}\")\n",
    "    print(f\"    Message: {str(e)[:60]}...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Initial Condition Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:19:08.808095716Z",
     "start_time": "2026-03-11T18:19:08.701985981Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.961206Z",
     "iopub.status.busy": "2026-03-11T19:33:19.960920Z",
     "iopub.status.idle": "2026-03-11T19:33:19.966744Z",
     "shell.execute_reply": "2026-03-11T19:33:19.965598Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test 6: Valid initial conditions\n",
      "  ✓ Initial conditions valid\n"
     ]
    }
   ],
   "source": [
    "# Test 6: Valid initial conditions\n",
    "print(\"Test 6: Valid initial conditions\")\n",
    "try:\n",
    "    model.validate_initial_conditions([990, 10, 0], totpop=1000)\n",
    "    print(\"  ✓ Initial conditions valid\")\n",
    "except ValidationError as e:\n",
    "    print(f\"  ✗ Validation failed: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:19:10.225179248Z",
     "start_time": "2026-03-11T18:19:10.125218517Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.969619Z",
     "iopub.status.busy": "2026-03-11T19:33:19.969367Z",
     "iopub.status.idle": "2026-03-11T19:33:19.975548Z",
     "shell.execute_reply": "2026-03-11T19:33:19.973892Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test 7: Negative initial condition\n",
      "  ✓ Caught error: ValidationError\n"
     ]
    }
   ],
   "source": [
    "# Test 7: Negative initial condition\n",
    "print(\"Test 7: Negative initial condition\")\n",
    "try:\n",
    "    model.validate_initial_conditions([990, -10, 0], totpop=1000)\n",
    "    print(\"  ✗ Should have raised error\")\n",
    "except ValidationError as e:\n",
    "    print(f\"  ✓ Caught error: {type(e).__name__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:19:13.936290998Z",
     "start_time": "2026-03-11T18:19:13.807591220Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.978631Z",
     "iopub.status.busy": "2026-03-11T19:33:19.978227Z",
     "iopub.status.idle": "2026-03-11T19:33:19.985714Z",
     "shell.execute_reply": "2026-03-11T19:33:19.983917Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test 8: Initial conditions exceed population\n",
      "  ✓ Caught error: ValidationError\n"
     ]
    }
   ],
   "source": [
    "# Test 8: Initial conditions exceed population\n",
    "print(\"Test 8: Initial conditions exceed population\")\n",
    "try:\n",
    "    model.validate_initial_conditions([1000, 100, 0], totpop=1000)\n",
    "    print(\"  ✗ Should have raised error\")\n",
    "except ValidationError as e:\n",
    "    print(f\"  ✓ Caught error: {type(e).__name__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Constraint Language\n",
    "\n",
    "The constraint evaluator supports various operators and expressions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:19:16.716050965Z",
     "start_time": "2026-03-11T18:19:16.582409647Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:19.989437Z",
     "iopub.status.busy": "2026-03-11T19:33:19.988757Z",
     "iopub.status.idle": "2026-03-11T19:33:19.998305Z",
     "shell.execute_reply": "2026-03-11T19:33:19.996934Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing constraint expressions:\n",
      "------------------------------------------------------------\n",
      "✓ 'beta > gamma' with {'beta': 0.5, 'gamma': 0.1}: True\n",
      "✓ 'beta / gamma > 1' with {'beta': 0.3, 'gamma': 0.1}: True\n",
      "✓ 'p + q <= 1' with {'p': 0.6, 'q': 0.3}: True\n",
      "✓ 'rate**2 > threshold' with {'rate': 2, 'threshold': 3}: True\n",
      "✓ 'x > 0 and y > 0' with {'x': 1, 'y': 1}: True\n",
      "✓ '(alpha + beta) / gamma > 2' with {'alpha': 0.2, 'beta': 0.3, 'gamma': 0.2}: True\n"
     ]
    }
   ],
   "source": [
    "# Test various constraint expressions\n",
    "constraints = [\n",
    "    (\"beta > gamma\", {'beta': 0.5, 'gamma': 0.1}),\n",
    "    (\"beta / gamma > 1\", {'beta': 0.3, 'gamma': 0.1}),\n",
    "    (\"p + q <= 1\", {'p': 0.6, 'q': 0.3}),\n",
    "    (\"rate**2 > threshold\", {'rate': 2, 'threshold': 3}),\n",
    "    (\"x > 0 and y > 0\", {'x': 1, 'y': 1}),\n",
    "    (\"(alpha + beta) / gamma > 2\", {'alpha': 0.2, 'beta': 0.3, 'gamma': 0.2}),\n",
    "]\n",
    "\n",
    "print(\"Testing constraint expressions:\")\n",
    "print(\"-\" * 60)\n",
    "for expr, params in constraints:\n",
    "    satisfied, msg = evaluate_constraint(expr, params)\n",
    "    status = \"✓\" if satisfied else \"✗\"\n",
    "    print(f\"{status} '{expr}' with {params}: {satisfied}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. More Complex Model Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:19:19.218039971Z",
     "start_time": "2026-03-11T18:19:19.021838287Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:20.001342Z",
     "iopub.status.busy": "2026-03-11T19:33:20.000959Z",
     "iopub.status.idle": "2026-03-11T19:33:20.017514Z",
     "shell.execute_reply": "2026-03-11T19:33:20.015685Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created SEQIAHR model\n",
      "Parameters: ['beta', 'p', 'chi', 'gamma', 'mu']\n",
      "Constraints: 2\n"
     ]
    }
   ],
   "source": [
    "class SEQIAHRWithValidation(BaseModel):\n",
    "    \"\"\"COVID-19 model with quarantine and hospitalization.\"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.model_type = \"SEQIAHR\"\n",
    "        \n",
    "        # Transmission parameters\n",
    "        self.define_parameter(ParameterSpec(\n",
    "            name=\"beta\",\n",
    "            symbol=r\"$\\beta$\",\n",
    "            description=\"Transmission rate\",\n",
    "            bounds=(0, None),\n",
    "            constraints=[\"value > 0\"],\n",
    "            units=\"1/day\"\n",
    "        ))\n",
    "        \n",
    "        # Probabilities (must be between 0 and 1)\n",
    "        self.define_parameter(ParameterSpec(\n",
    "            name=\"p\",\n",
    "            symbol=\"p\",\n",
    "            description=\"Proportion asymptomatic\",\n",
    "            bounds=(0, 1),\n",
    "            constraints=[\"value >= 0\", \"value <= 1\"],\n",
    "            units=\"dimensionless\"\n",
    "        ))\n",
    "        \n",
    "        self.define_parameter(ParameterSpec(\n",
    "            name=\"chi\",\n",
    "            symbol=r\"$\\chi$\",\n",
    "            description=\"Quarantine effectiveness\",\n",
    "            bounds=(0, 1),\n",
    "            constraints=[\"value >= 0\", \"value <= 1\"],\n",
    "            units=\"dimensionless\"\n",
    "        ))\n",
    "        \n",
    "        # Rates\n",
    "        self.define_parameter(ParameterSpec(\n",
    "            name=\"gamma\",\n",
    "            symbol=r\"$\\gamma$\",\n",
    "            description=\"Recovery rate\",\n",
    "            bounds=(0, None),\n",
    "            constraints=[\"value > 0\"],\n",
    "            units=\"1/day\"\n",
    "        ))\n",
    "        \n",
    "        self.define_parameter(ParameterSpec(\n",
    "            name=\"mu\",\n",
    "            symbol=r\"$\\mu$\",\n",
    "            description=\"Mortality rate\",\n",
    "            bounds=(0, None),\n",
    "            constraints=[\"value >= 0\"],\n",
    "            units=\"1/day\"\n",
    "        ))\n",
    "        \n",
    "        # Add constraint: mortality + recovery should be reasonable\n",
    "        self.add_constraint(ModelConstraint(\n",
    "            expression=\"mu + gamma < 1\",\n",
    "            description=\"Combined exit rate should be < 1/day for stability\",\n",
    "            severity=\"warning\"\n",
    "        ))\n",
    "        \n",
    "        # Add constraint: quarantine effectiveness\n",
    "        self.add_constraint(ModelConstraint(\n",
    "            expression=\"chi >= 0.5\",\n",
    "            description=\"Quarantine should be at least 50% effective\",\n",
    "            severity=\"warning\"\n",
    "        ))\n",
    "\n",
    "# Create and test model\n",
    "complex_model = SEQIAHRWithValidation()\n",
    "print(f\"Created {complex_model.model_type} model\")\n",
    "print(f\"Parameters: {list(complex_model.parameter_specs.keys())}\")\n",
    "print(f\"Constraints: {len(complex_model.model_constraints)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:19:20.599241314Z",
     "start_time": "2026-03-11T18:19:20.418980640Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:20.024056Z",
     "iopub.status.busy": "2026-03-11T19:33:20.023607Z",
     "iopub.status.idle": "2026-03-11T19:33:20.034706Z",
     "shell.execute_reply": "2026-03-11T19:33:20.033264Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validating complex model with valid parameters:\n",
      "  ✓ Validation passed with 0 warnings\n"
     ]
    }
   ],
   "source": [
    "# Test complex model validation\n",
    "params_valid = {\n",
    "    'beta': 0.3,\n",
    "    'p': 0.4,\n",
    "    'chi': 0.7,\n",
    "    'gamma': 0.1,\n",
    "    'mu': 0.01\n",
    "}\n",
    "\n",
    "print(\"Validating complex model with valid parameters:\")\n",
    "with warnings.catch_warnings(record=True) as w:\n",
    "    warnings.simplefilter(\"always\")\n",
    "    complex_model.validate_parameters(params_valid)\n",
    "    print(f\"  ✓ Validation passed with {len(w)} warnings\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:19:22.199221952Z",
     "start_time": "2026-03-11T18:19:22.101534666Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:20.038570Z",
     "iopub.status.busy": "2026-03-11T19:33:20.038273Z",
     "iopub.status.idle": "2026-03-11T19:33:20.044545Z",
     "shell.execute_reply": "2026-03-11T19:33:20.043674Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing with probability out of bounds (p=1.5):\n",
      "  ✓ Caught error: Parameter 'p' value 1.5 exceeds maximum bound 1\n",
      "Parameter 'p' value 1.5 violates...\n"
     ]
    }
   ],
   "source": [
    "# Test with probability out of bounds\n",
    "params_invalid = {\n",
    "    'beta': 0.3,\n",
    "    'p': 1.5,  # Invalid: > 1\n",
    "    'chi': 0.7,\n",
    "    'gamma': 0.1,\n",
    "    'mu': 0.01\n",
    "}\n",
    "\n",
    "print(\"Testing with probability out of bounds (p=1.5):\")\n",
    "try:\n",
    "    complex_model.validate_parameters(params_invalid)\n",
    "    print(\"  ✗ Should have raised error\")\n",
    "except ValidationError as e:\n",
    "    print(f\"  ✓ Caught error: {str(e)[:80]}...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Backward Compatibility\n",
    "\n",
    "Existing models without rich specs still work with simple validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:19:59.178600094Z",
     "start_time": "2026-03-11T18:19:59.021128740Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:20.049744Z",
     "iopub.status.busy": "2026-03-11T19:33:20.049384Z",
     "iopub.status.idle": "2026-03-11T19:33:20.063577Z",
     "shell.execute_reply": "2026-03-11T19:33:20.060908Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created legacy model with 2 parameters\n",
      "✓ Simple validation works\n",
      "✓ Simple validation catches missing parameters\n",
      "✓ Simple validation catches negative parameters\n"
     ]
    }
   ],
   "source": [
    "# Legacy-style model (backward compatible)\n",
    "class LegacyModel(BaseModel):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.model_type = \"Legacy\"\n",
    "        # Old-style simple dicts\n",
    "        self.parameters = {\"alpha\": r\"$\\alpha$\", \"delta\": r\"$\\delta$\"}\n",
    "        self.state_variables = {\"X\": \"Variable X\", \"Y\": \"Variable Y\"}\n",
    "\n",
    "legacy = LegacyModel()\n",
    "print(f\"Created legacy model with {len(legacy.parameters)} parameters\")\n",
    "\n",
    "# Simple validation still works\n",
    "legacy.validate_parameters({'alpha': 0.5, 'delta': 0.2})\n",
    "print(\"✓ Simple validation works\")\n",
    "\n",
    "try:\n",
    "    legacy.validate_parameters({'alpha': 0.5})\n",
    "    print(\"✗ Should have caught missing parameter\")\n",
    "except ValidationError:\n",
    "    print(\"✓ Simple validation catches missing parameters\")\n",
    "\n",
    "try:\n",
    "    legacy.validate_parameters({'alpha': -0.5, 'delta': 0.2})\n",
    "    print(\"✗ Should have caught negative parameter\")\n",
    "except ValidationError:\n",
    "    print(\"✓ Simple validation catches negative parameters\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Generating Documentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-11T18:21:14.429748265Z",
     "start_time": "2026-03-11T18:21:14.302743632Z"
    },
    "execution": {
     "iopub.execute_input": "2026-03-11T19:33:20.069404Z",
     "iopub.status.busy": "2026-03-11T19:33:20.069126Z",
     "iopub.status.idle": "2026-03-11T19:33:20.081704Z",
     "shell.execute_reply": "2026-03-11T19:33:20.080484Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# SIR Model Parameters\n",
      "\n",
      "## Parameters\n",
      "\n",
      "### $\\beta$ (beta)\n",
      "\n",
      "Transmission rate\n",
      "\n",
      "- **Bounds**: −∞ to ∞\n",
      "- **Units**: 1/day\n",
      "- **Typical range**: 0.1 to 1.0\n",
      "- **Constraints**: value > 0\n",
      "\n",
      "### $\\gamma$ (gamma)\n",
      "\n",
      "Recovery rate (1 / average infectious period)\n",
      "\n",
      "- **Bounds**: −∞ to ∞\n",
      "- **Units**: 1/day\n",
      "- **Typical range**: 0.05 to 0.5\n",
      "- **Constraints**: value > 0\n",
      "\n",
      "## Model Constraints\n",
      "\n",
      "- R0 > 1 required for epidemic spread (warning)\n",
      "  - Expression: `beta / gamma > 1`\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def generate_parameter_docs(model):\n",
    "    \"\"\"Generate markdown documentation from parameter specs.\"\"\"\n",
    "    if not model.parameter_specs:\n",
    "        return \"No parameter specifications defined.\"\n",
    "    \n",
    "    doc = f\"# {model.model_type} Model Parameters\\n\\n\"\n",
    "    \n",
    "    doc += \"## Parameters\\n\\n\"\n",
    "    for name, spec in model.parameter_specs.items():\n",
    "        doc += f\"### {spec.symbol} ({name})\\n\\n\"\n",
    "        doc += f\"{spec.description}\\n\\n\"\n",
    "        if spec.bounds:\n",
    "            min_val, max_val = spec.bounds\n",
    "            doc += f\"- **Bounds**: {min_val or '−∞'} to {max_val or '∞'}\\n\"\n",
    "        if spec.units:\n",
    "            doc += f\"- **Units**: {spec.units}\\n\"\n",
    "        if spec.typical_range:\n",
    "            doc += f\"- **Typical range**: {spec.typical_range[0]} to {spec.typical_range[1]}\\n\"\n",
    "        if spec.constraints:\n",
    "            doc += f\"- **Constraints**: {', '.join(spec.constraints)}\\n\"\n",
    "        doc += \"\\n\"\n",
    "    \n",
    "    if model.model_constraints:\n",
    "        doc += \"## Model Constraints\\n\\n\"\n",
    "        for constraint in model.model_constraints:\n",
    "            severity_str = f\" ({constraint.severity})\" if constraint.severity != \"error\" else \"\"\n",
    "            doc += f\"- {constraint.description}{severity_str}\\n\"\n",
    "            doc += f\"  - Expression: `{constraint.expression}`\\n\\n\"\n",
    "    \n",
    "    return doc\n",
    "\n",
    "# Generate documentation for SIR model\n",
    "print(generate_parameter_docs(model))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated:\n",
    "\n",
    "1. **ParameterSpec**: Define parameters with bounds, constraints, units, and documentation\n",
    "2. **VariableSpec**: Define state variables with validation rules\n",
    "3. **ModelConstraint**: Cross-parameter constraints with severity levels\n",
    "4. **Constraint Language**: Flexible expression evaluation\n",
    "5. **Validation Errors**: Clear error messages for constraint violations\n",
    "6. **Warning Constraints**: Soft constraints that warn but don't fail\n",
    "7. **Backward Compatibility**: Existing models work unchanged\n",
    "8. **Documentation Generation**: Automatic documentation from specs"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}