{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating positioned roughness elements"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read parameters from file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'bounding_box': {'end': (2400.0, 1000.0, 0.0),\n",
      "                  'start': (-500.0, -1000.0, 0.0)},\n",
      " 'element_params': {'height': 5.0, 'width': 5.0},\n",
      " 'spacing_params': {'line_offset': 10.0,\n",
      "                    'offset_direction': 'x',\n",
      "                    'spacing': (20.0, 20.0)},\n",
      " 'surfaces': {'disk': './fixtures/tests/roughness_gen/disk/disk.lnas',\n",
      "              'loft': './fixtures/tests/roughness_gen/loft/loft.lnas'}}\n"
     ]
    }
   ],
   "source": [
    "from cfdmod.use_cases.roughness_gen.parameters import PositionParams\n",
    "import pathlib\n",
    "\n",
    "import pprint\n",
    "\n",
    "pp = pprint.PrettyPrinter()\n",
    "\n",
    "cfg_file = pathlib.Path(\"./fixtures/tests/roughness_gen/position_params.yaml\")\n",
    "cfg = PositionParams.from_file(cfg_file)\n",
    "\n",
    "output_path = pathlib.Path(\"./output/roughness_gen\")\n",
    "\n",
    "pp.pprint(cfg.model_dump())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read mesh files and get surfaces bounding box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ -500.        , -1000.        ,   751.74133301],\n",
       "       [ 2400.        ,  1000.        ,     0.        ]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from lnas import LnasFormat\n",
    "import trimesh\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "surfaces_read: dict[str, LnasFormat] = {}\n",
    "surfaces_mesh: dict[str, trimesh.Trimesh] = {}\n",
    "\n",
    "mesh_bbox = [\n",
    "    (float(\"inf\"), float(\"inf\"), float(\"inf\")),\n",
    "    (float(\"-inf\"), float(\"-inf\"), float(\"-inf\")),\n",
    "]\n",
    "\n",
    "for sfc, sfc_path in cfg.surfaces.items():\n",
    "    lnas = LnasFormat.from_file(pathlib.Path(sfc_path))\n",
    "    surfaces_read[sfc] = lnas\n",
    "    surfaces_mesh[sfc] = trimesh.Trimesh(\n",
    "        vertices=lnas.geometry.vertices, faces=lnas.geometry.triangles\n",
    "    )\n",
    "\n",
    "    min_point = surfaces_read[sfc].geometry.vertices.min(axis=0)\n",
    "    max_point = surfaces_read[sfc].geometry.vertices.max(axis=0)\n",
    "\n",
    "    if any(min_point < mesh_bbox[0]):\n",
    "        mesh_bbox[0] = np.array([min_point, mesh_bbox[0]]).min(axis=0)\n",
    "    if any(max_point > mesh_bbox[1]):\n",
    "        mesh_bbox[1] = np.array([max_point, mesh_bbox[1]]).max(axis=0)\n",
    "\n",
    "s = np.array([cfg.bounding_box.start, mesh_bbox[0]]).max(axis=0)\n",
    "e = np.array([cfg.bounding_box.end, mesh_bbox[1]]).min(axis=0)\n",
    "bbox_to_use = np.array([s, e])\n",
    "\n",
    "bbox_to_use"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate generation parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(145, 80, 2900.0, 2000.0)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lx = bbox_to_use[1][0] - bbox_to_use[0][0]\n",
    "ly = bbox_to_use[1][1] - bbox_to_use[0][1]\n",
    "\n",
    "if cfg.spacing_params.offset_direction == \"x\":\n",
    "    Nx = int((lx - cfg.spacing_params.line_offset) // (cfg.spacing_params.spacing[0]) + 1)\n",
    "    Ny = int(\n",
    "        (ly + cfg.spacing_params.spacing[1])\n",
    "        // (cfg.element_params.width + cfg.spacing_params.spacing[1])\n",
    "    )\n",
    "else:\n",
    "    Nx = int(lx // (cfg.spacing_params.spacing[0]) + 1)\n",
    "    Ny = int(\n",
    "        (ly + cfg.spacing_params.spacing[1] - cfg.spacing_params.line_offset)\n",
    "        // (cfg.element_params.width + cfg.spacing_params.spacing[1])\n",
    "    )\n",
    "\n",
    "Nx, Ny, lx, ly"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup generation parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cfdmod.use_cases.roughness_gen.parameters import GenerationParams\n",
    "\n",
    "generation_params = GenerationParams(\n",
    "    N_elements_x=Nx,\n",
    "    N_elements_y=Ny,\n",
    "    element_params=cfg.element_params,\n",
    "    spacing_params=cfg.spacing_params,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build elements and apply linear patterns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Single element: Vertices count: 6 Triangles count: 2\n",
      "Single line Vertices count: 870 Triangles count: 290\n",
      "Replicated lines Vertices count: 69600 Triangles count: 23200\n"
     ]
    }
   ],
   "source": [
    "from cfdmod.use_cases.roughness_gen import build_single_element, linear_pattern\n",
    "\n",
    "triangles, normals = build_single_element(generation_params.element_params)\n",
    "\n",
    "single_line_triangles, single_line_normals = linear_pattern(\n",
    "    triangles,\n",
    "    normals,\n",
    "    direction=generation_params.spacing_params.offset_direction,\n",
    "    n_repeats=generation_params.single_line_elements,\n",
    "    spacing_value=generation_params.single_line_spacing,\n",
    ")\n",
    "\n",
    "full_triangles, full_normals = linear_pattern(\n",
    "    single_line_triangles,\n",
    "    single_line_normals,\n",
    "    direction=generation_params.perpendicular_direction,\n",
    "    n_repeats=generation_params.multi_line_elements,\n",
    "    spacing_value=generation_params.multi_line_spacing,\n",
    "    offset_value=generation_params.spacing_params.line_offset,\n",
    ")\n",
    "\n",
    "# Offset to match bounding box limits\n",
    "full_triangles[:, :, 0] += bbox_to_use[0][0]\n",
    "full_triangles[:, :, 1] += bbox_to_use[0][1]\n",
    "\n",
    "print(\n",
    "    \"Single element:\",\n",
    "    f\"Vertices count: {triangles.shape[0] * triangles.shape[1]}\",\n",
    "    f\"Triangles count: {len(triangles)}\",\n",
    ")\n",
    "print(\n",
    "    \"Single line\",\n",
    "    f\"Vertices count: {single_line_triangles.shape[0] * single_line_triangles.shape[1]}\",\n",
    "    f\"Triangles count: {len(single_line_triangles)}\",\n",
    ")\n",
    "print(\n",
    "    \"Replicated lines\",\n",
    "    f\"Vertices count: {full_triangles.shape[0] * full_triangles.shape[1]}\",\n",
    "    f\"Triangles count: {len(full_triangles)}\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create profiles with x normal planes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_vals = np.unique(full_triangles[:, :, 0].reshape(-1))\n",
    "\n",
    "profiles: dict[float, np.ndarray] = {}\n",
    "\n",
    "for sfc_label, mesh in surfaces_mesh.items():\n",
    "    x_vals_to_use = x_vals[\n",
    "        (x_vals >= mesh.vertices[:, 0].min()) & (x_vals <= mesh.vertices[:, 0].max())\n",
    "    ]\n",
    "\n",
    "    for x_val in x_vals_to_use:\n",
    "        section_slice = mesh.section(\n",
    "            plane_origin=[x_val, 0, 0],\n",
    "            plane_normal=[-1, 0, 0],\n",
    "        )\n",
    "        vertices = np.array(section_slice.to_dict()[\"vertices\"])\n",
    "        if x_val in profiles.keys():\n",
    "            profiles[x_val] = np.concatenate((profiles[x_val], vertices))\n",
    "        else:\n",
    "            profiles[x_val] = vertices\n",
    "\n",
    "for k, profile in profiles.items():\n",
    "    profile = np.array(sorted(profile, key=lambda x: x[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Offset geometry with surface information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "for tri in full_triangles:\n",
    "    left_bound = tri[:, 1].min()\n",
    "    right_bound = tri[:, 1].max()\n",
    "    prof = profiles[tri[0][0]][\n",
    "        (profiles[tri[0][0]][:, 1] >= left_bound) & (profiles[tri[0][0]][:, 1] <= right_bound)\n",
    "    ]\n",
    "    while len(prof) == 0:\n",
    "        value_offset = generation_params.element_params.width\n",
    "        left_bound -= value_offset\n",
    "        right_bound += value_offset\n",
    "        prof = profiles[tri[0][0]][\n",
    "            (profiles[tri[0][0]][:, 1] >= left_bound) & (profiles[tri[0][0]][:, 1] <= right_bound)\n",
    "        ]\n",
    "\n",
    "    offset_val = prof[:, 2].min()\n",
    "    tri[:, 2] += offset_val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Export generated geometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cfdmod.api.geometry.STL import export_stl\n",
    "\n",
    "export_stl(output_path / \"positioned_elements.stl\", full_triangles, full_normals)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cfdmod-XMkUSlb0-py3.10",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}