from collections.abc import Iterable
from typing import Any, Literal, Optional, Union

import matplotlib.pyplot as plt
import numpy as np
from qutip import Bloch, Qobj, wigner
from scipy.linalg import cosm, eigh, sinm

from .operators import (
    creation,
    destroy,
    ptrace,
    sigma_x,
    sigma_y,
    sigma_z,
    state_to_density_matrix,
)
from .qubit_base import QubitBase


class Ferbo(QubitBase):
    PARAM_LABELS = {
        "Ec": r"$E_C$",
        "El": r"$E_L$",
        "Ej": r"$E_J$",
        "Gamma": r"$\Gamma$",
        "delta_Gamma": r"$\delta \Gamma$",
        "er": r"$\epsilon_r$",
        "phase": r"$\Phi_{\mathrm{ext}} / \Phi_0$",
    }

    OPERATOR_LABELS = {
        "n_operator": r"\hat{n}",
        "phase_operator": r"\hat{\phi}",
        "d_hamiltonian_d_ng": r"\partial \hat{H} / \partial n_g",
        "d_hamiltonian_d_phase": r"\partial \hat{H} / \partial \phi_{{ext}}",
        "d_hamiltonian_d_EL": r"\partial \hat{H} / \partial E_L",
        "d_hamiltonian_d_deltaGamma": r"\partial \hat{H} / \partial \delta \Gamma",
        "d_hamiltonian_d_er": r"\partial \hat{H} / \partial \epsilon_r",
    }

    def __init__(
        self,
        Ec,
        El,
        Ej,
        Gamma,
        delta_Gamma,
        er,
        phase,
        dimension,
        flux_grouping: Literal["EL", "ABS"] = "ABS",
        Delta=40,
    ):
        """
        Initializes the Ferbo class with the given parameters.

        Parameters
        ----------
        Ec : float
            Charging energy.
        El : float
            Inductive energy.
        Ej : float
            Josephson energy.
        Gamma : float
            Coupling strength.
        delta_Gamma : float
            Coupling strength difference.
        er : float
            Energy relaxation rate.
        phase : float
            External magnetic phase.
        dimension : int
            Dimension of the Hilbert space.
        flux_grouping : Literal["EL", "ABS"], optional
            Flux grouping ('EL' or 'ABS') (default is 'ABS').
        Delta : float
            Superconducting gap.
        """
        if flux_grouping not in ["EL", "ABS"]:
            raise ValueError("Invalid flux grouping; must be 'EL' or 'ABS'.")

        self.Ec = Ec
        self.El = El
        self.Ej = Ej
        self.Gamma = Gamma
        self.delta_Gamma = delta_Gamma
        self.er = er
        self.phase = phase
        self.dimension = dimension // 2 * 2
        self.flux_grouping = flux_grouping
        self.Delta = Delta
        super().__init__(self.dimension)

    @property
    def phase_zpf(self) -> float:
        """
        Returns the zero-point fluctuation of the phase.

        Returns
        -------
        float
            Zero-point fluctuation of the phase.
        """
        return (2 * self.Ec / self.El) ** 0.25

    @property
    def n_zpf(self) -> float:
        """
        Returns the zero-point fluctuation of the charge number.

        Returns
        -------
        float
            Zero-point fluctuation of the charge number.
        """
        return 1 / 2 * (self.El / 2 / self.Ec) ** 0.25

    @property
    def lc_energy(self) -> float:
        """
        Returns the plasma energy.

        Returns
        -------
        float
            Plasma energy.
        """
        return np.sqrt(8 * self.Ec * self.El)

    @property
    def transparency(self) -> float:
        """
        Return the transparency of the weak link.

        Returns
        -------
        float
            Transparency of the weak link.
        """

        return (self.Gamma**2 - self.delta_Gamma**2) / (self.Gamma**2 + self.er**2)

    def phi_osc(self) -> float:
        """
        Returns the oscillator length for the LC oscillator composed of the inductance and capacitance.

        Returns
        -------
        float
            Oscillator length.
        """
        return (8.0 * self.Ec / self.El) ** 0.25

    def n_operator(self) -> np.ndarray:
        """
        Returns the charge number operator.

        Returns
        -------
        np.ndarray
            The charge number operator.
        """
        single_mode_n_operator = (
            1j
            * self.n_zpf
            * (creation(self.dimension // 2) - destroy(self.dimension // 2))
        )
        return np.kron(np.eye(2), single_mode_n_operator)

    def phase_operator(self) -> np.ndarray:
        """
        Returns the total phase operator.

        Returns
        -------
        np.ndarray
            The total phase operator.
        """
        single_mode_phase_operator = self.phase_zpf * (
            creation(self.dimension // 2) + destroy(self.dimension // 2)
        )
        return np.kron(np.eye(2), single_mode_phase_operator)

    def jrl_potential(self) -> np.ndarray:
        """
        Returns the Josephson Resonance Level potential.

        Returns
        -------
        np.ndarray
            The Josephson Resonance Level potential.
        """
        phase_op = self.phase_operator()[: self.dimension // 2, : self.dimension // 2]
        if self.flux_grouping == "ABS":
            phase_op -= self.phase * np.eye(self.dimension // 2)

        return (
            -self.Gamma * np.kron(sigma_z(), cosm(phase_op / 2))
            - self.delta_Gamma * np.kron(sigma_y(), sinm(phase_op / 2))
            + self.er * np.kron(sigma_x(), np.eye(self.dimension // 2))
        )

    def hamiltonian(self) -> np.ndarray:
        """
        Returns the Hamiltonian of the system.

        Returns
        -------
        np.ndarray
            The Hamiltonian of the system.
        """
        # n_x = self.delta_Gamma / 4 / (self.Gamma + self.Delta)
        n_op = (
            self.n_operator()
        )  # + n_x * np.kron(sigma_x(), np.eye(self.dimension // 2))

        charge_term = 4 * self.Ec * n_op @ n_op
        phase_op = self.phase_operator()

        if self.flux_grouping == "ABS":
            inductive_term = 0.5 * self.El * phase_op @ phase_op
            phase_op -= self.phase * np.eye(self.dimension)
            josephson_term = -self.Ej * cosm(phase_op)
        else:
            josephson_term = -self.Ej * cosm(phase_op)
            phase_op += self.phase * np.eye(self.dimension)
            inductive_term = 0.5 * self.El * phase_op @ phase_op

        potential = self.jrl_potential()
        return charge_term + inductive_term + potential + josephson_term

    def d_hamiltonian_d_EC(self) -> np.ndarray:
        """
        Returns the derivative of the Hamiltonian with respect to the charging energy.

        Returns
        -------
        np.ndarray
            The derivative of the Hamiltonian with respect to the charging energy.
        """
        n_x = self.delta_Gamma / 4 / (self.Gamma + self.Delta)
        n_op = self.n_operator() + n_x * np.kron(sigma_x(), np.eye(self.dimension // 2))

        return 8 * n_op @ n_op

    def d_hamiltonian_d_EL(self) -> np.ndarray:
        """
        Returns the derivative of the Hamiltonian with respect to the inductive energy.

        Returns
        -------
        np.ndarray
            The derivative of the Hamiltonian with respect to the inductive energy.
        """
        if self.flux_grouping == "EL":
            phase_op = self.phase_operator()
        elif self.flux_grouping == "ABS":
            phase_op = self.phase_operator() - self.phase * np.eye(self.dimension)

        return 1 / 2 * np.dot(phase_op, phase_op)

    def d_hamiltonian_d_EJ(self) -> np.ndarray:
        """
        Returns the derivative of the Hamiltonian with respect to the Josephson energy.

        Returns
        -------
        np.ndarray
            The derivative of the Hamiltonian with respect to the Josephson energy.
        """
        phase_op = self.phase_operator()
        if self.flux_grouping == "ABS":
            phase_op -= self.phase * np.eye(self.dimension)

        return -cosm(phase_op)

    def d_hamiltonian_d_Gamma(self) -> np.ndarray:
        """
        Returns the derivative of the Hamiltonian with respect to Gamma.

        Returns
        -------
        np.ndarray
            The derivative of the Hamiltonian with respect to Gamma.
        """
        phase_op = self.phase_operator()[: self.dimension // 2, : self.dimension // 2]
        if self.flux_grouping == "ABS":
            phase_op -= self.phase * np.eye(self.dimension // 2)

        return -np.kron(sigma_z(), sinm(phase_op / 2))

    def d_hamiltonian_d_er(self) -> np.ndarray:
        """
        Returns the derivative of the Hamiltonian with respect to the energy relaxation rate.

        Returns
        -------
        np.ndarray
            The derivative of the Hamiltonian with respect to the energy relaxation rate.
        """
        return +np.kron(sigma_x(), np.eye(self.dimension // 2))

    def d_hamiltonian_d_deltaGamma(self) -> np.ndarray:
        """
        Returns the derivative of the Hamiltonian with respect to the coupling strength difference.

        Returns
        -------
        np.ndarray
            The derivative of the Hamiltonian with respect to the coupling strength difference.
        """
        phase_op = self.phase_operator()[: self.dimension // 2, : self.dimension // 2]
        if self.flux_grouping == "ABS":
            phase_op -= self.phase * np.eye(self.dimension // 2)

        return -np.kron(sigma_y(), sinm(phase_op / 2))

    def d2_hamiltonian_d_Gamma2(self) -> np.ndarray:
        """
        Returns the second derivative of the Hamiltonian with respect to Gamma.

        Returns
        -------
        np.ndarray
            The second derivative of the Hamiltonian with respect to Gamma.
        """
        return np.zeros((self.dimension, self.dimension))

    def d2_hamiltonian_d_deltaGamma2(self) -> np.ndarray:
        """
        Returns the second derivative of the Hamiltonian with respect to the coupling strength difference.

        Returns
        -------
        np.ndarray
            The second derivative of the Hamiltonian with respect to the coupling strength difference.
        """
        return np.zeros((self.dimension, self.dimension))

    def d2_hamiltonian_d_er2(self) -> np.ndarray:
        """
        Returns the second derivative of the Hamiltonian with respect to the energy relaxation rate.

        Returns
        -------
        np.ndarray
            The second derivative of the Hamiltonian with respect to the energy relaxation rate.
        """
        return np.zeros((self.dimension, self.dimension))

    def d_hamiltonian_d_ng(self) -> np.ndarray:
        """
        Returns the derivative of the Hamiltonian with respect to the number of charge offset.

        Returns
        -------
        np.ndarray
            The derivative of the Hamiltonian with respect to the number of charge offset.

        """
        n_x = self.delta_Gamma / 4 / (self.Gamma + self.Delta)
        return (
            -8
            * self.Ec
            * (
                self.n_operator()
                + n_x * np.kron(sigma_x(), np.eye(self.dimension // 2))
            )
        )

    def d2_hamiltonian_d_ng2(self) -> np.ndarray:
        """
        Returns the second derivative of the Hamiltonian with respect to the number of charge offset.

        Returns
        -------
        np.ndarray
            The second derivative of the Hamiltonian with respect to the number of charge offset.

        """
        return 8 * self.Ec * np.eye(self.dimension)

    def d_hamiltonian_d_phase(self) -> np.ndarray:
        """
        Returns the derivative of the Hamiltonian with respect to the external magnetic phase.

        Returns
        -------
        np.ndarray
            The derivative of the Hamiltonian with respect to the external magnetic phase.
        """
        if self.flux_grouping == "EL":
            return self.El * (
                self.phase_operator() + self.phase * np.eye(self.dimension)
            )
        elif self.flux_grouping == "ABS":
            phase_op = self.phase_operator()[
                : self.dimension // 2, : self.dimension // 2
            ] - self.phase * np.eye(self.dimension // 2)
            return -self.Gamma / 2 * np.kron(
                sigma_z(), sinm(phase_op / 2)
            ) + self.delta_Gamma / 2 * np.kron(sigma_y(), cosm(phase_op / 2))
        else:
            raise ValueError(f"Unknown flux_grouping: {self.flux_grouping}")

    def d2_hamiltonian_d_phase2(self) -> np.ndarray:
        """
        Returns the second derivative of the Hamiltonian with respect to the external magnetic phase.

        Returns
        -------
        np.ndarray
            The second derivative of the Hamiltonian with respect to the external magnetic phase.
        """
        if self.flux_grouping == "EL":
            return self.El * np.eye(self.dimension)
        elif self.flux_grouping == "ABS":
            phase_op = self.phase_operator()[
                : self.dimension // 2, : self.dimension // 2
            ] - self.phase * np.eye(self.dimension // 2)
            return self.Gamma / 4 * np.kron(
                sigma_z(), cosm(phase_op / 2)
            ) + self.delta_Gamma / 4 * np.kron(sigma_y(), sinm(phase_op / 2))
        else:
            raise ValueError(f"Unknown flux_grouping: {self.flux_grouping}")

    def wigner(
        self,
        which: int = 0,
        phi_grid: Optional[np.ndarray] = None,
        n_grid: Optional[np.ndarray] = None,
        esys: Optional[tuple[np.ndarray, np.ndarray]] = None,
    ):
        """
        Computes the Wigner function for a given wavefunction.

        Parameters
        ----------
        which : int, optional
            Index of desired wavefunction (default is 0).
        phi_grid : np.ndarray, optional
            Custom grid for phi; if None, a default grid is used.
        n_grid : np.ndarray, optional
            Custom grid for n; if None, a default grid is used.
        esys : Tuple[np.ndarray, np.ndarray], optional
            Precomputed eigenvalues and eigenvectors.

        Returns
        -------
        np.ndarray
            The Wigner function.
        """
        rho_reduced = self.reduced_density_matrix(which=which, esys=esys, subsys=0)

        if phi_grid is None:
            phi_grid = np.linspace(-5, 5, 151)
        if n_grid is None:
            n_grid = np.linspace(-5, 5, 151)

        rho_reduced_qobj = Qobj(rho_reduced)

        wigner_func = wigner(rho_reduced_qobj, phi_grid, n_grid)
        return wigner_func

    def reduced_density_matrix(
        self,
        which: int = 0,
        esys: tuple[np.ndarray, np.ndarray] = None,
        subsys: int = 0,
    ) -> np.ndarray:
        """
        Computes the reduced density matrix for a given wavefunction.

        Parameters
        ----------
        which : int, optional
            Index of desired wavefunction (default is 0).
        esys : Tuple[np.ndarray, np.ndarray], optional
            Precomputed eigenvalues and eigenvectors.
        subsys : int, optional
            Subsystem to compute the reduced density matrix for (default is 0).
            0 for the tracing out the Fock states, 1 for the Andreev states.

        Returns
        -------
        np.ndarray
            The reduced density matrix.
        """

        if esys is None:
            evals_count = max(which + 1, 3)
            _, evecs = self.eigensys(evals_count)
        else:
            _, evecs = esys

        rho = state_to_density_matrix(evecs[:, which])
        rho_reduced = ptrace(rho, dims=(2, self.dimension // 2), subsys=subsys)

        return rho_reduced

    def wavefunction(
        self,
        which: int = 0,
        phi_grid: Optional[np.ndarray] = None,
        esys: Optional[tuple[np.ndarray, np.ndarray]] = None,
        basis: Literal["phase", "charge"] = "phase",
        andreev_basis: Literal["static", "adiabatic"] = "static",
    ) -> dict[str, Any]:
        """
        Returns a wave function in the phi basis.

        Parameters
        ----------
        which : int, optional
            Index of desired wave function (default is 0).
        phi_grid : np.ndarray, optional
            Custom grid for phi; if None, a default grid is used.
        basis : Literal["phase", "charge"], optional
            Basis in which to return the wavefunction ('phase' or 'charge') (default is 'phase').
        andreev_basis : Literal["static", "adiabatic"], optional
            Andreev basis used for the two-component wavefunction. Use 'static'
            for the basis used to define the Hamiltonian matrix, or 'adiabatic'
            for the local phi-dependent eigenbasis of the Andreev sector.
        Returns
        -------
        Dict[str, Any]
            Wave function data containing basis labels, amplitudes, and energy.
        """
        if andreev_basis not in {"static", "adiabatic"}:
            raise ValueError("Invalid andreev_basis; must be 'static' or 'adiabatic'.")

        if esys is None:
            evals_count = max(which + 1, 3)
            evals, evecs = self.eigensys(evals_count)
        else:
            evals, evecs = esys

        dim = self.dimension // 2
        evecs = evecs.T

        if basis == "phase":
            l_osc = self.phase_zpf
        elif basis == "charge":
            l_osc = self.n_zpf

        if phi_grid is None:
            phi_grid = np.linspace(-5 * np.pi, 5 * np.pi, 151)

        phi_basis_labels = phi_grid
        wavefunc_osc_basis_amplitudes = evecs[which, :]
        phi_wavefunc_amplitudes = np.zeros((2, len(phi_grid)), dtype=np.complex128)

        # Compute the wavefunction in the static Andreev basis used in the Hamiltonian.
        for n in range(dim):
            phi_wavefunc_amplitudes[0] += wavefunc_osc_basis_amplitudes[
                n
            ] * self.harm_osc_wavefunction(n, phi_basis_labels, l_osc)
            phi_wavefunc_amplitudes[1] += wavefunc_osc_basis_amplitudes[
                self.dimension // 2 + n
            ] * self.harm_osc_wavefunction(n, phi_basis_labels, l_osc)

        # Optionally rotate from the static Andreev basis to the adiabatic
        # phi-dependent Andreev basis.
        if andreev_basis == "adiabatic":
            # Get rotation matrices for each phi point
            _, rotation_matrices = self.potential(phi_grid, return_evecs=True)

            # Rotate wavefunction at each phi point
            rotated_amplitudes = np.zeros_like(phi_wavefunc_amplitudes)
            for i in range(len(phi_grid)):
                # rotation_matrices[i] transforms from computational to Andreev eigenbasis
                # We need to apply the adjoint to transform the wavefunction
                state_vector = np.array(
                    [phi_wavefunc_amplitudes[0, i], phi_wavefunc_amplitudes[1, i]]
                )
                rotated_state = rotation_matrices[i].conj().T @ state_vector
                rotated_amplitudes[0, i] = rotated_state[0]
                rotated_amplitudes[1, i] = rotated_state[1]
            phi_wavefunc_amplitudes = rotated_amplitudes

        if basis == "charge":
            phi_wavefunc_amplitudes[0] /= np.sqrt(self.n_zpf)
            phi_wavefunc_amplitudes[1] /= np.sqrt(self.n_zpf)
        elif basis == "phase":
            phi_wavefunc_amplitudes[0] /= np.sqrt(self.phase_zpf)
            phi_wavefunc_amplitudes[1] /= np.sqrt(self.phase_zpf)

        # Fix global phase: ensure the wavefunction is mostly positive
        # For each Andreev component, choose phase so that the sum of real parts is positive
        for j in range(2):
            real_sum = np.sum(np.real(phi_wavefunc_amplitudes[j]))
            if real_sum < 0:
                phi_wavefunc_amplitudes[j] *= -1

        return {
            "basis_labels": phi_basis_labels,
            "amplitudes": phi_wavefunc_amplitudes,
            "energy": evals[which],
        }

    def potential(
        self,
        phi: Union[float, np.ndarray],
        return_evecs: bool = False,
        include_berry: bool = True,
    ) -> Union[np.ndarray, tuple[np.ndarray, np.ndarray]]:
        """
        Calculates the potential energy for given values of phi.

        Parameters
        ----------
        phi : Union[float, np.ndarray]
            The phase values at which to calculate the potential.
        return_evecs : bool, optional
            If True, returns both eigenvalues and eigenvectors (default is False).
        include_berry : bool, optional
            If True, includes the scalar Berry correction in the potential
            eigenvalues (default is True).

        Returns
        -------
        Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]
            If return_evecs is False: The potential energy values (eigenvalues).
            If return_evecs is True: A tuple of (eigenvalues, eigenvectors).
            Eigenvalues shape: (len(phi), 2)
            Eigenvectors shape: (len(phi), 2, 2) where each [i, :, :] is the rotation matrix at phi[i]
        """
        phi_array = np.atleast_1d(phi)
        evals_array = np.zeros((len(phi_array), 2))
        if return_evecs:
            evecs_array = np.zeros((len(phi_array), 2, 2), dtype=complex)

        for i, phi_val in enumerate(phi_array):
            # Construct only the Andreev part
            if self.flux_grouping == "ABS":
                andreev_phase = phi_val - self.phase
                andreev_term = (
                    -self.Gamma * np.cos(andreev_phase / 2) * sigma_z()
                    - self.delta_Gamma * np.sin(andreev_phase / 2) * sigma_y()
                    + self.er * sigma_x()
                )
                # Inductive and Josephson contributions (diagonal, add to eigenvalues)
                inductive_contribution = 0.5 * self.El * phi_val**2
                josephson_contribution = -self.Ej * np.cos(phi_val - self.phase)
            elif self.flux_grouping == "EL":
                andreev_phase = phi_val
                andreev_term = (
                    -self.Gamma * np.cos(andreev_phase / 2) * sigma_z()
                    - self.delta_Gamma * np.sin(andreev_phase / 2) * sigma_y()
                    + self.er * sigma_x()
                )
                # Inductive and Josephson contributions (diagonal, add to eigenvalues)
                inductive_contribution = 0.5 * self.El * (phi_val + self.phase) ** 2
                josephson_contribution = -self.Ej * np.cos(phi_val)
            else:
                raise ValueError(f"Unknown flux_grouping: {self.flux_grouping}")

            berry_contribution = (
                self._berry_contribution(andreev_phase) if include_berry else 0.0
            )

            # Diagonalize only the Andreev part
            if return_evecs:
                andreev_evals, andreev_evecs = eigh(
                    andreev_term,
                    check_finite=False,
                )
                # Add diagonal contributions to eigenvalues
                evals_array[i] = (
                    andreev_evals
                    + inductive_contribution
                    + josephson_contribution
                    + berry_contribution
                )

                # Gauge fixing: ensure phase continuity with previous point
                if i > 0:
                    for j in range(2):  # For each eigenstate
                        # Compute overlap with previous eigenstate
                        overlap = np.vdot(evecs_array[i - 1, :, j], andreev_evecs[:, j])
                        # If overlap is negative or has large imaginary part, fix phase
                        if np.real(overlap) < 0:
                            andreev_evecs[:, j] *= -1

                evecs_array[i] = andreev_evecs
            else:
                andreev_evals = eigh(
                    andreev_term,
                    eigvals_only=True,
                    check_finite=False,
                )
                # Add diagonal contributions to eigenvalues
                evals_array[i] = (
                    andreev_evals
                    + inductive_contribution
                    + josephson_contribution
                    + berry_contribution
                )

        if return_evecs:
            return evals_array, evecs_array
        return evals_array

    def _berry_contribution(
        self, phi: Union[float, np.ndarray]
    ) -> Union[float, np.ndarray]:
        """
        Returns the scalar Berry contribution in the adiabatic Andreev basis.

        The correction is
        1/4 * [ (d chi / d phi)^2 + (d theta / d phi)^2 ],
        multiplied by the identity in the Andreev subspace.
        """
        phi_array = np.asarray(phi, dtype=float)
        b_perp_sq = (
            self.Gamma**2 * np.cos(phi_array / 2) ** 2
            + self.delta_Gamma**2 * np.sin(phi_array / 2) ** 2
        )
        b_sq = b_perp_sq + self.er**2
        b_perp = np.sqrt(b_perp_sq)

        with np.errstate(divide="ignore", invalid="ignore"):
            dchi_dphi = self.Gamma * self.delta_Gamma / (2 * b_perp_sq)
            dtheta_dphi = (
                self.er
                * (self.delta_Gamma**2 - self.Gamma**2)
                * np.sin(phi_array)
                / (4 * b_perp * b_sq)
            )
            berry_contribution = 0.25 * (dchi_dphi**2 + dtheta_dphi**2)

        if np.isscalar(phi):
            return float(berry_contribution)
        return berry_contribution

    def tphi_1_over_f_flux(
        self,
        A_noise: float = 1e-6,
        esys: tuple[np.ndarray, np.ndarray] = None,
        get_rate: bool = False,
        **kwargs,
    ) -> float:
        return self.tphi_1_over_f(
            A_noise,
            ["d_hamiltonian_d_phase", "d2_hamiltonian_d_phase2"],
            esys=esys,
            get_rate=get_rate,
            **kwargs,
        )

    def plot_wavefunction(
        self,
        which: Union[int, Iterable[int]] = 0,
        phi_grid: np.ndarray = None,
        esys: tuple[np.ndarray, np.ndarray] = None,
        scaling: Optional[float] = 1,
        plot_potential: bool = False,
        basis: Literal["phase", "charge"] = "phase",
        andreev_basis: Literal["static", "adiabatic"] = "static",
        mode: Literal["abs", "abs2", "real", "imag"] = "abs",
        **kwargs,
    ) -> tuple[plt.Figure, plt.Axes]:
        """
        Plot the wave function in the phi basis.

        Parameters
        ----------
        which : Union[int, Iterable[int]], optional
            Index or indices of desired wave function(s) (default is 0).
        phi_grid : np.ndarray, optional
            Custom grid for phi; if None, a default grid is used.
        esys : Tuple[np.ndarray, np.ndarray], optional
            Precomputed eigenvalues and eigenvectors.
        scaling : float, optional
            Scaling factor for the wavefunction (default is 1).
        plot_potential : bool, optional
            Whether to plot the potential (default is False).
        basis : Literal["phase", "charge"], optional
            Basis in which to return the wavefunction ('phase' or 'charge') (default is 'phase').
        andreev_basis : Literal["static", "adiabatic"], optional
            Andreev basis used for plotting. Use 'static' for the basis used to
            define the Hamiltonian matrix, or 'adiabatic' for the local
            phi-dependent Andreev eigenbasis.
        mode : Literal["abs", "abs2", "real", "imag"], optional
            Mode of the wavefunction ('abs', 'abs2', 'real', or 'imag') (default is 'abs').
        **kwargs
            Additional arguments for plotting. Can include:
            - fig_ax: Tuple[plt.Figure, plt.Axes], optional
                Figure and axes to use for plotting. If not provided, a new figure and axes are created.

        Returns
        -------
        Tuple[plt.Figure, plt.Axes]
            The figure and axes of the plot.
        """
        if isinstance(which, int):
            which = [which]

        if andreev_basis not in {"static", "adiabatic"}:
            raise ValueError("Invalid andreev_basis; must be 'static' or 'adiabatic'.")

        if phi_grid is None:
            phi_grid = np.linspace(-5 * np.pi, 5 * np.pi, 151)

        fig_ax = kwargs.get("fig_ax")
        if fig_ax is None:
            fig, ax = plt.subplots()
            fig.suptitle(self._generate_suptitle())
        else:
            fig, ax = fig_ax

        if plot_potential:
            potential = self.potential(phi=phi_grid)
            ax.plot(
                phi_grid / 2 / np.pi, potential[:, 0], color="black", label="Potential"
            )
            ax.plot(phi_grid / 2 / np.pi, potential[:, 1], color="black")

        for idx in which:
            wavefunc_data = self.wavefunction(
                which=idx,
                phi_grid=phi_grid,
                esys=esys,
                basis=basis,
                andreev_basis=andreev_basis,
            )
            phi_basis_labels = wavefunc_data["basis_labels"]
            wavefunc_amplitudes = wavefunc_data["amplitudes"]
            wavefunc_energy = wavefunc_data["energy"]

            if mode == "abs":
                y_values = np.abs(wavefunc_amplitudes[0])
                y_values_down = np.abs(wavefunc_amplitudes[1])
            elif mode == "abs2":
                y_values = np.abs(wavefunc_amplitudes[0]) ** 2
                y_values_down = np.abs(wavefunc_amplitudes[1]) ** 2
            elif mode == "real":
                y_values = wavefunc_amplitudes[0].real
                y_values_down = wavefunc_amplitudes[1].real
            elif mode == "imag":
                y_values = wavefunc_amplitudes[0].imag
                y_values_down = wavefunc_amplitudes[1].imag
            else:
                raise ValueError("Invalid mode; must be 'abs', 'real', or 'imag'.")

            ax.plot(
                phi_basis_labels / 2 / np.pi,
                wavefunc_energy + scaling * y_values,
                label=rf"$\Psi_{idx} \uparrow $",
            )
            ax.plot(
                phi_basis_labels / 2 / np.pi,
                wavefunc_energy + scaling * y_values_down,
                label=rf"$\Psi_{idx} \downarrow $",
            )

        if basis == "phase":
            ax.set_xlabel(r"$\Phi / \Phi_0$")
            ax.set_ylabel(r"$\psi(\varphi)$, Energy [GHz]")
        elif basis == "charge":
            ax.set_xlabel(r"$n$")
            ax.set_ylabel(r"$\psi(n)$, Energy [GHz]")

        ax.legend()
        ax.grid(True)

        return fig, ax

    def plot_state(
        self,
        which: int = 0,
        phi_grid: np.ndarray = None,
        n_grid: np.ndarray = None,
        wigner_func: bool = False,
        esys: tuple[np.ndarray, np.ndarray] = None,
        plot_bloch: bool = False,
        **kwargs,
    ) -> tuple[plt.Figure, plt.Axes]:
        """
        Plot the Wigner function of the state and the Bloch sphere.

        Parameters
        ----------
        which : int, optional
            Index of desired wavefunction (default is 0).
        phi_grid : np.ndarray, optional
            Custom grid for phi; if None, a default grid is used.
        n_grid : np.ndarray, optional
            Custom grid for n; if None, a default grid is used.
        wigner_func : bool, optional
            Precomputed wigner_func function (default is False).
        esys : Tuple[np.ndarray, np.ndarray], optional
            Precomputed eigenvalues and eigenvectors.
        plot_bloch : bool, optional
            Whether to plot the Bloch sphere (default is False).

        **kwargs : dict, optional
            Additional arguments for plotting. Can include:
            - fig_ax: Tuple[plt.Figure, plt.Axes], optional
                Figure and axes to use for plotting. If not provided, a new figure and axes are created.
            - cmap: str, optional
                Colormap to use for the Wigner function (default is 'seismic').
            - bloch_view: Tuple[float, float], optional
                Tuple with (elevation, azimuth) for Bloch sphere view (default is (-30, 60)).
            - bloch_position: Tuple[float, float, float, float], optional
                Position of the Bloch sphere inset in figure coordinates (left, bottom, width, height).
                If not provided, a default position is calculated.

        Returns
        -------
        Tuple[plt.Figure, plt.Axes]
            The figure and axes of the plot.
        """
        if phi_grid is None:
            phi_grid = np.linspace(-5, 5, 151)
        if n_grid is None:
            n_grid = np.linspace(-5, 5, 151)

        if wigner_func is False:
            wigner_func = self.wigner(
                which=which, phi_grid=phi_grid, n_grid=n_grid, esys=esys
            )

        fig_ax = kwargs.get("fig_ax")
        if fig_ax is None:
            fig, ax = plt.subplots()
            fig.suptitle(self._generate_suptitle())
        else:
            fig, ax = fig_ax

        cmap = kwargs.get("cmap", "seismic")

        map = ax.imshow(
            wigner_func,
            aspect="auto",
            origin="lower",
            extent=[phi_grid[0], phi_grid[-1], n_grid[0], n_grid[-1]],
            cmap=cmap,
        )

        vmax = np.max(np.abs(wigner_func))
        vmin = -vmax
        map.set_clim(vmin, vmax)

        ax.set_xlabel(r"$\varphi/2\pi$")
        ax.set_ylabel(r"$n$")
        # ax.set_aspect('equal')

        if plot_bloch:
            bloch_view = kwargs.get("bloch_view", (30, -60))
            bloch_position = kwargs.get(
                "bloch_position"
            )  # Custom position for Bloch sphere

            bbox = (
                ax.get_position()
            )  # posición del Axes principal en coordenadas de figura
            if bloch_position is None:
                inset_width = 0.3 * bbox.width
                inset_height = 0.3 * bbox.height
                inset_left = bbox.width - inset_width * 1.01
                inset_bottom = bbox.height - inset_height * 1.01
                bloch_position = [inset_left, inset_bottom, inset_width, inset_height]

            bloch_position[0] += bbox.x0
            bloch_position[1] += bbox.y0
            inset_ax = fig.add_axes(bloch_position, projection="3d")
            inset_ax.view_init(elev=bloch_view[0], azim=bloch_view[1])
            rho_reduced = self.reduced_density_matrix(which=which, esys=esys, subsys=1)
            rho_reduced_qobj = Qobj(rho_reduced)

            b = Bloch(fig=fig, axes=inset_ax)
            b.vector_width = 1.5
            b.xlabel = ["", ""]
            b.ylabel = ["", ""]
            b.zlabel = ["", ""]

            b.add_states(state=rho_reduced_qobj, colors="blue")
            b.render()

        # fig.colorbar(map, ax=ax, label=r"$W(\Phi,n)$", shrink=0.8)
        # fig.tight_layout()

        return fig, ax
