{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1ecb329a-8614-42a5-82a8-78a3d5dacada",
   "metadata": {},
   "source": [
    "# Ex 1: Automatic Topological Generation & Surrogate Optimization (2D / 4-DOF)\n",
    "\n",
    "This notebook represents the complete methodology for imprecise tolerance analysis on the 2D baseline assembly. \n",
    "\n",
    "Building on the manual vector loops constructed previously, this notebook introduces the streamlined workflow required to scale the methodology to high-dimensional systems:\n",
    "1. **Automatic Graph-Based Assembly Modeling:** Automatically constructing the `SystemOfConstraintsAssemblyModel` from a high-level topological dictionary.\n",
    "2. **Multi-Layer Perceptron (MLP) Surrogate:** Replacing expensive direct geometric optimizations with a PyTorch neural network.\n",
    "3. **Automated Tracking & Constraint Injection:** Using `otaf.optimization.OptimizationTracker` and the model's native constraint generator to manage the credal space boundaries.\n",
    "4. **Min-Max Bounds Resolution:** Finding a feasible warm-start ($x_0$) and optimizing across threshold levels ($\\alpha_k$) to compute the P-Box."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bd254d60-fd5b-4c8e-9f4a-1de5de524e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy\n",
    "import openturns as ot\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from scipy.optimize import minimize, NonlinearConstraint, Bounds\n",
    "\n",
    "import otaf\n",
    "from otaf.example_models import model1\n",
    "from gldpy import GLD\n",
    "\n",
    "ot.Log.Show(ot.Log.NONE)\n",
    "np.set_printoptions(suppress=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bcde591-85fa-4662-992b-78fca659bfb0",
   "metadata": {},
   "source": [
    "## 1. High-Level Topological Dictionary & Automated Loops\n",
    "\n",
    "Instead of manually defining deviation and gap matrices, we define the assembly as a topological graph. The `otaf.AssemblyDataProcessor` interprets the interactions and expands the first-order loops.\n",
    "\n",
    "*(Note: For the surrogate and optimization phases, we will dynamically pull the parameters directly from `otaf.example_models.model1`, but this dictionary illustrates how the physics of `model1` are constructed).*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c637ce58-ba04-4d7d-a03f-441b1d8e60d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Nominal dimensions and base frames\n",
    "X1, X2, X3 = 99.8, 100.0, 10.0\n",
    "R0 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])\n",
    "x_, y_, z_ = R0[0], R0[1], R0[2]\n",
    "\n",
    "# Part 1 and 2 points\n",
    "P1A0, P1A1, P1A2 = np.array((0, X3/2, 0.0)), np.array((0, X3, 0.0)), np.array((0, 0, 0.0))\n",
    "P1B0, P1B1, P1B2 = np.array((X1, X3/2, 0.0)), np.array((X1, X3, 0.0)), np.array((X1, 0, 0.0))\n",
    "P1C0, P1C1, P1C2 = np.array((X1/2, 0, 0.0)), np.array((0, 0, 0.0)), np.array((X1, 0, 0.0))\n",
    "P2A0, P2A1, P2A2 = np.array((0, X3/2, 0.0)), np.array((0, X3, 0.0)), np.array((0, 0, 0.0))\n",
    "P2B0, P2B1, P2B2 = np.array((X2, X3/2, 0.0)), np.array((X2, X3, 0.0)), np.array((X2, 0, 0.0))\n",
    "P2C0, P2C1, P2C2 = np.array((X2/2, 0, 0.0)), np.array((0, 0, 0.0)), np.array((X2, 0, 0.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5ed640e-40fc-48fe-ab9a-b46b5fd97b61",
   "metadata": {},
   "source": [
    "## 2. High-Level Topological Dictionary\n",
    "\n",
    "Instead of manually defining deviation and gap matrices, we define the assembly as a topological graph. The `otaf.AssemblyDataProcessor` interprets the interactions and automatically expands the first-order loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8e1248a3-5053-4030-b165-c48863d8d2b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_data = {\n",
    "    \"PARTS\" : {\n",
    "        '1' : {\n",
    "            \"a\" : {\"FRAME\": np.array([-1*x_, -1*y_, z_]), \"POINTS\": {'A0': P1A0, 'A1': P1A1, 'A2': P1A2}, \"TYPE\": \"plane\", \"INTERACTIONS\": ['P2a'], \"CONSTRAINTS_D\": [\"PERFECT\"], \"CONSTRAINTS_G\": [\"FLOATING\"]},\n",
    "            \"b\" : {\"FRAME\": R0, \"POINTS\": {'B0': P1B0, 'B1': P1B1, 'B2': P1B2}, \"TYPE\": \"plane\", \"INTERACTIONS\": ['P2b'], \"CONSTRAINTS_D\": [\"NONE\"], \"CONSTRAINTS_G\": [\"FLOATING\"]},\n",
    "            \"c\" : {\"FRAME\": np.array([-y_, x_, z_]), \"POINTS\": {'C0': P1C0, 'C1': P1C1, 'C2': P1C2}, \"TYPE\": \"plane\", \"INTERACTIONS\": ['P2c'], \"CONSTRAINTS_D\": [\"PERFECT\"], \"CONSTRAINTS_G\": [\"SLIDING\"]},\n",
    "        },\n",
    "        '2' : {\n",
    "            \"a\" : {\"FRAME\": R0, \"POINTS\": {'A0': P2A0, 'A1': P2A1, 'A2': P2A2}, \"TYPE\": \"plane\", \"INTERACTIONS\": ['P1a'], \"CONSTRAINTS_D\": [\"PERFECT\"], \"CONSTRAINTS_G\": [\"FLOATING\"]},\n",
    "            \"b\" : {\"FRAME\": np.array([-1*x_, -1*y_, z_]), \"POINTS\": {'B0': P2B0, 'B1': P2B1, 'B2': P2B2}, \"TYPE\": \"plane\", \"INTERACTIONS\": ['P1b'], \"CONSTRAINTS_D\": [\"NONE\"], \"CONSTRAINTS_G\": [\"FLOATING\"]},\n",
    "            \"c\" : {\"FRAME\": np.array([y_, -1*x_, z_]), \"POINTS\": {'C0': P2C0, 'C1': P2C1, 'C2': P2C2}, \"TYPE\": \"plane\", \"INTERACTIONS\": ['P1c'], \"CONSTRAINTS_D\": [\"PERFECT\"], \"CONSTRAINTS_G\": [\"SLIDING\"]},\n",
    "        }  \n",
    "    },\n",
    "    \"LOOPS\": {\n",
    "        \"COMPATIBILITY\": {\n",
    "            \"L0\": \"P1cC0 -> P2cC0 -> P2aA0 -> P1aA0\",\n",
    "            \"L1\": \"P1cC0 -> P2cC0 -> P2bB0 -> P1bB0\",\n",
    "        },\n",
    "    },\n",
    "    \"GLOBAL_CONSTRAINTS\": \"2D_NZ\",\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "eda4b960-047d-47fa-a3ae-0a131bef54a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total DOFs: 4\n",
      "Variables: [u_d_4, gamma_d_4, u_d_5, gamma_d_5]\n"
     ]
    }
   ],
   "source": [
    "# Process the topology\n",
    "SDA = otaf.AssemblyDataProcessor(system_data)\n",
    "SDA.generate_expanded_loops()\n",
    "\n",
    "# Extract constraints automatically and embed the slack variable (s)\n",
    "CLH = otaf.CompatibilityLoopHandling(SDA)\n",
    "ILH = otaf.InterfaceLoopHandling(SDA, CLH, circle_resolution=20)\n",
    "SOCAM = otaf.SystemOfConstraintsAssemblyModel(\n",
    "    CLH.get_compatibility_expressions(), \n",
    "    ILH.get_interface_expressions()\n",
    ")\n",
    "SOCAM.embedOptimizationVariable()\n",
    "\n",
    "print(f\"Total DOFs: {len(SOCAM.deviation_symbols)}\")\n",
    "print(f\"Variables: {SOCAM.deviation_symbols}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba18adf2-2501-4ab0-a3d4-d02d6ff30ed5",
   "metadata": {},
   "source": [
    "## 2. Surrogate Training (MLP)\n",
    "\n",
    "To bypass the expensive direct optimizations, we fit a PyTorch surrogate. We draw the distribution parameters directly from the `model1` definition, taking advantage of the `multiply_composed_distribution_standard_with_constants` utility to push the sampling into the failure space for robust training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "296e3160-d1fe-4f84-ab1b-804aac69c5bb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54fed309236b4be4a43f71765491f50a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ratio of failed simulations in sample: 0.14665\n"
     ]
    }
   ],
   "source": [
    "# Fetch distributions and parameters natively from model1\n",
    "jointDist, symbols, max_std_vect, mu_vect = model1.getDistributionParams(tol=0.31, capa=1.0)\n",
    "dim = model1.dim\n",
    "\n",
    "# Artificial dispersion multiplier to ensure limit-state representation\n",
    "mult = 1.35 \n",
    "expanded_dist = otaf.distribution.multiply_composed_distribution_standard_with_constants(jointDist, [mult]*dim)\n",
    "\n",
    "# Generate Sample\n",
    "sample_size = 100000\n",
    "np.random.seed(420)\n",
    "TRAIN_SAMPLE = np.array(expanded_dist.getSample(sample_size), dtype=\"float32\")\n",
    "\n",
    "# Compute exact gap optimizations for training\n",
    "TRAIN_RESULTS = otaf.uncertainty.compute_gap_optimizations_on_sample_batch(\n",
    "    SOCAM, TRAIN_SAMPLE, bounds=None, n_cpu=-2, progress_bar=True, batch_size=500, dtype=\"float32\"\n",
    ")\n",
    "\n",
    "failure_ratio = np.where(TRAIN_RESULTS[:, -1] < 0, 1, 0).sum() / sample_size\n",
    "print(f\"Ratio of failed simulations in sample: {failure_ratio:.5f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "661bb8f1-85e1-4386-8e2a-a95e5645a2ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Neural Network Initialization and Training\n",
    "neural_model = otaf.surrogate.NeuralRegressorNetwork(\n",
    "    input_dim=dim, output_dim=1, X=TRAIN_SAMPLE, y=TRAIN_RESULTS[:, -1], \n",
    "    clamping=True, finish_critertion_epoch=5, loss_finish=1e-6, \n",
    "    metric_finish=0.9998, max_epochs=500, batch_size=30000, \n",
    "    compile_model=False, train_size=0.6, display_progress_disable=False\n",
    ")\n",
    "\n",
    "neural_model.model = torch.nn.Sequential(\n",
    "    *otaf.surrogate.get_custom_mlp_layers([dim, 100, 70, 30, 1], activation_class=torch.nn.GELU)\n",
    ")\n",
    "neural_model.optimizer = torch.optim.AdamW(neural_model.parameters(), lr=0.003, weight_decay=1e-4)\n",
    "otaf.surrogate.initialize_model_weights(neural_model)\n",
    "neural_model.scheduler = torch.optim.lr_scheduler.ExponentialLR(neural_model.optimizer, 1.0001)\n",
    "neural_model.loss_fn = torch.nn.MSELoss()\n",
    "\n",
    "neural_model.train_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82716e8f-1d74-4ace-ab6c-7baded11854b",
   "metadata": {},
   "source": [
    "## 3. Defining the Evaluation Closure and Optimization Logic\n",
    "\n",
    "We define a closure `get_model_evaluator` to map the normalized inputs ($x \\in [0,1]$) directly to the un-normalized model evaluation. \n",
    "\n",
    "We also wrap the optimization execution using the `otaf.optimization.OptimizationTracker` to reliably store all configuration evaluations, successful GLD parameters, and failure probabilities for later visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b954e704-8bfd-4f23-96ba-44692ecac705",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_model_evaluator(sample, mu_vect, surrogate_model): \n",
    "    \"\"\"Closure mapping scaled input x to unstandardized surrogate space.\"\"\"\n",
    "    def evaluate(x):\n",
    "        x_transformed = (sample - mu_vect) * x + mu_vect\n",
    "        prediction = surrogate_model.evaluate_model_non_standard_space(x_transformed)\n",
    "        return np.squeeze(prediction.numpy())\n",
    "    return evaluate\n",
    "\n",
    "@otaf.optimization.scaling(scale_factor=1.0)\n",
    "def optimization_function(\n",
    "        x, failure_slack=0.0, gld=None, model=None, \n",
    "        experiment_key=None, tracker=None, logprob=False, minimize=True):\n",
    "    \n",
    "    multiplier = 1 if minimize else -1\n",
    "    slack = model(x)\n",
    "    gld_params = gld.fit_LMM(slack, disp_fit=False, disp_optimizer=False)\n",
    "    \n",
    "    fp_slack = np.where(slack < failure_slack, 1, 0).mean()\n",
    "    if np.any(np.isnan(gld_params)):\n",
    "        fp_out = fp_slack\n",
    "    else:\n",
    "        fp_out = gld.CDF_num(failure_slack, gld_params, xtol=1e-6)\n",
    "        \n",
    "    tracker.update_objective_data(\n",
    "        exp_key=experiment_key, x=x, fp_gld=fp_out, \n",
    "        fp_slack=fp_slack, gld_params=gld_params, failure_slack=failure_slack\n",
    "    )\n",
    "    \n",
    "    if logprob:\n",
    "        return multiplier * np.log(1e-16 + fp_out)\n",
    "    return multiplier * fp_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69344ac5-4131-45b1-afb5-700e3cee16ee",
   "metadata": {},
   "source": [
    "## 4. Constraint Factory and Feasible Warm Start ($x_0$)\n",
    "\n",
    "Instead of manually constructing tolerance equations, we pull the encapsulated credal set constraint function natively from `model1`. We use SLSQP to push a generic uniform array to the closest valid point on the boundary, serving as a reliable warm start for COBYQA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f318bc00-42e7-4b24-a3c5-addfbd847565",
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimize_scaling_vector(constraint_fn, n_vars, x_warm=None):\n",
    "    \"\"\"Finds the independent scaling factors x closest to 1.0 satisfying the credal constraints.\"\"\"\n",
    "    if x_warm is None:\n",
    "        x_warm = np.full(n_vars, 0.5)\n",
    "        \n",
    "    result = scipy.optimize.minimize(\n",
    "        fun=lambda x: np.sum((x - 1.0) ** 2),\n",
    "        jac=lambda x: 2.0 * (x - 1.0),\n",
    "        x0=x_warm, method=\"SLSQP\",\n",
    "        bounds=Bounds(lb=1e-7, ub=1.0, keep_feasible=True),\n",
    "        constraints={\"type\": \"ineq\", \"fun\": lambda x: -np.array(constraint_fn(x))},\n",
    "        options={\"ftol\": 1e-10, \"maxiter\": 500, \"disp\": False},\n",
    "    )\n",
    "    return result.x\n",
    "\n",
    "def create_directional_constraint(tr, expKey, is_minimization):\n",
    "    lb = 0.0 if is_minimization else -np.inf\n",
    "    ub = np.inf if is_minimization else 0.0\n",
    "    return NonlinearConstraint(\n",
    "        fun=model1.getScaledCredalSetConstraintsFunction(max_std_vect, tr, expKey),\n",
    "        lb=lb, ub=ub, keep_feasible=True\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f1a4b79-e366-416b-982a-aa4a7e0b464b",
   "metadata": {},
   "source": [
    "## 5. Min-Max Global Optimization Loop\n",
    "\n",
    "With the tracker, closure, and constraints configured, we iterate over the alpha levels ($\\alpha_k$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7931fd9-0536-4ab4-81c8-cebb80edc302",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pf_min_max_optimizer(failure_slack, tracker, experiment_key, model_eval_fn, x0, dim):\n",
    "    gld = GLD('VSL')\n",
    "    normalized_bounds = Bounds(1e-9, 1.0, keep_feasible=True)\n",
    "    \n",
    "    print(f\"\\nOptimizing bounds for Slack = {failure_slack}...\")\n",
    "    \n",
    "    res_maxi = scipy.optimize.minimize(\n",
    "        optimization_function, x0,\n",
    "        args=(failure_slack, gld, model_eval_fn, experiment_key, tracker, True, False),\n",
    "        method=\"COBYQA\", bounds=normalized_bounds,\n",
    "        constraints=create_directional_constraint(tracker, experiment_key, False),\n",
    "        options={\"f_target\": -np.inf, \"maxiter\": 1000, \"feasibility_tol\": 1e-5, \"disp\": False}\n",
    "    )\n",
    "    \n",
    "    res_mini = scipy.optimize.minimize(\n",
    "        optimization_function, x0, \n",
    "        args=(failure_slack, gld, model_eval_fn, experiment_key, tracker, True, True),\n",
    "        method=\"COBYQA\", bounds=normalized_bounds,\n",
    "        constraints=create_directional_constraint(tracker, experiment_key, True),\n",
    "        options={\"f_target\": -np.inf, \"maxiter\": 1000, \"feasibility_tol\": 1e-5, \"disp\": False}\n",
    "    )\n",
    "    return res_mini, res_maxi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33d2bd61-cb30-483b-b3c3-cee568ec236d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execution\n",
    "sample_gld = np.array(jointDist.getSample(100000))\n",
    "evaluator = get_model_evaluator(sample_gld, mu_vect, neural_model)\n",
    "tracker = otaf.optimization.OptimizationTracker(bounds=Bounds(1e-7, 1.0), constraint_tolerance=1e-5, precision_decimals=8)\n",
    "\n",
    "# Obtain feasible starting point\n",
    "print(\"Finding feasible start (x0)...\")\n",
    "base_constraint = model1.getScaledCredalSetConstraintsFunction(max_std_vect)\n",
    "x0 = optimize_scaling_vector(base_constraint, dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7b7f9e8-f0c4-4e15-8e12-73526fc7bfa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sweep alpha levels\n",
    "alpha_levels = [0.0, 0.025, 0.05, 0.075]\n",
    "for slack in alpha_levels:\n",
    "    pf_min_max_optimizer(\n",
    "        failure_slack=slack, \n",
    "        tracker=tracker, \n",
    "        experiment_key=f\"model1_slack_{slack}\", \n",
    "        model_eval_fn=evaluator, \n",
    "        x0=x0, \n",
    "        dim=dim\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f9ac320-85fb-4f9f-a8ee-a18711cdae98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the centralized tracker dictionary to CSV for analysis\n",
    "df = tracker.to_dataframe()\n",
    "df.to_csv(\"OptimizationResults_model1_4_dof.csv\")\n",
    "print(\"\\nOptimization Complete. Data saved to OptimizationResults_model1_4_dof.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bd5aeeb-1f6c-467d-b662-9ae356706e2b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}