{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers_merging import *\n",
    "\n",
    "initial_power = 50* si.uW\n",
    "initial_waist = 1.1*si.uW\n",
    "initial_distance = 3*si.um\n",
    "\n",
    "trap: DoubleTweezer = DoubleTweezer(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z= 0*si.G/si.cm,\n",
    "    grad_r=0,\n",
    "    power_tweezer1 = initial_power,     #stationary\n",
    "    power_tweezer2 = initial_power*1.18,     #transfer tweezer\n",
    "    waist_tweezer1 = initial_waist,     #stationary\n",
    "    waist_tweezer2 = initial_waist,     #transfer tweezer\n",
    "    distance_tweezers = initial_distance,\n",
    "\n",
    "    a=180*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "    wvl = 532 * si.nm,\n",
    "\n",
    "    g = 0,\n",
    ")\n",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63864 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -0.5*initial_distance-3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "right_cutoff = 0.5*initial_distance+3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "\n",
    "trap[trap.distance_tweezers] = initial_distance\n",
    "plot_solutions(trap,7,left_cutoff,right_cutoff,display_plot=-1,state_mult=40,n_pot_steps=2000,plot=True,ret_results=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "metadata": {},
   "outputs": [],
   "source": [
    "distances = np.linspace(initial_distance,0*si.um,500)\n",
    "energies, states, potentials = loop_distances(trap, distances,n_levels=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {},
   "outputs": [],
   "source": [
    "index_top = find_crossovers_topstate(energies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {},
   "outputs": [],
   "source": [
    "index, swap_index = find_crossovers(energies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x2a0beb158b0>,\n",
       " <matplotlib.lines.Line2D at 0x2a0beb169f0>,\n",
       " <matplotlib.lines.Line2D at 0x2a0beb145c0>,\n",
       " <matplotlib.lines.Line2D at 0x2a0cb532c30>,\n",
       " <matplotlib.lines.Line2D at 0x2a0cb533fe0>,\n",
       " <matplotlib.lines.Line2D at 0x2a0cb5330b0>,\n",
       " <matplotlib.lines.Line2D at 0x2a0cb532a20>,\n",
       " <matplotlib.lines.Line2D at 0x2a0cb4f4770>,\n",
       " <matplotlib.lines.Line2D at 0x2a0cb4f60c0>,\n",
       " <matplotlib.lines.Line2D at 0x2a0cb4f5100>]"
      ]
     },
     "execution_count": 197,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 45988 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prob_left = np.sum(states[:,:,:int(len(states[0,0])/2)]**2,axis=2)\n",
    "diff = np.abs(np.gradient(prob_left,axis=0))\n",
    "plt.plot(prob_left)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'n_levels = len(energies[0,:])\\nindex = np.array([],dtype=int)\\nswap_index =np.empty((0,2),dtype=int)\\n                    \\nfor i in range(n_levels):\\n    for j in range(i+1,n_levels):\\n        diff = np.abs(energies[:,i] - energies[:,j])\\n        plt.plot(-diff)\\n        #mask = np.where(diff < 5e-2*(np.max(diff)-np.min(diff)),diff,5e-2*(np.max(diff)-np.min(diff)))\\n        #plt.plot(mask,color=\"k\")\\n        peaks = find_peaks(-diff,height=(-1e-33,0))[0]\\n        index = np.append(index,peaks)\\n        for k in range(len(peaks)):\\n            swap_index = np.append(swap_index,[[i,j]],axis=0)'"
      ]
     },
     "execution_count": 174,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"n_levels = len(energies[0,:])\n",
    "index = np.array([],dtype=int)\n",
    "swap_index =np.empty((0,2),dtype=int)\n",
    "                    \n",
    "for i in range(n_levels):\n",
    "    for j in range(i+1,n_levels):\n",
    "        diff = np.abs(energies[:,i] - energies[:,j])\n",
    "        plt.plot(-diff)\n",
    "        #mask = np.where(diff < 5e-2*(np.max(diff)-np.min(diff)),diff,5e-2*(np.max(diff)-np.min(diff)))\n",
    "        #plt.plot(mask,color=\"k\")\n",
    "        peaks = find_peaks(-diff,height=(-1e-33,0))[0]\n",
    "        index = np.append(index,peaks)\n",
    "        for k in range(len(peaks)):\n",
    "            swap_index = np.append(swap_index,[[i,j]],axis=0)\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[17 24 29 34 37 43]\n",
      "[359 351 342]\n",
      "[[2 3]\n",
      " [4 5]\n",
      " [6 7]]\n"
     ]
    }
   ],
   "source": [
    "print(index_top)\n",
    "print(index)\n",
    "print(swap_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'arr = np.abs(np.gradient(np.gradient(energies[:,-1])))\\nplt.plot(arr)\\nindex = find_peaks(arr,width=(1,4))[0]\\nfor i,ind in enumerate(index):\\n    plt.axvline(ind,color=\"k\",alpha=0.7)'"
      ]
     },
     "execution_count": 176,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"arr = np.abs(np.gradient(np.gradient(energies[:,-1])))\n",
    "plt.plot(arr)\n",
    "index = find_peaks(arr,width=(1,4))[0]\n",
    "for i,ind in enumerate(index):\n",
    "    plt.axvline(ind,color=\"k\",alpha=0.7)\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of top state at 4.83um\n",
      "crossover of top state at 4.76um\n",
      "crossover of top state at 4.71um\n",
      "crossover of top state at 4.66um\n",
      "crossover of top state at 4.63um\n",
      "crossover of top state at 4.57um\n",
      "crossover of states [[6 7]] at 1.57um\n",
      "crossover of states [[4 5]] at 1.48um\n",
      "crossover of states [[2 3]] at 1.40um\n"
     ]
    }
   ],
   "source": [
    "new_energies = np.full((energies.shape[0],energies.shape[1]+len(index_top)),np.nan)\n",
    "new_states = np.full((states.shape[0],states.shape[1]+len(index_top),states.shape[2]),np.nan)\n",
    "\n",
    "swapped_index = np.arange(0,energies.shape[1])\n",
    "for i, dist in enumerate(distances): \n",
    "    if np.any(i == index):\n",
    "        j = np.where(index == i)\n",
    "        swapped_index[swap_index[j]] = swapped_index[np.roll(swap_index[j],1)]\n",
    "        print(f\"crossover of states {swap_index[j]} at {dist/si.um:.2f}um\")\n",
    "        \n",
    "    elif np.any(i == index_top):\n",
    "        print(f\"crossover of top state at {dist/si.um:.2f}um\")\n",
    "        swapped_index[-1] = np.max(swapped_index)+1\n",
    "        \n",
    "            \n",
    "    new_energies[i,swapped_index] = energies[i]\n",
    "    new_states[i,swapped_index] = states[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 64424 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot eigenenergies\n",
    "plt.axvline(float(trap.subs(trap.waist_tweezer1))/si.um, color=\"black\",label=\"tweezer waist\")\n",
    "for i in range(len(new_energies[0])):\n",
    "    #differentiate states that are coming from right and left tweezer\n",
    "    if np.sum(new_states[0,i,:int(new_states.shape[2]/2)]**2) < 0.5:\n",
    "        plt.plot(distances/si.um,new_energies[:,i],color=\"red\",label=f\"state {i}\")\n",
    "    else:\n",
    "        plt.plot(distances/si.um,new_energies[:,i],color=\"blue\",label=f\"state {i}\")\n",
    "\n",
    "plt.title(\"Eigenenergies of lowest states at different tweezer distances\")\n",
    "plt.xlabel(\"tweezer distance [um]\")\n",
    "plt.ylabel(\"energy [a.u.]\")\n",
    "#plt.xlim(1.1,2.7)\n",
    "#plt.ylim(-4.0e-29,-1.5e-29)\n",
    "\n",
    "plt.grid()\n",
    "#plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 74396 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot eigenenergies\n",
    "plt.axvline(float(trap.subs(trap.waist_tweezer1))/si.um, color=\"black\",label=\"tweezer waist\")\n",
    "for i in range(len(new_energies[0])):\n",
    "    if np.any(i == np.arange(len(new_energies[0]))[0:20]):\n",
    "        plt.plot(distances/si.um,new_energies[:,i],label=f\"state {i}\")\n",
    "\n",
    "plt.title(\"Eigenenergies of lowest states at different tweezer distances\")\n",
    "plt.xlabel(\"tweezer distance [um]\")\n",
    "plt.ylabel(\"energy [a.u.]\")\n",
    "#plt.xlim(1.1,2.7)\n",
    "#plt.ylim(-4.0e-29,-1.5e-29)\n",
    "\n",
    "plt.grid()\n",
    "#plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
