import warnings
from collections.abc import Callable
from functools import cache
from itertools import permutations
from numbers import Number
from typing import Any
import datetime
import pickle

import numpy as np
import numpy.typing as npt
import sympy as sp
from scipy import constants as const
from scipy.sparse import save_npz
from scipy.optimize import minimize_scalar
from sympy import LT, default_sort_key, topological_sort
from sympy.core.numbers import Zero

import trap_units as si
from trap_numerics import (nstationary_solution,
                           nstationary_solution_export,
                           discrete_hamiltonian,
                           functional_hamiltonian)
import trap_potential as pot
import trap_properties as p
from trap_typing import float_or_array

Numeric = int | float | complex | np.number | sp.Number
ExprNum = sp.Expr | Numeric


class PancakeTrap:
    #: Spacial coordinates
    x, y, z = sp.symbols("x, y, z", real=True, finite=True)
    #: Offset from the waist location in z
    d_0 = sp.symbols("d_0", real=True, finite=True)
    #: Beam waist in x-direction (collimated).
    waist_x0 = sp.Symbol("W_{x0}", real=True, positive=True, finite=True, nonzero=True)
    #: Beam waist in z-direction at the focus.
    waist_z0 = sp.Symbol("W_{z0}", real=True, positive=True, finite=True, nonzero=True)
    #: Waist of the tweezer (at the focus)
    waist_tweezer = sp.Symbol(
        "W_t", real=True, positive=True, finite=True, nonzero=True
    )
    #: Power per beam. Total power is 2x.
    power = sp.Symbol("P", real=True, positive=True, finite=True, nonzero=True)
    #: Power of the tweezer beam
    power_tweezer = sp.Symbol(
        "P_t", real=True, positive=True, finite=True, nonzero=True
    )
    #magnetic offset field
    B_x, B_y, B_z = sp.symbols(
        "B_x, B_y, B_z", real=True, finite=True
        )
    #: Magnetic field gradients.
    grad_r, grad_z = sp.symbols(
        "\\nabla{B_x}, \\nabla{B_z}", real=True, finite=True, nonzero=True
    )
    #: Real part of polarizability :math:`\alpha`
    a = sp.Symbol("a", real=True, finite=True, nonzero=True, positive=True)
    #: Pancake spacing
    d = sp.Symbol("d", real=True, positive=True, finite=True, nonzero=True)
    #: Bohr magneton
    mu_b = sp.Symbol("mu_b", real=True, positive=True, finite=True, nonzero=True)
    #: Wavelength
    wvl = sp.Symbol("lambda", real=True, positive=True, finite=True, nonzero=True)
    #: Angle of the incident beams
    theta = sp.Symbol("theta", real=True, positive=True, finite=True, nonzero=True)
    #: Mass of the trapped atoms
    m = sp.Symbol("m", real=True, positive=True, finite=True, nonzero=True)
    #: contact-interaction scattering length of atoms
    a_s = sp.Symbol("a_s", real=True, positive=True, finite=True, nonzero=True)
    #: Trap frequency in each spacial direction
    omega_x, omega_y, omega_z = sp.symbols(
        "omega_x, omega_y, omega_z", real=True, positive=True, finite=True, nonzero=True
    )
    #: Trap frequency of the tweezer
    omega_r_tweezer, omega_ax_tweezer = sp.symbols(
        "omega_t_r, omega_t_ax", real=True, positive=True, finite=True, nonzero=True
    )
    #: Resonance frequency
    omega_0 = sp.Symbol("omega_0", real=True, positive=True, finite=True, nonzero=True)
    #: Laser frequency
    omega_l = sp.Symbol("omega_l", real=True, positive=True, finite=True, nonzero=True)
    #: Decay rate gamma
    gamma = sp.Symbol("Gamma", real=True, positive=True, finite=True, nonzero=True)
    #: Angle of the magnetic field gradient
    beta = sp.Symbol("beta", real=True, positive=True, finite=True)
    #: gravitational acceleration (set 0 if not needed)
    g = sp.Symbol("g", real=True, positive=True, finite=True)

    #: Default substitutions applied using :meth:`.subs`
    _defaults: dict[str, Any] = {
        #"a": (
        #    (3 * sp.pi * const.c**2)
        #    / (2 * omega_0**3)
        #    * (gamma / (omega_0 - omega_l) + gamma / (omega_0 + omega_l))
        #),
        "d": wvl / (2 * sp.sin(theta)),
        "omega_l": 2 * np.pi * const.c / wvl,
        #"omega_0": 2 * np.pi * const.c / (626 * si.nm),
        "d_0": 0,
        "wvl": 532 * si.nm, #1064 * si.nm,
        "theta": 0, #np.deg2rad(7.35),
        "beta": 0,
        "gamma": 2 * np.pi * 135 * const.kilo, #2 * np.pi * 5.8724 * const.mega,
        "grad_r": 0.0066 * grad_z,  # see labbook 2023-11-21
        "grad_z": 85 * si.G / si.cm,
        "B_x": 0*si.G,
        "B_y": 0*si.G,
        "B_z": 1e-14*si.G, #apply small field to avoid zero

        "m": 164 * const.value("atomic mass constant"), #6.0151228 * const.value("atomic mass constant"),
        "a_s": 85* const.value("Bohr radius"),
        "mu_b": 9.93 * const.value("Bohr magneton" ), #using mu_b as mu for the atom
        "g": const.g,
    }

    def __init__(self, **values: Numeric):
        self._values: dict[sp.Symbol, Numeric] = {
            self._get_key_symbol(key): value for key, value in self._defaults.items()
        }
        for key, value in values.items():
            self[key] = value
        self._omega_solution = {}

    def subs(self, expr: sp.Expr, symbolic_only: bool = False) -> sp.Expr:
        # Only substitute symbols that have a value of type Expr,
        # i.e. they are expressed through other symbols.
        symbolic_values: list[tuple[sp.Symbol, sp.Expr | Numeric]] = [
            (key, value)
            for key, value in self._values.items()
            if isinstance(value, sp.Expr) or value == 0
        ]
        # Find all edges where an expression contains the symbol of
        # a different subsitution.
        edges = [
            (i, j)
            for i, j in permutations(symbolic_values, 2)
            if isinstance(i[1], sp.Expr) and i[1].has(j[0])
        ]
        # Sort the substitutions topologically such that
        sorted_symbolic_values = topological_sort(
            [symbolic_values, edges], default_sort_key
        )
        expr_symbolic = expr.subs(sorted_symbolic_values)
        if symbolic_only:
            return expr_symbolic
        else:
            # Also substitute all non-symbolic values
            return expr_symbolic.subs(
                [
                    (key, value)
                    for key, value in self._values.items()
                    if not isinstance(value, sp.Expr)
                ]
            )

    def __setitem__(self, key: str | sp.Symbol, value):
        self._values[self._get_key_symbol(key)] = value

    def __getitem__(self, item: str | sp.Symbol):
        return self._values[self._get_key_symbol(item)]

    def _get_key_symbol(self, key: str | sp.Symbol) -> sp.Symbol:
        if isinstance(key, sp.Symbol):
            return key
        if not isinstance(key, str):
            raise ValueError(
                f"Expected type 'str' or 'sympy.Symbol' but got {type(key)!s}"
            )
        if not hasattr(self, key):
            raise ValueError(f"Unknown attribute '{key!s}'")
        key_symbol = getattr(self, key)
        if not isinstance(getattr(self, key), sp.Symbol):
            raise ValueError(f"Attribute '{key!s}' is not a symbol")
        return key_symbol

    def get_tweezer_rayleigh(self) -> sp.Expr:
        return sp.pi * self.waist_tweezer**2 / self.wvl

    def get_tweezer_intensity(self) -> sp.Expr:
        rayleigh_length = self.get_tweezer_rayleigh()
        waist = self.waist_tweezer * sp.sqrt(1 + (self.z / rayleigh_length) ** 2)
        i_0 = 2 * self.power_tweezer / (sp.pi * self.waist_tweezer**2)
        r = sp.sqrt(self.x**2 + self.y**2)
        return i_0 * (self.waist_tweezer / waist) ** 2 * sp.exp(-2 * r**2 / (waist**2))

    def get_tweezer_potential(self) -> sp.Expr:
        return -self.a * self.get_tweezer_intensity()

    def get_omega_tweezer(self, omega: sp.Symbol) -> sp.Expr:
        """Get tweezer trap frequency"""
        return self.get_omega(omega, with_tweezer=True, with_2dtrap=False)

    def get_omega_r_tweezer(self) -> sp.Expr:
        return self.get_omega_tweezer(self.omega_r_tweezer)

    def get_omega_ax_tweezer(self) -> sp.Expr:
        return self.get_omega_tweezer(self.omega_ax_tweezer)

    def get_beam_waist_z(self, length: sp.Symbol | Number) -> sp.Expr:
        """The beam is focused in z-direction, leading to a waist that
        depends on the distance :math:`l` from the lens.

        Depending on which beam (upper or lower), the distance from the
        lens :math:`l` is given by
        :math:`y \\cos(\\theta) \\mp z \\sin(\\theta)`

        :param length: Distance :math:`l` from the lens.
        """
        z_r = sp.pi * self.waist_z0**2 / self.wvl
        return self.waist_z0 * sp.sqrt(1 + ((length - self.d_0) / z_r) ** 2)

    def get_upper_beam_intensity(self) -> sp.Expr:
        """Intensity distribution of the upper beam"""
        length = self.y * sp.cos(self.theta) - self.z * sp.sin(self.theta)
        waist_z = self.get_beam_waist_z(length)
        # Intensity of th0e beam at the origin
        i_0 = 2 * self.power / (sp.pi * self.waist_x0 * self.waist_z0)
        # Turn off automatic formatting to preserve readability
        # fmt: off
        return i_0 * self.waist_z0 / waist_z * sp.exp(
            - 2 * self.x ** 2 / self.waist_x0 ** 2
            - 2 * (self.z * sp.cos(self.theta) + self.y * sp.sin(self.theta)) ** 2
            / waist_z ** 2
        )  # fmt: on

    def get_lower_beam_intensity(self) -> sp.Expr:
        """Intensity distribution of the upper beam"""
        length = self.y * sp.cos(self.theta) + self.z * sp.sin(self.theta)
        waist_z = self.get_beam_waist_z(length)
        # Intensity of the beam at the origin
        i_0 = 2 * self.power / (sp.pi * self.waist_x0 * self.waist_z0)
        # Turn off automatic formatting to preserve readability
        # fmt: off
        return i_0 * self.waist_z0 / waist_z * sp.exp(
            - 2 * self.x ** 2 / self.waist_x0 ** 2
            - 2 * (self.z * sp.cos(self.theta) - self.y * sp.sin(self.theta)) ** 2
            / waist_z ** 2
        )  # fmt: on

    def get_total_intensity(
        self, with_tweezer: bool = True, with_2dtrap: bool = True
    ) -> sp.Expr:
        """Total intensity distribution of both beams"""
        i_1 = self.get_upper_beam_intensity()
        i_2 = self.get_lower_beam_intensity()
        i_tweezer = int(with_tweezer) * self.get_tweezer_intensity()
        i_2dtrap = int(with_2dtrap) * (
            i_1 + i_2 + 2 * sp.sqrt(i_1 * i_2) * sp.cos(2 * sp.pi * self.z / self.d)
        )
        return i_2dtrap + i_tweezer

    def get_potential(
        self,
        with_tweezer: bool = True,
        with_2dtrap: bool = True,
        with_grad: bool = True,
        apply_zero_offset: bool = True,
    ) -> sp.Expr:
        """Get the potential of the 2d trap.
        The total potential consists of the laser induced optical dipole
        potential and the magnetic field potential.
        """
        grad_r_coordinate = self.x * sp.cos(self.beta) - self.y * sp.sin(self.beta)
        pot = -self.a * self.get_total_intensity(
            with_tweezer=with_tweezer, with_2dtrap=with_2dtrap
        ) - int(with_grad) * self.mu_b * (
            self.grad_r * grad_r_coordinate + self.grad_z * self.z
        ) -  self.m * self.g * self.z
        offset = int(apply_zero_offset) * pot.subs({self.x: 0, self.y: 0, self.z: 0})
        return pot - offset

    @cache
    def get_omega(
        self,
        omega: sp.Symbol,
        with_tweezer: bool = True,
        with_2dtrap: bool = True,
    ) -> sp.Expr:
        """Calculate trap frequency for one spacial direction.

        Uses :func:`sympy.series`, equates the 2nd order term with
        the model of a harmonic oscillator and solves for omega.

        :param omega: Which spacial trap frequency to calculate.
            One of :attr:`.omega_x`, :attr:`.omega_y`,
            or :attr:`.omega_z`.
        :param force: Force re-evaluation of the trap frequency. The
            analytic solution contains many symbols and is thus rather
            lengthy. To save time, the result is cached in the
            :attr:`._omega_solution` attribute.
        :param with_tweezer: Whether or not to include tweezer potential
            in the calculation.
        :param with_2dtrap: Whether or not to include 2D trap potential
            in the calculation.
        :return: Trap frequency as a sympy expression.
        """
        others: set[sp.Symbol] = {self.x, self.y, self.z}
        variable: sp.Symbol
        if omega == self.omega_x or omega == self.omega_r_tweezer:
            variable = self.x
        elif omega == self.omega_y:
            variable = self.y
        elif omega == self.omega_z or omega == self.omega_ax_tweezer:
            variable = self.z
        else:
            raise ValueError(f"Unknown omega '{omega!r}'")
        others.remove(variable)
        series = self.subs(
            self.get_total_intensity(
                with_tweezer=with_tweezer, with_2dtrap=with_2dtrap
            ).subs({symbol: 0 for symbol in others}),
            # Only substitute symbols with other symbols
            symbolic_only=True,
        ).series(x=variable, x0=0, n=3)
        quad_term = sp.LT(series.removeO(), gens=(variable,))
        equation = sp.Eq(-self.a * quad_term, self.m * omega**2 * variable**2 / 2)
        solution = sp.solve(equation, omega)
        if len(solution) == 0:
            raise RuntimeError(f"Unable to solve for '{omega}'")
        return solution[0]

    def get_omega_x(
        self, with_tweezer: bool = True, with_2dtrap: bool = True
    ) -> sp.Expr:
        """Calculate trap frequency in x.

        See :meth:`.get_omega`.
        """
        return self.get_omega(
            self.omega_x, with_tweezer=with_tweezer, with_2dtrap=with_2dtrap
        )

    def get_omega_y(
        self, with_tweezer: bool = True, with_2dtrap: bool = True
    ) -> sp.Expr:
        """Calculate trap frequency in y.

        See :meth:`.get_omega`.
        """
        return self.get_omega(
            self.omega_y, with_tweezer=with_tweezer, with_2dtrap=with_2dtrap
        )

    def get_omega_z(
        self, with_tweezer: bool = True, with_2dtrap: bool = True
    ) -> sp.Expr:
        """Calculate trap frequency in z.

        See :meth:`.get_omega`.
        """
        return self.get_omega(
            self.omega_z, with_tweezer=with_tweezer, with_2dtrap=with_2dtrap
        )

    def waist_z_to_waist_y(self) -> sp.Expr:
        """Calculate waist in y from :math:`W_z`"""
        return self.waist_z0 / sp.sin(self.theta)

    def waist_z_to_waist_ax(self) -> sp.Expr:
        """Calculates the axial waist from :math:`W_z`"""
        return self.waist_z0 / sp.cos(self.theta)

    def nstationary_solution(
        self,
        dims: tuple[sp.Symbol, ...],
        extend: tuple[tuple[ExprNum, ExprNum], ...] | tuple[ExprNum, ExprNum],
        size: tuple[int, ...],
        k: int,
        method="sparse",
        export=False
    ):
        """Numeric stationary solution"""
        if isinstance(dims, sp.Expr):
            dims = (dims,)
        ndims = len(dims)
        if not (np.shape(extend) != (2,) or np.shape(extend) != (ndims, 2)):
            raise ValueError(
                f"Shape of parameter 'extend' has to be either (2,) or ({ndims}, 2), "
                f"but is {np.shape(extend)}."
            )
        if not (isinstance(size, int) or len(size) == ndims):
            raise ValueError(
                "Parameter 'size' has to be either of type 'int' or "
                f"a tuple of length {ndims}."
            )

        # Make sure 'size' is always a tuple, even if all dimensions
        # have the same size.
        if isinstance(size, int):
            size = (size,) * ndims

        others: set[sp.Symbol] = {self.x, self.y, self.z}
        coordinates: dict[sp.Symbol, npt.NDArray] = {}
        for i, dim in enumerate(dims):
            if dim not in others:
                raise ValueError(f"Symbols '{dim}' unsupported or used multiple times")
            others.remove(dim)
            # Check if a single extend is used for all dimensions. This
            # boolean is later used to determine if the extend should
            # be overwritten. This avoids re-evaluating the same extend
            # multiple times.
            single_extend = np.shape(extend) == (2,)
            # The extend is casted to a list to make the items mutable
            dim_extend: list[ExprNum, ExprNum] = (
                list(extend) if single_extend else list(extend[i])
            )
            # Make sure that all limits in the extend are numbers that
            # numpy can work with.
            for j, lim in enumerate(dim_extend):
                if isinstance(lim, sp.Expr):
                    new_lim = self.subs(lim).evalf()
                    if not isinstance(new_lim, (float, sp.Float)):
                        raise ValueError(
                            f"Unable to convert '{lim}' to float (using substitutions)"
                        )
                    dim_extend[j] = float(new_lim)
            if single_extend:
                # To avoid re-evaluation of the symbolic extends, fix
                # the extends to the calculated value
                extend = tuple(dim_extend)  # cast list to tuple again
            coordinates[dim] = np.linspace(*dim_extend, num=size[i])

        # Calculate size of the volume elements
        """
        coords = np.array(list(coordinates.values()))
        dvol: npt.NDArray = np.mean(coords[:, 1:] - coords[:, :-1], axis=1)
        assert dvol.shape == (ndims,)"""
        dvol = np.array([])
        for coord in coordinates.values():
            dvol = np.append(dvol,np.mean(np.diff(coord)))

        # Create potential array
        pot = self.subs(self.get_potential(apply_zero_offset=False).subs({k: 0 for k in others}))
        pot_numpy = sp.lambdify(tuple(coordinates.keys()), pot, cse=True)
        if ndims > 1:
            meshgrid = np.meshgrid(*coordinates.values(),indexing="ij")
            potential = pot_numpy(*meshgrid)
        else:
            potential = pot_numpy(*coordinates.values())

        # Perform the diagonalization of the Hamiltonian
        energies, states = nstationary_solution(
            tuple(dvol), potential, float(self.subs(self.m).evalf()), k=k,
            method=method)
            
        states = states.reshape((k, *size))

        #export hamiltonian and parameters
        if export:
            #check if we are trying to solve DoubleTweezer and save corr. parameters
            current_time = datetime.datetime.now()
            
            #add to overview the key parameters of this run
            if isinstance(self,DoubleTweezer):
                distance_tweezers = float(self.subs(self.distance_tweezers).evalf())
                power1 = float(self.subs(self.power_tweezer1).evalf())
                power2 = float(self.subs(self.power_tweezer2).evalf())
                waist1 = float(self.subs(self.waist_tweezer1).evalf())
                waist2 = float(self.subs(self.waist_tweezer2).evalf())
                wvl = float(self.subs(self.wvl).evalf())
                omega_x1, omega_x2 = self.get_both_omega(self.x)
                omega_x1 = float(self.subs(omega_x1).evalf())
                omega_x2 = float(self.subs(omega_x2).evalf())
                with open('data/overview.txt', 'a') as f:
                    f.write(f'{current_time.strftime("%Y-%m-%d_%H-%M-%S")}; {size}; {wvl}; {power1}; {power2}; {waist1}; {waist2}; {distance_tweezers}; {omega_x1/2/np.pi}; {omega_x2/2/np.pi} \n')

            elif isinstance(self,TwoSiteLattice):
                lattice_spacing_x = float(self.subs(self.lattice_spacing_x).evalf())
                lattice_spacing_y = float(self.subs(self.lattice_spacing_y).evalf())
                lattice_spacing_z = float(self.subs(self.lattice_spacing_z).evalf())
                omega_x1, omega_x2 = self.get_both_omega(self.x)
                omega_z1, omega_z2 = self.get_both_omega(self.z)
                aspect_ratio = float(self.subs(sp.sqrt(omega_x1/omega_z1)))
                wvl = float(self.subs(self.wvl).evalf())
                with open('data/overview.txt', 'a') as f:
                    f.write(f'{current_time.strftime("%Y-%m-%d_%H-%M-%S")}; {size}; {wvl}; {lattice_spacing_x}; {lattice_spacing_y}; {lattice_spacing_z}; {aspect_ratio}; TwoSiteLattice\n')
                
            #save results as npz
            x3D,y3D,z3D = np.meshgrid(coordinates[self.x],coordinates[self.y],coordinates[self.z],indexing="ij")
            np.savez(f"data/{current_time.strftime("%Y-%m-%d_%H-%M-%S")}_results.npz",
                            energies=energies, states=states, dvol=dvol,
                            k=k, size=size, extend=extend, pot=pot_numpy(x3D,y3D,z3D),
                            x=coordinates[self.x],y=coordinates[self.y],z=coordinates[self.z])               
            #save hamiltonian as sparse matrix
            save_npz(f"data/{current_time.strftime("%Y-%m-%d_%H-%M-%S")}_hamiltonian.npz",
                                    discrete_hamiltonian(tuple(dvol), potential, float(self.subs(self.m).evalf())))
            #save trap object
            with open(f"data/{current_time.strftime("%Y-%m-%d_%H-%M-%S")}_trap.npz", 'wb') as file:
                pickle.dump(self, file)
                
            print(f"files saved with ...._{current_time.strftime("%Y-%m-%d_%H-%M-%S")}")

        return energies, states, pot_numpy, coordinates
    
    def nstationary_solution_export(
        self,path,
        dims: tuple[sp.Symbol, ...],
        extend: tuple[tuple[ExprNum, ExprNum], ...] | tuple[ExprNum, ExprNum],
        size: tuple[int, ...],
        k: int,
    ):
        """Numeric stationary solution"""
        if isinstance(dims, sp.Expr):
            dims = (dims,)
        ndims = len(dims)
        if not (np.shape(extend) != (2,) or np.shape(extend) != (ndims, 2)):
            raise ValueError(
                f"Shape of parameter 'extend' has to be either (2,) or ({ndims}, 2), "
                f"but is {np.shape(extend)}."
            )
        if not (isinstance(size, int) or len(size) == ndims):
            raise ValueError(
                "Parameter 'size' has to be either of type 'int' or "
                f"a tuple of length {ndims}."
            )

        # Make sure 'size' is always a tuple, even if all dimensions
        # have the same size.
        if isinstance(size, int):
            size = (size,) * ndims

        others: set[sp.Symbol] = {self.x, self.y, self.z}
        coordinates: dict[sp.Symbol, npt.NDArray] = {}
        for i, dim in enumerate(dims):
            if dim not in others:
                raise ValueError(f"Symbols '{dim}' unsupported or used multiple times")
            others.remove(dim)
            # Check if a single extend is used for all dimensions. This
            # boolean is later used to determine if the extend should
            # be overwritten. This avoids re-evaluating the same extend
            # multiple times.
            single_extend = np.shape(extend) == (2,)
            # The extend is casted to a list to make the items mutable
            dim_extend: list[ExprNum, ExprNum] = (
                list(extend) if single_extend else list(extend[i])
            )
            # Make sure that all limits in the extend are numbers that
            # numpy can work with.
            for j, lim in enumerate(dim_extend):
                if isinstance(lim, sp.Expr):
                    new_lim = self.subs(lim).evalf()
                    if not isinstance(new_lim, (float, sp.Float)):
                        raise ValueError(
                            f"Unable to convert '{lim}' to float (using substitutions)"
                        )
                    dim_extend[j] = float(new_lim)
            if single_extend:
                # To avoid re-evaluation of the symbolic extends, fix
                # the extends to the calculated value
                extend = tuple(dim_extend)  # cast list to tuple again
            coordinates[dim] = np.linspace(*dim_extend, num=size[i])

        # Calculate size of the volume elements
        """
        coords = np.array(list(coordinates.values()))
        dvol: npt.NDArray = np.mean(coords[:, 1:] - coords[:, :-1], axis=1)
        assert dvol.shape == (ndims,)"""
        dvol = np.array([])
        for coord in coordinates.values():
            dvol = np.append(dvol,np.mean(np.diff(coord)))

        # Create potential array
        pot = self.subs(self.get_potential(apply_zero_offset=False).subs({k: 0 for k in others}))
        pot_numpy = sp.lambdify(tuple(coordinates.keys()), pot, cse=True)
        if ndims > 1:
            meshgrid = np.meshgrid(*coordinates.values(),indexing="ij")
            potential = pot_numpy(*meshgrid)
        else:
            potential = pot_numpy(*coordinates.values())

        # Perform the diagonalization of the Hamiltonian
        e_min = nstationary_solution_export(path,
            tuple(dvol), potential, float(self.subs(self.m).evalf()), k=k)
        
        x3D,y3D,z3D = np.meshgrid(coordinates[self.x],coordinates[self.y],coordinates[self.z],indexing="ij")

        #check if we are trying to solve DoubleTweezer and save corr. parameters
        if isinstance(self,DoubleTweezer):
            np.savez(path + r"\parameters.npz",e_min=e_min, k=k, size=size, pot=pot_numpy(x3D,y3D,z3D),
                                            x=coordinates[self.x],y=coordinates[self.y],z=coordinates[self.z],
                                            mass=self.m, mu=self.mu_b,distance_tweezers=self.distance_tweezers,
                                            power1=self.power_tweezer1, power2=self.power_tweezer2,
                                            waist1=self.waist_tweezer1, waist2=self.waist_tweezer2)
        else:
            np.savez(path + r"\parameters.npz",e_min=e_min, k=k, size=size, pot=pot_numpy(x3D,y3D,z3D),
                                            x=coordinates[self.x],y=coordinates[self.y],z=coordinates[self.z],
                                            mass=self.m, mu=self.mu_b)
        print("files saved under " + path)


class TrueHarmonicTweezer(PancakeTrap):
    """Pancake Trap but with a truly harmonic tweezer (i.e. quadratic
    in z and r). Note that spilling is not possible, as the gradient
    simply shifts the center of the trap.
    """

    def get_tweezer_intensity(self) -> sp.Expr:
        i_xz = super().get_tweezer_intensity().subs({self.y: 0})
        series_x = i_xz.subs({self.z: 0}).series(x=self.x, x0=0, n=3).removeO()
        series_z = i_xz.subs({self.x: 0}).series(x=self.z, x0=0, n=3).removeO()
        series_r = series_x.subs(self.x, sp.sqrt(self.x**2 + self.y**2))
        quad_term_z = sp.LT(series_z, gens=(self.z,))
        return series_r + quad_term_z


class DoubleTweezer(PancakeTrap):
    """
    Pancake trap but with two tweezers.
    Each waist and power, and their spacing can be adjusted.
    """
    def __init__(self,**values):
        self.power_tweezer1 = sp.Symbol("P_t1", real=True, positive=True, finite=True, nonzero=True)
        self.power_tweezer2 = sp.Symbol("P_t2", real=True, positive=True, finite=True, nonzero=True)
        self.waist_tweezer1 = sp.Symbol("W_t1", real=True, positive=True, finite=True, nonzero=True)
        self.waist_tweezer2 = sp.Symbol("W_t2", real=True, positive=True, finite=True, nonzero=True)
        self.distance_tweezers = sp.Symbol("d_t", real=True, positive=True, finite=True, nonzero=True)
        super().__init__(**values)
    
    def get_tweezer_rayleigh1(self):
        return sp.pi * self.waist_tweezer1**2 / self.wvl
    
    def get_tweezer_rayleigh2(self):
        return sp.pi * self.waist_tweezer2**2 / self.wvl

    def get_tweezer_intensity1(self,x_offset) -> sp.Expr:
        rayleigh_length = self.get_tweezer_rayleigh1()
        waist = self.waist_tweezer1 * sp.sqrt(1 + (self.z / rayleigh_length) ** 2)
        i_0 = 2 * self.power_tweezer1 / (sp.pi * self.waist_tweezer1**2)
        r = sp.sqrt((self.x-x_offset)**2 + self.y**2)
        return i_0 * (self.waist_tweezer1 / waist) ** 2 * sp.exp(-2 * r**2 / (waist**2))
    
    def get_tweezer_intensity2(self,x_offset) -> sp.Expr:
        rayleigh_length = self.get_tweezer_rayleigh2()
        waist = self.waist_tweezer2 * sp.sqrt(1 + (self.z / rayleigh_length) ** 2)
        i_0 = 2 * self.power_tweezer2 / (sp.pi * self.waist_tweezer2**2)
        r = sp.sqrt((self.x-x_offset)**2 + self.y**2)
        return i_0 * (self.waist_tweezer2 / waist) ** 2 * sp.exp(-2 * r**2 / (waist**2))

    def get_tweezer_intensity(self) -> sp.Expr:
        return self.get_tweezer_intensity1(-self.distance_tweezers/2) + self.get_tweezer_intensity2(self.distance_tweezers/2)
    
    def get_both_omega(self, dim: sp.Symbol):
        """
        Returns omega_1, omega_2 the frequency in direction dim (sympy coord)
        of the left(1) and right(2) tweezer
        """
        V = self.subs(self.get_potential(apply_zero_offset=False))
        a = float(self.subs(self.distance_tweezers))

        #find minima of potential
        def V_func(x):
            return float(V.subs(self.x,x).subs(self.y,0).subs(self.z,0))
        x_right = minimize_scalar(V_func,bracket=[0,a]).x
        x_left = minimize_scalar(V_func,bracket=[-a,0]).x

        #catch case where both potentials have already merged
        tunneling_dist = abs(x_right-x_left)
        if tunneling_dist < 1e-20:
            raise Exception("potential has only one minmum")
    
        #trapping frequencies through second derivative
        V_prime = sp.diff(V, dim)
        V_double_prime = sp.diff(V_prime, dim)

        omega_1 = sp.sqrt(V_double_prime.subs({self.x:x_left, self.y:0, self.z:0})/self.m)
        omega_2 = sp.sqrt(V_double_prime.subs({self.x:x_left, self.y:0, self.z:0})/self.m)

        return omega_1, omega_2

class TwoSiteLattice(PancakeTrap):
    """
    Trap with lattice potential that has just two sites
    """
    def __init__(self,**values):
        self.lattice_depth_x = sp.Symbol("V0_x", real=True, positive=True, finite=True, nonzero=True)
        self.lattice_depth_y = sp.Symbol("V0_y", real=True, positive=True, finite=True, nonzero=True)
        self.lattice_depth_z = sp.Symbol("V0_z", real=True, positive=True, finite=True, nonzero=True)
        #lattice spacings are normaly half the wavelength in that direction
        self.lattice_spacing_x = sp.Symbol("d_x", real=True, positive=True, finite=True, nonzero=True)
        self.lattice_spacing_y = sp.Symbol("d_y", real=True, positive=True, finite=True, nonzero=True)
        self.lattice_spacing_z = sp.Symbol("d_z", real=True, positive=True, finite=True, nonzero=True)
        super().__init__(**values)


    def get_potential(
        self,
        with_tweezer: bool = True,
        with_2dtrap: bool = True,
        with_lattice: bool = True,
        with_grad: bool = True,
        apply_zero_offset: bool = True,
    ) -> sp.Expr:
        """Get the potential of the 2d trap.
        The total potential consists of the laser induced optical dipole
        potential and the magnetic field potential.
        """
        a = self.lattice_spacing_x
        b = self.lattice_spacing_y
        c = self.lattice_spacing_z
        #use sine in x direction to have two wells around the origin
        pot_lattice = (self.lattice_depth_x* sp.sin(np.pi/a*self.x)**2 + 
                       self.lattice_depth_y* sp.cos(np.pi/b*self.y)**2 +
                       self.lattice_depth_z* sp.cos(np.pi/c*self.z)**2)

        grad_r_coordinate = self.x * sp.cos(self.beta) - self.y * sp.sin(self.beta)
        pot = -self.a * self.get_total_intensity(
            with_tweezer=with_tweezer, with_2dtrap=with_2dtrap
        ) - int(with_grad) * self.mu_b * (
            self.grad_r * grad_r_coordinate + self.grad_z * self.z
        ) - int(with_lattice)* (pot_lattice
        ) -  self.m * self.g * self.z
        offset = int(apply_zero_offset) * pot.subs({self.x: 0, self.y: 0, self.z: 0})
        return pot - offset
    
    def get_both_omega(self, dim: sp.Symbol):
        """
        Returns omega_1, omega_2 the frequency in direction dim (sympy coord)
        of the left(1) and right(2) tweezer
        """
        V = self.subs(self.get_potential(apply_zero_offset=False))
        a = float(self.subs(self.lattice_spacing_x))

        #find minima of potential
        x_right = a/2
        x_left = -a/2

        #catch case where both potentials have already merged
        tunneling_dist = abs(x_right-x_left)
        if tunneling_dist < 1e-20:
            raise Exception("potential has only one minmum")
    
        #trapping frequencies through second derivative
        V_prime = sp.diff(V, dim)
        V_double_prime = sp.diff(V_prime, dim)

        omega_1 = sp.sqrt(V_double_prime.subs({self.x:x_left, self.y:0, self.z:0})/self.m)
        omega_2 = sp.sqrt(V_double_prime.subs({self.x:x_left, self.y:0, self.z:0})/self.m)

        return omega_1, omega_2
    


def _optional_a(
    a: float_or_array | None, wavelength: float_or_array | None
) -> float_or_array:
    if a is None:
        if wavelength is None:
            raise ValueError("Either 'a' or 'wavelength' has to be specified.")
        a = get_a(wavelength)
    elif wavelength is not None:
        warnings.warn(
            "Both 'a' and 'wavelength' were provided. "
            "Keyword argument 'wavelength' will be ignored.",
            RuntimeWarning,
            stacklevel=2,
        )
    return a


def get_d(
    waist_r: float_or_array | None = None,
    waist_z: float_or_array | None = None,
    wavelength: float_or_array = p.wavelength_trap,
    theta: float_or_array | None = p.theta,  # in rad
) -> float_or_array:
    """Get fringe spacing (in z-direction) for the 2D trap pancakes.
    If the waists are specified, the angle theta is ignored as it can
    be substituted using the waists.
    Either both waists or the angle have to be specified.

    :param waist_r:
    :param waist_z:
    :param wavelength:
    :param theta:
    :return:
    """
    if waist_r is not None and waist_z is not None:
        return wavelength / (2 * waist_z) * waist_r
    elif theta is not None:
        return wavelength / (2 * np.cos(theta))
    else:
        raise ValueError(
            "Either 'waist_r' and 'waist_z', or 'theta' have to be specified."
        )


def get_a(wavelength: float_or_array) -> float_or_array:
    """Get real part of the polarizability (Re{alpha}) for a

    :param wavelength: Laser wavelength
    :return:
    """
    omega_laser = 2 * np.pi * const.c / wavelength
    return (
        (3 * np.pi * const.c**2)
        / (2 * p.omega_0**3)
        * (p.gamma / (p.omega_0 - omega_laser) + p.gamma / (p.omega_0 + omega_laser))
    )


def get_power_radial_gaussian(
    omega: float_or_array,
    waist: float_or_array,
    a: float_or_array | None = None,
    wavelength: float_or_array | None = None,
) -> float_or_array:
    """Calculate the required power for a radial gaussian beam with a
    given desired trap frequency and waist.

    :param omega: Desired trap frequency to compensate.
    :param waist: Waist of the repulsion beam.
    :param a: Pre-factor 'a' for the trap depth.
    :param wavelength: Optional to calculate 'a'.
    :return: Power of the repulsion beam.
    """
    a = _optional_a(a, wavelength)
    return np.abs(omega**2 * np.pi * waist**4 * p.mass_lithium6 / (8 * a))


def get_omega_r_trap(
    power: float_or_array,
    waist_r: float_or_array,
    waist_z: float_or_array,
    a: float_or_array | None = None,
    wavelength: float_or_array | None = None,
) -> float_or_array:
    """
    :param power: Power of the trapping beam.
    :param waist_r: Radial waist, denoted :math:`W_{0x}` in Ram-Janik's
        thesis. Optional.
    :param waist_z: Waist in z-direction, denoted :math:`W_{0z}`.
        Optional.
    :param a: Pre-factor 'a' for the trap depth.
    :param wavelength: Optional to calculate 'a'.
    :return: Trapping frequency trap :math:`\\omega_{x}`.
    """
    a = _optional_a(a, wavelength)
    return np.sqrt(32 * a * power / (np.pi * p.mass_lithium6 * waist_r**3 * waist_z))


def get_omega_z_trap(
    power_trap: float_or_array,
    waist_r: float_or_array,
    waist_z: float_or_array,
    theta: float_or_array = p.theta,  # in rad
    wavelength: float_or_array = p.wavelength_trap,
) -> float_or_array:
    a = get_a(wavelength)
    return np.sqrt(
        16
        * a
        * power_trap
        / (np.pi * p.mass_lithium6 * waist_r * waist_z)
        * (
            2 * np.sin(theta) ** 2 / waist_z**2
            + (wavelength * np.sin(theta) / np.pi) ** 2
            * (1 / waist_r**4 + 1 / waist_z**4)
            + (np.pi / get_d(waist_r, waist_z, wavelength)) ** 2
        )
    )


def get_power_omega_r(
    omega_r: float_or_array,
    waist_r: float_or_array = p.waist_trap_r,
    waist_z: float_or_array = p.waist_trap_z,
    a: float_or_array | None = None,
    wavelength: float_or_array | None = None,
) -> float_or_array:
    """Calculate the power of the trap based upon the radial properties.

    :param omega_r: Radial trapping frequency.
    :param waist_r: Radial waist, denoted :math:`W_{0x}` in Ram-Janik's
        thesis. Optional.
    :param waist_z: Waist in z-direction, denoted :math:`W_{0z}`.
        Optional.
    :param a: Pre-factor 'a' for the trap depth.
    :param wavelength: Optional to calculate 'a'.
    :return: Power of the trapping beam :math:`P_{1,2}`.
    """
    intensity = pot.TotalIntensity()
    series_x = intensity.subs(pot.z, 0).subs(pot.y, 0).series(x=pot.x, x0=0, n=3)
    a = _optional_a(a, wavelength)
    return omega_r**2 * np.pi * waist_r**3 * waist_z * p.mass_lithium6 / (32 * a)


def get_power_omega_z(
    omega_z: float_or_array,
    waist_r: float_or_array = p.waist_trap_r,
    waist_z: float_or_array = p.waist_trap_z,
    a: float_or_array | None = None,
    wavelength: float_or_array | None = None,
) -> float_or_array:
    """Calculate the power of the trap based upon the axial properties.

    :param omega_z: Axial trapping frequency.
    :param waist_r: Radial waist, denoted :math:`W_{0x}` in Ram-Janik's
        thesis. Optional.
    :param waist_z: Waist in z-direction, denoted :math:`W_{0z}`.
        Optional.
    :param a: Pre-factor 'a' for the trap depth.
    :param wavelength: Optional to calculate 'a'.
    :return: Power of the trapping beam :math:`P_{1,2}`.
    """
    a = _optional_a(a, wavelength)
    return


def gaussian_beam_1d(
    r: float_or_array, power: float_or_array, waist: float_or_array
) -> float_or_array:
    """Gaussian beam profile in 1d.

    :param r: Radial coordinate
    :param power: Power of the beam
    :param waist: Waist
    """
    return 2 * power / (np.pi * waist**2) * np.exp(-2 * r**2 / waist**2)
