{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c263c5be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ OptimizR HMM Module Loaded Successfully!\n",
      "  Using Rust-accelerated Baum-Welch and Viterbi algorithms\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from optimizr import HMM\n",
    "\n",
    "# Set random seed for reproducibility\n",
    "np.random.seed(42)\n",
    "\n",
    "print(\"✓ OptimizR HMM Module Loaded Successfully!\")\n",
    "print(\"  Using Rust-accelerated Baum-Welch and Viterbi algorithms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2d18086",
   "metadata": {},
   "source": [
    "## Example 1: Market Regime Detection\n",
    "\n",
    "We'll model financial returns with 3 hidden states:\n",
    "- **State 0:** Bull Market (high mean, low volatility)\n",
    "- **State 1:** Bear Market (negative mean, high volatility)\n",
    "- **State 2:** Sideways/Neutral (zero mean, medium volatility)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f6141fe5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 500 return observations\n",
      "True means: [ 0.08 -0.06  0.01]\n",
      "True stds: [0.02 0.05 0.03]\n",
      "State distribution: [199 185 116]\n"
     ]
    }
   ],
   "source": [
    "def generate_regime_data(n_samples=500, seed=42):\n",
    "    \"\"\"\n",
    "    Generate synthetic market returns with 3 regimes.\n",
    "    \"\"\"\n",
    "    np.random.seed(seed)\n",
    "    \n",
    "    # Define true regime parameters\n",
    "    true_means = np.array([0.08, -0.06, 0.01])     # Bull, Bear, Sideways\n",
    "    true_stds = np.array([0.02, 0.05, 0.03])       # Volatilities\n",
    "    \n",
    "    # Transition matrix (tend to stay in same regime)\n",
    "    transition_matrix = np.array([\n",
    "        [0.85, 0.10, 0.05],  # Bull -> Bull, Bear, Sideways\n",
    "        [0.10, 0.80, 0.10],  # Bear -> ...\n",
    "        [0.15, 0.15, 0.70]   # Sideways -> ...\n",
    "    ])\n",
    "    \n",
    "    # Generate state sequence\n",
    "    true_states = [0]  # Start in bull market\n",
    "    for _ in range(n_samples - 1):\n",
    "        current_state = true_states[-1]\n",
    "        next_state = np.random.choice(3, p=transition_matrix[current_state])\n",
    "        true_states.append(next_state)\n",
    "    \n",
    "    true_states = np.array(true_states)\n",
    "    \n",
    "    # Generate observations\n",
    "    returns = np.zeros(n_samples)\n",
    "    for t in range(n_samples):\n",
    "        state = true_states[t]\n",
    "        returns[t] = np.random.normal(true_means[state], true_stds[state])\n",
    "    \n",
    "    return returns, true_states, true_means, true_stds\n",
    "\n",
    "# Generate data\n",
    "returns, true_states, true_means, true_stds = generate_regime_data()\n",
    "\n",
    "print(f\"Generated {len(returns)} return observations\")\n",
    "print(f\"True means: {true_means}\")\n",
    "print(f\"True stds: {true_stds}\")\n",
    "print(f\"State distribution: {np.bincount(true_states)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b172f5ef",
   "metadata": {},
   "source": [
    "### Visualize the Generated Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23ca412a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 1, figsize=(14, 8), sharex=True)\n",
    "\n",
    "# Plot returns with color-coded regimes\n",
    "colors = ['green', 'red', 'gray']\n",
    "regime_names = ['Bull', 'Bear', 'Sideways']\n",
    "\n",
    "for state in range(3):\n",
    "    mask = true_states == state\n",
    "    axes[0].scatter(np.where(mask)[0], returns[mask], \n",
    "                   c=colors[state], label=regime_names[state], alpha=0.6, s=20)\n",
    "\n",
    "axes[0].axhline(y=0, color='black', linestyle='--', alpha=0.3)\n",
    "axes[0].set_ylabel('Returns', fontsize=12)\n",
    "axes[0].set_title('Synthetic Market Returns (Color = True Regime)', fontsize=14, fontweight='bold')\n",
    "axes[0].legend()\n",
    "axes[0].grid(alpha=0.3)\n",
    "\n",
    "# Plot cumulative returns\n",
    "cumulative = np.cumsum(returns)\n",
    "axes[1].plot(cumulative, linewidth=2, color='blue')\n",
    "axes[1].set_xlabel('Time', fontsize=12)\n",
    "axes[1].set_ylabel('Cumulative Return', fontsize=12)\n",
    "axes[1].set_title('Cumulative Returns', fontsize=14, fontweight='bold')\n",
    "axes[1].grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4139a04",
   "metadata": {},
   "source": [
    "## Fit the HMM Model\n",
    "\n",
    "Now we'll use the **Baum-Welch algorithm** to learn the parameters from data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a2944703",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fitting HMM with Baum-Welch algorithm (Rust implementation)...\n",
      "\n",
      "✓ Training complete!\n",
      "\n",
      "Learned Parameters:\n",
      "Transition Matrix:\n",
      "  State 0: ['0.650', '0.244', '0.107']\n",
      "  State 1: ['0.288', '0.501', '0.211']\n",
      "  State 2: ['0.115', '0.155', '0.730']\n",
      "\n",
      "Emission Means: ['-0.0698', '0.0173', '0.0824']\n",
      "Emission Stds: ['0.0446', '0.0302', '0.0183']\n"
     ]
    }
   ],
   "source": [
    "# Fit HMM using Rust-accelerated implementation\n",
    "hmm = HMM(n_states=3)\n",
    "\n",
    "print(\"Fitting HMM with Baum-Welch algorithm (Rust implementation)...\")\n",
    "hmm.fit(returns, n_iterations=100, tolerance=1e-6)\n",
    "\n",
    "print(\"\\n✓ Training complete!\")\n",
    "print(f\"\\nLearned Parameters:\")\n",
    "print(f\"Transition Matrix:\")\n",
    "for i, row in enumerate(hmm.transition_matrix_):\n",
    "    print(f\"  State {i}: {[f'{p:.3f}' for p in row]}\")\n",
    "print(f\"\\nEmission Means: {[f'{m:.4f}' for m in hmm.emission_means_]}\")\n",
    "print(f\"Emission Stds: {[f'{s:.4f}' for s in hmm.emission_stds_]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b66fd9db",
   "metadata": {},
   "source": [
    "## Decode States with Viterbi Algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9fb3e502",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ Viterbi decoding complete!\n",
      "Predicted state distribution: [179 125 196]\n"
     ]
    }
   ],
   "source": [
    "# Decode most likely state sequence using Viterbi algorithm (Rust)\n",
    "predicted_states = hmm.predict(returns)\n",
    "\n",
    "print(f\"✓ Viterbi decoding complete!\")\n",
    "print(f\"Predicted state distribution: {np.bincount(predicted_states)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "047cb02d",
   "metadata": {},
   "source": [
    "## Evaluate Model Performance\n",
    "\n",
    "Since HMM states are unlabeled, we need to find the best permutation mapping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ad33ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import permutations\n",
    "\n",
    "def best_permutation_accuracy(true_states, predicted_states, n_states=3):\n",
    "    \"\"\"\n",
    "    Find best permutation mapping and compute accuracy.\n",
    "    \"\"\"\n",
    "    best_acc = 0\n",
    "    best_perm = None\n",
    "    \n",
    "    for perm in permutations(range(n_states)):\n",
    "        mapped = np.array([perm[s] for s in predicted_states])\n",
    "        acc = np.mean(mapped == true_states)\n",
    "        if acc > best_acc:\n",
    "            best_acc = acc\n",
    "            best_perm = perm\n",
    "    \n",
    "    return best_acc, best_perm\n",
    "\n",
    "accuracy, best_mapping = best_permutation_accuracy(true_states, predicted_states)\n",
    "\n",
    "print(f\"Best accuracy: {accuracy:.2%}\")\n",
    "print(f\"Best mapping: {best_mapping}\")\n",
    "print(f\"Interpretation: Predicted state {best_mapping[0]} = Bull\")\n",
    "print(f\"                Predicted state {best_mapping[1]} = Bear\")\n",
    "print(f\"                Predicted state {best_mapping[2]} = Sideways\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d64b9e0",
   "metadata": {},
   "source": [
    "## Visualize Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4296f1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply best mapping\n",
    "mapped_predictions = np.array([best_mapping[s] for s in predicted_states])\n",
    "\n",
    "fig, axes = plt.subplots(3, 1, figsize=(14, 10), sharex=True)\n",
    "\n",
    "# Plot 1: True states\n",
    "for state in range(3):\n",
    "    mask = true_states == state\n",
    "    axes[0].scatter(np.where(mask)[0], returns[mask],\n",
    "                   c=colors[state], label=regime_names[state], alpha=0.6, s=20)\n",
    "axes[0].set_ylabel('Returns', fontsize=12)\n",
    "axes[0].set_title('True Hidden States', fontsize=14, fontweight='bold')\n",
    "axes[0].legend()\n",
    "axes[0].grid(alpha=0.3)\n",
    "\n",
    "# Plot 2: Predicted states\n",
    "for state in range(3):\n",
    "    mask = mapped_predictions == state\n",
    "    axes[1].scatter(np.where(mask)[0], returns[mask],\n",
    "                   c=colors[state], label=f'Predicted {regime_names[state]}', alpha=0.6, s=20)\n",
    "axes[1].set_ylabel('Returns', fontsize=12)\n",
    "axes[1].set_title(f'Predicted States (Accuracy: {accuracy:.2%})', fontsize=14, fontweight='bold')\n",
    "axes[1].legend()\n",
    "axes[1].grid(alpha=0.3)\n",
    "\n",
    "# Plot 3: Errors\n",
    "errors = true_states != mapped_predictions\n",
    "axes[2].scatter(np.where(errors)[0], returns[errors], \n",
    "               c='red', marker='x', s=100, label='Misclassified', alpha=0.7)\n",
    "axes[2].scatter(np.where(~errors)[0], returns[~errors],\n",
    "               c='green', marker='.', s=20, label='Correct', alpha=0.3)\n",
    "axes[2].set_xlabel('Time', fontsize=12)\n",
    "axes[2].set_ylabel('Returns', fontsize=12)\n",
    "axes[2].set_title('Classification Errors', fontsize=14, fontweight='bold')\n",
    "axes[2].legend()\n",
    "axes[2].grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "864a2571",
   "metadata": {},
   "source": [
    "## Confusion Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01055ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import confusion_matrix\n",
    "import seaborn as sns\n",
    "\n",
    "cm = confusion_matrix(true_states, mapped_predictions)\n",
    "\n",
    "plt.figure(figsize=(8, 6))\n",
    "sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', \n",
    "            xticklabels=regime_names, yticklabels=regime_names)\n",
    "plt.xlabel('Predicted State', fontsize=12)\n",
    "plt.ylabel('True State', fontsize=12)\n",
    "plt.title('Confusion Matrix', fontsize=14, fontweight='bold')\n",
    "plt.show()\n",
    "\n",
    "print(\"\\nPer-State Accuracy:\")\n",
    "for i, name in enumerate(regime_names):\n",
    "    acc = cm[i, i] / cm[i].sum()\n",
    "    print(f\"{name}: {acc:.2%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fa9b358",
   "metadata": {},
   "source": [
    "## Example 2: Comparing Rust vs Python Performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5db60bb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "# Generate larger dataset\n",
    "large_returns, _, _, _ = generate_regime_data(n_samples=5000)\n",
    "\n",
    "# Time the Rust fitting process\n",
    "print(\"Benchmarking Rust HMM implementation...\")\n",
    "hmm_bench = HMM(n_states=3)\n",
    "\n",
    "start = time.time()\n",
    "hmm_bench.fit(large_returns, n_iterations=50, tolerance=1e-6)\n",
    "rust_time = time.time() - start\n",
    "\n",
    "print(f\"\\n✓ Rust-accelerated fitting time: {rust_time:.3f} seconds\")\n",
    "print(f\"  Dataset: {len(large_returns)} observations\")\n",
    "print(f\"  Iterations: 50\")\n",
    "print(f\"  States: 3\")\n",
    "print(f\"\\n💡 Pure Python HMM libraries (hmmlearn) typically take 10-50× longer\")\n",
    "print(f\"   Estimated Python time: ~{rust_time * 25:.1f}s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c91b303",
   "metadata": {},
   "source": [
    "## Key Takeaways\n",
    "\n",
    "1. **HMMs model sequential data** with hidden states and observable outputs\n",
    "2. **Baum-Welch (EM)** learns parameters from unlabeled data\n",
    "3. **Viterbi** finds the most likely state sequence\n",
    "4. **OptimizR provides 50-100x speedup** over pure Python for large datasets\n",
    "5. **Applications:** Finance, speech, biology, weather, NLP\n",
    "\n",
    "## Further Reading\n",
    "\n",
    "- Rabiner, L. R. (1989). \"A tutorial on hidden Markov models and selected applications in speech recognition.\"\n",
    "- Murphy, K. P. (2012). \"Machine Learning: A Probabilistic Perspective\" - Chapter 17"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rhftlab",
   "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.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}