{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b01c5c92-e134-42f9-bf8e-5c3f5b553ef7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pyright: reportUnknownArgumentType=false\n",
    "from rich.theme import Theme\n",
    "from rich.console import Console\n",
    "\n",
    "from finesse.model import Model\n",
    "from finesse.analysis.actions import (\n",
    "    TemporaryParameters,\n",
    "    Change,\n",
    "    Maximize,\n",
    "    Minimize,\n",
    "    Series,\n",
    "    FrequencyResponse,\n",
    "    Xaxis,\n",
    "    Noxaxis,\n",
    ")\n",
    "from finesse.solutions import SeriesSolution\n",
    "\n",
    "from matplotlib.pyplot import figure, show\n",
    "from matplotlib.axes import Axes\n",
    "\n",
    "from numpy import geomspace, linspace\n",
    "\n",
    "from numpy.typing import NDArray\n",
    "from typing import Any, NamedTuple\n",
    "\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0b4199a-b9b0-4969-97c7-497faa662b94",
   "metadata": {},
   "outputs": [],
   "source": [
    "from gettext import install\n",
    "from logging import getLogger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71f985ab-ce47-4621-86c5-d4dc41a267d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "install(__name__)\n",
    "logger = getLogger(__name__)\n",
    "theme = Theme(\n",
    "    {\n",
    "        \"strong\": \"cyan underline\",\n",
    "        \"result\": \"red bold\",\n",
    "        \"error\": \"red underline bold\",\n",
    "    }\n",
    ")\n",
    "console = Console(theme=theme)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b5dbd67-3baa-4cba-a5a2-368d306735a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "C_DARK_FRINGE = 8e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6999b95-7956-4d2a-a97e-93f3d12f5c74",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib ipympl\n",
    "model_file = Path(\"model.kat\")\n",
    "model = Model()\n",
    "model.phase_config(zero_k00=False, zero_tem00_gouy=True)\n",
    "model.modes(modes=\"off\")  # pyright: ignore[reportUnusedCallResult]\n",
    "model.parse(model_file.read_text())\n",
    "model.lambda0 = model.get(\"wavelength\")\n",
    "model.plot_graph()  # pyright: ignore[reportUnusedCallResult]\n",
    "show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2b2ac68-0730-4f68-a4b1-ca036eadc880",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = model.run(\n",
    "    TemporaryParameters(\n",
    "        Series(\n",
    "            Change(\n",
    "                {\n",
    "                    \"SR.misaligned\": True,\n",
    "                    \"PR.misaligned\": True,\n",
    "                    \"eom1.midx\": 0,\n",
    "                    \"eom2.midx\": 0,\n",
    "                    \"eom3.midx\": 0,\n",
    "                    \"eom4.midx\": 0,\n",
    "                }\n",
    "            ),\n",
    "            Maximize(\n",
    "                model.get(\"NE_p1\"),\n",
    "                model.get(\"NORTH_ARM.DC\"),\n",
    "                bounds=[-180, 180],\n",
    "                tol=1e-14,\n",
    "            ),\n",
    "            Maximize(\n",
    "                model.get(\"WE_p1\"),\n",
    "                model.get(\"WEST_ARM.DC\"),\n",
    "                bounds=[-180, 180],\n",
    "                tol=1e-14,\n",
    "            ),\n",
    "            Minimize(\n",
    "                model.get(\"SR_p2\"), model.get(\"MICH.DC\"), bounds=[-180, 180], tol=1e-14\n",
    "            ),\n",
    "            Change(\n",
    "                {\n",
    "                    \"PR.misaligned\": False,\n",
    "                }\n",
    "            ),\n",
    "            Maximize(\n",
    "                model.get(\"PR_p2\"), model.get(\"PRCL.DC\"), bounds=[-180, 180], tol=1e-14\n",
    "            ),\n",
    "            Change(\n",
    "                {\n",
    "                    \"SR.misaligned\": False,\n",
    "                }\n",
    "            ),\n",
    "            Maximize(\n",
    "                model.get(\"B1_DC\"), model.get(\"SRCL.DC\"), bounds=[-180, 180], tol=1e-14\n",
    "            ),\n",
    "            Change(\n",
    "                {\n",
    "                    \"SRCL.DC\": -90,\n",
    "                },\n",
    "                relative=True,\n",
    "            ),\n",
    "        ),\n",
    "        exclude=[\n",
    "            \"NE.phi\",\n",
    "            \"NI.phi\",\n",
    "            \"WE.phi\",\n",
    "            \"WI.phi\",\n",
    "            \"SR.phi\",\n",
    "            \"PR.phi\",\n",
    "            \"NORTH_ARM.DC\",\n",
    "            \"WEST_ARM.DC\",\n",
    "            \"DARM.DC\",\n",
    "            \"MICH.DC\",\n",
    "            \"PRCL.DC\",\n",
    "            \"SRCL.DC\",\n",
    "            \"SR.misaligned\",\n",
    "            \"eom1.midx\",\n",
    "            \"eom2.midx\",\n",
    "            \"eom3.midx\",\n",
    "            \"eom4.midx\",\n",
    "        ],\n",
    "    ),\n",
    ")\n",
    "model._settings.phase_config.zero_k00 = False\n",
    "model.fsig.f = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7504821e-b33d-4396-b86c-06e90477eb8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_solutions(\n",
    "    model: Model, DOF: str, padding: float, nb: int = 10000\n",
    ") -> SeriesSolution:\n",
    "    return model.run(\n",
    "        Xaxis(\n",
    "            model.get(DOF).DC,\n",
    "            \"lin\",\n",
    "            model.get(DOF).DC - padding,\n",
    "            model.get(DOF).DC + padding,\n",
    "            nb,\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "def display_ax(\n",
    "    ax: Axes,\n",
    "    solution: SeriesSolution,\n",
    "    model: Model,\n",
    "    DOF: str,\n",
    "    padding: float,\n",
    "    nb: int = 10000,\n",
    ") -> Axes:\n",
    "    x = linspace(model.get(DOF).DC - padding, model.get(DOF).DC + padding, nb + 1)\n",
    "    _ = ax.semilogy(x, solution[\"SR_p2\"], label=\"dark fringe\")\n",
    "    _ = ax.semilogy(x, solution[\"NE_p1\"], label=\"north cavity\")\n",
    "    _ = ax.semilogy(x, solution[\"WE_p1\"], label=\"west cavity\")\n",
    "    _ = ax.vlines(\n",
    "        [model.get(DOF).DC],\n",
    "        min(solution[\"SR_p2\"]),\n",
    "        max(solution[\"NE_p1\"]),\n",
    "        colors=\"red\",\n",
    "    )\n",
    "    _ = ax.set_ylabel(\"power (W)\")\n",
    "    ax.grid()\n",
    "    _ = ax.legend()\n",
    "    return ax\n",
    "\n",
    "\n",
    "class DisplayData(NamedTuple):\n",
    "    DOF: str\n",
    "    padding: float\n",
    "\n",
    "\n",
    "data: list[DisplayData] = [\n",
    "    DisplayData(\"NORTH_ARM\", 10),\n",
    "    DisplayData(\"WEST_ARM\", 10),\n",
    "    DisplayData(\"PRCL\", 10),\n",
    "    DisplayData(\"MICH\", 10),\n",
    "    DisplayData(\"DARM\", 10),\n",
    "    DisplayData(\"CARM\", 10),\n",
    "]\n",
    "\n",
    "Figure = figure(figsize=(13, 10))\n",
    "nb = int(1e4)\n",
    "\n",
    "for i in range(len(data)):\n",
    "    element: DisplayData = data[i]\n",
    "    ax = Figure.add_subplot(3, 2, i + 1)\n",
    "    solution = compute_solutions(model, element.DOF, element.padding, nb)\n",
    "    _ = display_ax(ax, solution, model, element.DOF, element.padding, nb).set_xlabel(\n",
    "        \"{} value\".format(element.DOF)\n",
    "    )\n",
    "show()\n",
    "\n",
    "solution = model.run(Noxaxis())\n",
    "result = solution[\"B1_DC\"]\n",
    "start, stop, nb = 0, 1, 0\n",
    "while (abs(result - C_DARK_FRINGE) > 1e-4) and (nb < 100):\n",
    "    nb += 1\n",
    "    temp = start + (stop - start) / 2\n",
    "\n",
    "    model.DARM.DC = temp\n",
    "    solution = model.run(Noxaxis())\n",
    "    result = solution[\"B1_DC\"]\n",
    "    if result > C_DARK_FRINGE:\n",
    "        stop = temp\n",
    "    else:\n",
    "        start = temp\n",
    "console.print(\n",
    "    \"Degré de liberté [result]{dof}[/result] trouvé en [strong]{nb} pas[/strong] pour avoir une puissance de [result]{result} W[/result] sur B1\".format(\n",
    "        nb=nb, dof=model.DARM.DC, result=result\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fbf07f1-22bc-445b-95ad-18a30adcc85d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_evolution(parameter: str, model: Model, values: NDArray[Any], TFs: list[str]):\n",
    "    model.SNEB.phi = model.NE.phi - 45\n",
    "    model.SWEB.phi = model.WE.phi - 45\n",
    "    model.SDB1.phi = model.SR.phi + 45\n",
    "\n",
    "    if len(TFs) == 0:\n",
    "        console.print(\"[error]Nothing to show[/error]\")\n",
    "\n",
    "    Figure = figure(figsize=(14, 5 * len(TFs)))\n",
    "    Figure.suptitle(\"TF in function of {}\".format(parameter))\n",
    "\n",
    "    for i in range(len(TFs)):\n",
    "        _ = Figure.add_subplot(len(TFs), 1, i + 1)\n",
    "\n",
    "    temp_value = model.get(parameter).eval()\n",
    "    for value in values:\n",
    "        index = 0\n",
    "        model.set(parameter, value)\n",
    "\n",
    "        DARM = model.run(\n",
    "            FrequencyResponse(geomspace(5, 10000, 1000), [\"DARM\"], [\"B1.I\"])\n",
    "        )\n",
    "        for bench in [\"SNEB\", \"SWEB\", \"SDB1\"]:\n",
    "            if bench in TFs:\n",
    "                result = model.run(\n",
    "                    FrequencyResponse(\n",
    "                        geomspace(5, 10000, 1000), [\"{}_z\".format(bench)], [\"B1.I\"]\n",
    "                    )\n",
    "                )\n",
    "                _ = Figure.get_axes()[index].set_title(bench)\n",
    "                _ = Figure.get_axes()[index].loglog(\n",
    "                    result.f,\n",
    "                    abs(result[\"B1.I\", \"{}_z\".format(bench)])\n",
    "                    / abs(DARM[\"B1.I\", \"DARM\"])\n",
    "                    / model.space_NI_NE.L.eval(),\n",
    "                    label=\"{}\".format(value),\n",
    "                )\n",
    "\n",
    "                index += 1\n",
    "    for i in range(len(TFs)):\n",
    "        _ = Figure.get_axes()[i].grid(True, \"both\", \"both\")\n",
    "        _ = Figure.get_axes()[i].legend()\n",
    "    show()\n",
    "    model.set(parameter, temp_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2b8bc8d-dd73-4a45-af9d-179b3f7a3c7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_evolution(\"NE.T\", model, geomspace(1e-7, 1e-3, 10), [\"SNEB\", \"SWEB\", \"SDB1\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a514ded-0777-41a7-82ce-380894d7be44",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_evolution(\"WE.T\", model, geomspace(1e-7, 1e-3, 10), [\"SWEB\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a7d30d4-4d9f-4715-a1d6-a0330e8450a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_evolution(\"SR.T\", model, linspace(0.30, 0.50, 10), [\"SDB1\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7dca854-5dff-483f-8a72-efaf0b80d87d",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_evolution(\"NI.T\", model, linspace(1.35e-2, 1.39e-2, 10), [\"SNEB\", \"SWEB\", \"SDB1\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b31a0f43-0a07-44b4-9264-b5001350eb9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_evolution(\"WI.T\", model, linspace(1.35e-2, 1.39e-2, 10), [\"SNEB\", \"SWEB\", \"SDB1\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27b2b8da-2334-4aaa-9e6b-537490598f9b",
   "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.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}