{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors as mcolors\n",
    "from IPython.display import Math, display\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from scipy import constants as const\n",
    "from scipy.optimize import minimize_scalar, curve_fit\n",
    "from scipy import sparse\n",
    "import qutip as qt\n",
    "\n",
    "#add relative path to backend\n",
    "import sys\n",
    "sys.path.append('../../clean_diag/backend')\n",
    "\n",
    "import trap_units as si\n",
    "from twod_trap import DoubleTweezer, TwoSiteLattice\n",
    "from boson_helpers import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Choose data to analyse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 334\n",
    "end_line = 342\n",
    "\n",
    "n_lines = end_line - start_line + 1\n",
    "lines = np.arange(start_line,end_line+1,dtype=int)\n",
    "state_nrs = np.zeros((n_lines,2),dtype=int)\n",
    "state_nrs[:,1] = 1\n",
    "\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 334\n",
      "separation: 0.7354 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 42596 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 335\n",
      "separation: 0.7379 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 42516 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 336\n",
      "separation: 0.7404 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41392 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 337\n",
      "separation: 0.7429 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41108 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 338\n",
      "separation: 0.7454 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40316 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 339\n",
      "separation: 0.7479 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 39592 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 340\n",
      "separation: 0.7504 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 39116 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 341\n",
      "separation: 0.7529 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 39176 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 342\n",
      "separation: 0.7554 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38548 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "    \n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    print(f\"line: {line}\")\n",
    "    print(f\"separation: {trap.subs(trap.distance_tweezers)/si.um:.4f} um\")\n",
    "\n",
    "    #plot wavefunctions\n",
    "    n_states = 2\n",
    "\n",
    "    states_plot = [GS_left_1.real, GS_right_1.real]\n",
    "    #states_plot = [res[\"states\"][state0].real, res[\"states\"][state1].real]\n",
    "    #states_plot = np.angle(states)\n",
    "    #states_plot = states.imag\n",
    "    #states_plot = np.abs(states)**2\n",
    "\n",
    "    # Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "    fig, axes = plt.subplots(1, n_states, figsize=(6, 4))  # Adjust the size as needed\n",
    "\n",
    "    # Loop over the state numbers from 0 to n_states-1\n",
    "    for state_number, ax in zip(range(n_states), axes):\n",
    "        # Slice through the y-direction and rotate the x-z plane (by swapping x and z)\n",
    "        im = ax.imshow(states_plot[state_number][:, int(res[\"size\"][1]/2), :].T,\n",
    "                    extent=[*res[\"extend\"][0], *res[\"extend\"][2]], origin=\"lower\",\n",
    "                    vmin=np.min(states_plot[state_number]), vmax=np.max(states_plot[state_number]))\n",
    "\n",
    "        # Set labels for each subplot\n",
    "        ax.set_xlabel(\"x\")\n",
    "        ax.set_ylabel(\"z\")\n",
    "        ax.set_title(f\"State {state_number}\")\n",
    "        fig.colorbar(im, ax=ax)  \n",
    "\n",
    "    # Adjust layout for better spacing\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "frequencies = np.zeros_like(lines,dtype=float)\n",
    "distance = np.zeros_like(lines,dtype=float)\n",
    "J = np.zeros_like(lines,dtype=float)\n",
    "#tilt = np.zeros_like(lines,dtype=float)\n",
    "U_s = np.zeros_like(lines,dtype=float)\n",
    "U_dds = np.zeros((len(lines),50),dtype=float)\n",
    "angles = np.zeros((len(lines),50),dtype=float)\n",
    "V_lrs = np.zeros((len(lines),50),dtype=float)\n",
    "DeltaJs = np.zeros((len(lines),50),dtype=float)\n",
    "\n",
    "E_exc = np.zeros_like(lines,dtype=float)\n",
    "E_barr = np.zeros_like(lines,dtype=float)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "\n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    frequencies[i] = float(trap.subs(sp.sqrt(2*trap.power_tweezer1*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer1**2))\n",
    "    distance[i] = float(trap.subs(trap.distance_tweezers))\n",
    "    J[i], U_s[i], U_dds[i], angles[i], V_lrs[i], DeltaJs[i] = analyse_diagonalisation(line,\n",
    "                                                                state_nr0=state0,\n",
    "                                                                state_nr1=state1,\n",
    "                                                                degenerate=degeneracy[i],\n",
    "                                                                n_angles=20)\n",
    "    \n",
    "    #estimate energies of 0th and 1st state\n",
    "    E0 = 0.5*(res[\"energies\"][0] + res[\"energies\"][1])\n",
    "    E1 = 0.5*(res[\"energies\"][2] + res[\"energies\"][3])\n",
    "    E_exc[i] = E1-E0\n",
    "\n",
    "    #estimate heigth of barrier\n",
    "    V = trap.subs(trap.get_potential(apply_zero_offset=False))\n",
    "    a = float(trap.subs(trap.distance_tweezers))\n",
    "\n",
    "    #find minima of potential\n",
    "    def V_func(x):\n",
    "        return float(V.subs({trap.x:x, trap.y:0, trap.z:0}))\n",
    "\n",
    "    x_right = minimize_scalar(V_func,bracket=[0,a]).x\n",
    "    x_left = minimize_scalar(V_func,bracket=[-a,0]).x\n",
    "\n",
    "    #catch case where both potentials have already merged\n",
    "    tunneling_dist = abs(x_right-x_left)\n",
    "    if tunneling_dist < 1e-15:\n",
    "        raise Exception(\"potential has only one minmum\")\n",
    "\n",
    "    #depts of both tweezers\n",
    "    V_1 = abs(V_func(x_left) - V_func(0))\n",
    "    V_2 = abs(V_func(x_right) - V_func(0)) \n",
    "\n",
    "    E_barr[i] = np.max([V_1,V_2])\n",
    "\n",
    "    # these states are basically 1 tweezer\n",
    "    if state1 !=1:\n",
    "        J[i] = np.nan\n",
    "        U_s[i] = np.nan\n",
    "        U_dds[i] = np.nan\n",
    "        V_lrs[i] = np.nan\n",
    "        DeltaJs[i] = np.nan\n",
    "        E_exc[i] = np.nan\n",
    "        E_barr[i] = np.nan\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 315436 characters]",
      "text/plain": [
       "<Figure size 1000x1500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create a figure with 4 subplots (stacked vertically)\n",
    "fig, axes = plt.subplots(4, 1, figsize=(10, 15), sharex=True)\n",
    "\n",
    "# Get unique frequencies\n",
    "unique_freqs = np.unique(frequencies)\n",
    "\n",
    "# Define colors for consistency\n",
    "colors = [\"orange\", \"blue\", \"green\", \"purple\", \"orange\"]\n",
    "\n",
    "### 1️⃣ Hopping parameter J ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    axes[0].plot(distance[mask] / si.nm, J[mask] / const.h, color=color, \n",
    "                 label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "#axes[0].axhspan(50, 150, alpha=0.3, color=\"gray\", label=\"desired region of 50-100Hz\")\n",
    "axes[0].set_ylabel(r\"$J/h$ [Hz]\")\n",
    "axes[0].xaxis.set_label_position('top')\n",
    "axes[0].set_xlabel(\"tweezer spacing [nm]\")\n",
    "axes[0].legend()\n",
    "# 🔹 Add top x-axis ticks\n",
    "axes[0].tick_params(axis=\"x\", labeltop=True, top=True)\n",
    "axes[0].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[0].minorticks_on()\n",
    "axes[0].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "### 2️⃣ Interaction parameter U ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    x = distance[mask] / si.nm\n",
    "    y = U_s[mask] / const.h\n",
    "    delta = U_dds[mask] / const.h\n",
    "\n",
    "    axes[1].plot(x, y, color=color, label=fr\"$U_s$ for $\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "    axes[1].fill_between(x, y + np.min(delta, axis=1), y + np.max(delta, axis=1), \n",
    "                         alpha=0.3, color=color, edgecolor=None, label=r\"$\\Delta U_{dd}$\")\n",
    "\n",
    "axes[1].set_ylabel(r\"$U/h$ [Hz]\")\n",
    "axes[1].legend()\n",
    "axes[1].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[1].minorticks_on()\n",
    "axes[1].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "### 3️⃣ First excited state energy ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    axes[2].plot(distance[mask] / si.nm, E_exc[mask] / const.h, color=color, \n",
    "                 label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "axes[2].set_ylabel(r\"$E_{exc}/h$ [Hz]\")\n",
    "axes[2].legend()\n",
    "axes[2].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[2].minorticks_on()\n",
    "axes[2].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "### 4️⃣ Barrier height ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    axes[3].plot(distance[mask] / si.nm, E_barr[mask] / const.h, color=color, \n",
    "                 label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "axes[3].set_ylabel(r\"$E_{barr}/h$ [Hz]\")\n",
    "axes[3].set_xlabel(\"tweezer spacing [nm]\")\n",
    "#axes[3].set_ylim(0, 30000)\n",
    "axes[3].legend()\n",
    "axes[3].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[3].minorticks_on()\n",
    "axes[3].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "# Improve layout\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
