import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from IPython.display import Math, display
from matplotlib.axes import Axes
from scipy import constants as const
from scipy.integrate import quad
from scipy.optimize import root_scalar
from scipy.signal import argrelmax,argrelmin
from tqdm import tqdm

import fewfermions.analysis.units as si
from fewfermions.simulate.traps.twod.trap import PancakeTrap
from fewfermions.style import FIGS_PATH, setup

colors, colors_alpha = setup()


def plot_solutions(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4):
    """Plot the potential and solutions for a given trap object
        (diplay_plot=-1 for all, 0,...,n for individual ones and -2 for none)   
    """

    x, y, z = trap.x, trap.y, trap.z
    axial_width = trap.get_tweezer_rayleigh()

    pot_ax = trap.subs(trap.get_potential())
    pot_diff_ax = sp.diff(pot_ax, trap.z)
    pot_diff2_ax = sp.diff(pot_diff_ax, trap.z)
    pot_diff3_ax = sp.diff(pot_diff2_ax, trap.z)
    pot_diff_ax_numpy = sp.lambdify(trap.z, pot_diff_ax.subs({x: 0, y: 0}))
    pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({x: 0, y: 0}))
    pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({x: 0, y: 0}))

    barrier = root_scalar(
        pot_diff_ax_numpy,
        x0=1.5 * float(trap.subs(axial_width)),
        fprime=pot_diff2_ax_numpy,
        xtol=1e-18,
        fprime2=pot_diff3_ax_numpy,
    ).root
    minimum = root_scalar(
            pot_diff_ax_numpy,
            x0=0,
            fprime=pot_diff2_ax_numpy,
            xtol=1e-28,
            fprime2=pot_diff3_ax_numpy,
        ).root

    # Solve the hamiltonian numerically in axial direction
    energies, states, potential, coords = trap.nstationary_solution(
        trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
    )

    # States that are below the potential barrier
    bound_states = energies < potential(barrier)


    # Density of states is larger on the left than on the right
    # Likely that the state in question is a true bound state
    true_bound_states = np.logical_and(
        bound_states,
        np.sum(states[:, coords[z] < barrier] ** 2, axis=1)
        > np.sum(states[:, coords[z] > barrier] ** 2, axis=1),
    )

    transmission_probability = np.full_like(energies, np.nan, dtype=float)
    for j, energy in enumerate(energies):
        if not true_bound_states[j]:
            continue
        intersect_end = root_scalar(
            lambda x: potential(x) - energy,
            bracket=(barrier, 3 * float(trap.subs(axial_width))),
        ).root
        intersect_start = root_scalar(
            lambda x: potential(x) - energy,
            bracket=(minimum, barrier),
        ).root
        barrier_interval = np.logical_and(
            coords[z] > intersect_start, coords[z] < intersect_end
        )
        s = quad(
            lambda x: np.sqrt(
                2
                * float(trap.subs(trap.m))
                * np.clip(potential(x) - energy, a_min=0, a_max=None)
            )
            / const.hbar,
            intersect_start,
            intersect_end,
        )
        transmission_probability[j] = sp.exp(-2 * s[0])
    lifetime = (
        const.h / (transmission_probability * np.abs(energies - potential(minimum)))
    )
    print(f"{lifetime[true_bound_states]} s")

    z_np = np.linspace(left_cutoff, right_cutoff, num=n_pot_steps)

    ax: plt.Axes
    fig, ax = plt.subplots(figsize=(2.5, 2.5))
    # ax.set_title("Axial")
    abs_min = np.min(potential(z_np))
    print(abs_min)
    ax.fill_between(
        z_np / si.um,
        potential(z_np) / const.h / si.kHz,
        abs_min / const.h / si.kHz,
        fc=colors_alpha["red"],
        alpha=0.5,
    )
    # ax2 = ax.twinx()

    count = 0
    for i, bound in enumerate(true_bound_states):
        if not bound:
            continue
        energy = energies[i]
        state = states[i]
        ax.plot(
            z_np / si.um,
            np.where(
                (energy > potential(z_np)) & (z_np < barrier),
                energy / const.h / si.kHz,
                np.nan,
            ),
            c="k",
            lw=0.5,
            marker="None",
        )
        if display_plot == -1 or display_plot == count:
            ax.plot(z_np/si.um, state**2 *state_mult, marker="None", c="k")
        count += 1



    ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz, marker="None")
    ax.vlines(barrier/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
    ax.vlines(minimum/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
    ax.set_title(f"{np.sum(true_bound_states)} bound solutions, power: {trap.subs(trap.power_tweezer)/si.mW}mW")
    ax.set_xlabel(r"z ($\mathrm{\mu m}$)")
    ax.set_ylabel(r"E / h (kHz)")

def plot_solutions_ax(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4):
    """Plot the potential and solutions in the z-direction for a given trap object with no spilling
        (diplay_plot=-1 for all, 0,...,n for individual ones and -2 for none)   
    """
    #turn off magnetic gradient and gravity
    trap[trap.g] = 0
    initial_gradient = trap.grad_z
    trap[trap.grad_z] = 0

    trap
    x, y, z = trap.x, trap.y, trap.z
    axial_width = trap.get_tweezer_rayleigh()

    pot_ax = trap.subs(trap.get_potential())
    pot_diff_ax = sp.diff(pot_ax, trap.z)
    pot_diff2_ax = sp.diff(pot_diff_ax, trap.z)
    pot_diff3_ax = sp.diff(pot_diff2_ax, trap.z)
    pot_diff_ax_numpy = sp.lambdify(trap.z, pot_diff_ax.subs({x: 0, y: 0}))
    pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({x: 0, y: 0}))
    pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({x: 0, y: 0}))

    """
    barrier = root_scalar(
        pot_diff_ax_numpy,
        x0=1.5 * float(trap.subs(axial_width)),
        fprime=pot_diff2_ax_numpy,
        xtol=1e-18,
        fprime2=pot_diff3_ax_numpy,
    ).root
    minimum = root_scalar(
            pot_diff_ax_numpy,
            x0=0,
            fprime=pot_diff2_ax_numpy,
            xtol=1e-28,
            fprime2=pot_diff3_ax_numpy,
        ).root
        """

    # Solve the hamiltonian numerically in axial direction
    energies, states, potential, coords = trap.nstationary_solution(
        trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
    )
    
    # States that are below the potential barrier
    #bound_states = energies < potential(barrier)
    bound_states = energies < potential(left_cutoff)

    # Density of states is larger on the left than on the right
    # Likely that the state in question is a true bound state
    """true_bound_states = np.logical_and(
        bound_states,
        np.sum(states[:, coords[x] < barrier] ** 2, axis=1)
        > np.sum(states[:, coords[x] > barrier] ** 2, axis=1),
    )"""
    true_bound_states = bound_states
    """
    transmission_probability = np.full_like(energies, np.nan, dtype=float)
    for j, energy in enumerate(energies):
        if not true_bound_states[j]:
            continue
        intersect_end = root_scalar(
            lambda x: potential(x) - energy,
            bracket=(barrier, 3 * float(trap.subs(axial_width))),
        ).root
        intersect_start = root_scalar(
            lambda x: potential(x) - energy,
            bracket=(minimum, barrier),
        ).root
        barrier_interval = np.logical_and(
            coords[x] > intersect_start, coords[x] < intersect_end
        )
        s = quad(
            lambda x: np.sqrt(
                2
                * float(trap.subs(trap.m))
                * np.clip(potential(x) - energy, a_min=0, a_max=None)
            )
            / const.hbar,
            intersect_start,
            intersect_end,
        )
        transmission_probability[j] = sp.exp(-2 * s[0])
    lifetime = (
        const.h / (transmission_probability * np.abs(energies - potential(minimum)))
    )
    print(f"{lifetime[true_bound_states]} s")"""

    z_np = np.linspace(left_cutoff, right_cutoff, num=n_pot_steps)

    ax: plt.Axes
    fig, ax = plt.subplots(figsize=(2.5, 2.5))
    # ax.set_title("Axial")
    abs_min = np.min(potential(z_np))
    print(abs_min)
    ax.fill_between(
        z_np / si.um,
        potential(z_np) / const.h / si.kHz,
        abs_min / const.h / si.kHz,
        fc=colors_alpha["red"],
        alpha=0.5,
    )
    # ax2 = ax.twinx()
    
    count = 0
    for i, bound in enumerate(true_bound_states):
        if not bound:
            continue
        energy = energies[i]
        print((np.max(potential(z_np))-energy)/ const.h / si.kHz)
        state = states[i]
        ax.plot(
            z_np / si.um,
            np.where(
                (energy > potential(z_np)), #& (z_np < barrier),
                energy / const.h / si.kHz,
                np.nan,
            ),
            c="k",
            lw=0.5,
            marker="None",
        )
        if display_plot == -1 or display_plot == count:
            ax.plot(z_np/si.um, state**2 *state_mult, marker="None", c="k")
        count += 1



    ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz, marker="None")
    #ax.vlines(barrier/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
    #ax.vlines(minimum/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
    #ax.set_title(f"{np.sum(true_bound_states)} bound solutions, power: {trap.subs(trap.power_tweezer)/si.mW}mW")
    ax.set_xlabel(r"z ($\mathrm{\mu m}$)")
    ax.set_ylabel(r"E / h (kHz)")

    #turn on magnetic gradient and gravity
    trap[trap.g] = const.g
    trap[trap.grad_z] = initial_gradient

def plot_solutions_rad(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4):
    """Plot the potential and solutions in the x-direction for a given trap object
        (diplay_plot=-1 for all, 0,...,n for individual ones and -2 for none)   
    """

    x, y, z = trap.x, trap.y, trap.z
    axial_width = trap.get_tweezer_rayleigh()

    pot_ax = trap.subs(trap.get_potential())
    pot_diff_ax = sp.diff(pot_ax, trap.x)
    pot_diff2_ax = sp.diff(pot_diff_ax, trap.x)
    pot_diff3_ax = sp.diff(pot_diff2_ax, trap.x)
    pot_diff_ax_numpy = sp.lambdify(trap.x, pot_diff_ax.subs({x: 0, y: 0}))
    pot_diff2_ax_numpy = sp.lambdify(trap.x, pot_diff2_ax.subs({x: 0, y: 0}))
    pot_diff3_ax_numpy = sp.lambdify(trap.x, pot_diff3_ax.subs({x: 0, y: 0}))

    """
    barrier = root_scalar(
        pot_diff_ax_numpy,
        x0=1.5 * float(trap.subs(axial_width)),
        fprime=pot_diff2_ax_numpy,
        xtol=1e-18,
        fprime2=pot_diff3_ax_numpy,
    ).root
    minimum = root_scalar(
            pot_diff_ax_numpy,
            x0=0,
            fprime=pot_diff2_ax_numpy,
            xtol=1e-28,
            fprime2=pot_diff3_ax_numpy,
        ).root
        """

    # Solve the hamiltonian numerically in axial direction
    energies, states, potential, coords = trap.nstationary_solution(
        trap.x, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
    )
    
    # States that are below the potential barrier
    #bound_states = energies < potential(barrier)
    bound_states = energies < potential(left_cutoff)

    # Density of states is larger on the left than on the right
    # Likely that the state in question is a true bound state
    """true_bound_states = np.logical_and(
        bound_states,
        np.sum(states[:, coords[x] < barrier] ** 2, axis=1)
        > np.sum(states[:, coords[x] > barrier] ** 2, axis=1),
    )"""
    true_bound_states = bound_states
    """
    transmission_probability = np.full_like(energies, np.nan, dtype=float)
    for j, energy in enumerate(energies):
        if not true_bound_states[j]:
            continue
        intersect_end = root_scalar(
            lambda x: potential(x) - energy,
            bracket=(barrier, 3 * float(trap.subs(axial_width))),
        ).root
        intersect_start = root_scalar(
            lambda x: potential(x) - energy,
            bracket=(minimum, barrier),
        ).root
        barrier_interval = np.logical_and(
            coords[x] > intersect_start, coords[x] < intersect_end
        )
        s = quad(
            lambda x: np.sqrt(
                2
                * float(trap.subs(trap.m))
                * np.clip(potential(x) - energy, a_min=0, a_max=None)
            )
            / const.hbar,
            intersect_start,
            intersect_end,
        )
        transmission_probability[j] = sp.exp(-2 * s[0])
    lifetime = (
        const.h / (transmission_probability * np.abs(energies - potential(minimum)))
    )
    print(f"{lifetime[true_bound_states]} s")"""

    z_np = np.linspace(left_cutoff, right_cutoff, num=n_pot_steps)

    ax: plt.Axes
    fig, ax = plt.subplots(figsize=(2.5, 2.5))
    # ax.set_title("Axial")
    abs_min = np.min(potential(z_np))
    print(abs_min)
    ax.fill_between(
        z_np / si.um,
        potential(z_np) / const.h / si.kHz,
        abs_min / const.h / si.kHz,
        fc=colors_alpha["red"],
        alpha=0.5,
    )
    # ax2 = ax.twinx()
    
    count = 0
    for i, bound in enumerate(true_bound_states):
        if not bound:
            continue
        energy = energies[i]
        print((np.max(potential(z_np))-energy)/ const.h / si.kHz)
        state = states[i]
        ax.plot(
            z_np / si.um,
            np.where(
                (energy > potential(z_np)), #& (z_np < barrier),
                energy / const.h / si.kHz,
                np.nan,
            ),
            c="k",
            lw=0.5,
            marker="None",
        )
        if display_plot == -1 or display_plot == count:
            ax.plot(z_np/si.um, state**2 *state_mult, marker="None", c="k")
        count += 1



    ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz, marker="None")
    #ax.vlines(barrier/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
    #ax.vlines(minimum/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
    #ax.set_title(f"{np.sum(true_bound_states)} bound solutions, power: {trap.subs(trap.power_tweezer)/si.mW}mW")
    ax.set_xlabel(r"z ($\mathrm{\mu m}$)")
    ax.set_ylabel(r"E / h (kHz)")

def plot_occupation(trap,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,n_spill_steps=100,power_fac_down=0.4,power_fac_up=0.7,n_pot_steps=1000, results=False):
    """
    Plotting the occupation of the tweezer when varying the tweezer power.
    """
    x, y, z = trap.x, trap.y, trap.z
    axial_width = trap.get_tweezer_rayleigh()

    initial_power = trap[trap.power_tweezer]
    spill_power_factor = np.linspace(power_fac_up, power_fac_down, num=n_spill_steps)
    powers = trap[trap.power_tweezer] * spill_power_factor
    atom_number = np.zeros_like(powers,dtype=float)

    # Resolution of the potential when solving numerically
    #n_pot_steps = 1000

    for i, power in enumerate(tqdm(powers)):
        trap[trap.power_tweezer] = power
        # Solve the hamiltonian numerically in axial direction
        energies, states, potential, coords = trap.nstationary_solution(
            trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
        )

        # Determine the potential and its derivatives
        pot_ax = trap.subs(trap.get_potential())
        pot_diff_ax = sp.diff(pot_ax, trap.z)
        pot_diff2_ax = sp.diff(pot_diff_ax, trap.z)
        pot_diff3_ax = sp.diff(pot_diff2_ax, trap.z)
        pot_diff_ax_numpy = sp.lambdify(trap.z, pot_diff_ax.subs({x: 0, y: 0}))
        pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({x: 0, y: 0}))
        pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({x: 0, y: 0}))

        barrier = root_scalar(
            pot_diff_ax_numpy,
            x0=1.5 * float(trap.subs(axial_width)),
            fprime=pot_diff2_ax_numpy,
            xtol=1e-28,
            fprime2=pot_diff3_ax_numpy,
        ).root
        minimum = root_scalar(
            pot_diff_ax_numpy,
            x0=0,
            fprime=pot_diff2_ax_numpy,
            xtol=1e-28,
            fprime2=pot_diff3_ax_numpy,
        ).root
        # States that are below the potential barrier
        bound_states = energies < potential(barrier)

        # Density of states is larger on the left than on the right
        # Likely that the state in question is a true bound state
        true_bound_states = np.logical_and(
            bound_states,
            np.sum(states[:, coords[z] < barrier] ** 2, axis=1)
            > np.sum(states[:, coords[z] > barrier] ** 2, axis=1),
        )
        
        if t_spill == 0:
            atom_number[i] = np.sum(true_bound_states)
            continue

        transmission_probability = np.full_like(energies, np.nan, dtype=float)
        for j, energy in enumerate(energies):
            if not true_bound_states[j]:
                continue
            intersect_end = root_scalar(
                lambda x: potential(x) - energy,
                bracket=(barrier, 3 * float(trap.subs(axial_width))),
            ).root
            intersect_start = root_scalar(
                lambda x: potential(x) - energy,
                bracket=(minimum, barrier),
            ).root
            barrier_interval = np.logical_and(
                coords[z] > intersect_start, coords[z] < intersect_end
            )
            s = quad(
                lambda x: np.sqrt(
                    2
                    * float(trap.subs(trap.m))
                    * np.clip(potential(x) - energy, a_min=0, a_max=None)
                )
                / const.hbar,
                intersect_start,
                intersect_end,
            )
            transmission_probability[j] = sp.exp(-2 * s[0])
        tunneling_rate = (
            transmission_probability * np.abs(energies - potential(minimum)) / const.h
        )
        atom_number[i] = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))

        ax: plt.Axes
    fig: plt.Figure
    fig, ax = plt.subplots(figsize=(2.5, 2.5))

    ax.set_title(f"initial power:{initial_power/si.mW}mW, B':{trap[trap.grad_z]/si.G*si.cm}G/cm, w_0:{trap[trap.waist_tweezer]/si.um}um, wvl:{trap[trap.wvl]/si.nm}nm")
    ax.set_xlabel("rel. tweezer power (a.u.)")
    ax.set_ylabel("atom number")
    ax.plot(spill_power_factor, atom_number, marker="None")
    #ax.fill_between(spill_power_factor, atom_number, fc=colors_alpha["red"], alpha=0.5)

    if results:
        return spill_power_factor, atom_number


def plot_occupation_grad(trap,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,n_spill_steps=100,grad_fac_down=0.4,grad_fac_up=0.7,n_pot_steps=1000):
    """
    Plotting the occupation of the tweezer when varying the magnetic gradient.
    """

    x, y, z = trap.x, trap.y, trap.z
    axial_width = trap.get_tweezer_rayleigh()

    initial_power = trap[trap.power_tweezer]
    initial_grad = trap[trap.grad_z]
    spill_grad_factor = np.linspace(grad_fac_up, grad_fac_down, num=n_spill_steps)
    grads = trap[trap.grad_z] * spill_grad_factor
    atom_number = np.zeros_like(grads)

    # Resolution of the potential when solving numerically
    #n_pot_steps = 1000

    for i, grad in enumerate(tqdm(grads)):
        trap[trap.grad_z] = grad
        # Solve the hamiltonian numerically in axial direction
        energies, states, potential, coords = trap.nstationary_solution(
            trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
        )

        # Determine the potential and its derivatives
        pot_ax = trap.subs(trap.get_potential())
        pot_diff_ax = sp.diff(pot_ax, trap.z)
        pot_diff2_ax = sp.diff(pot_diff_ax, trap.z)
        pot_diff3_ax = sp.diff(pot_diff2_ax, trap.z)
        pot_diff_ax_numpy = sp.lambdify(trap.z, pot_diff_ax.subs({x: 0, y: 0}))
        pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({x: 0, y: 0}))
        pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({x: 0, y: 0}))

        barrier = root_scalar(
            pot_diff_ax_numpy,
            x0=1.5 * float(trap.subs(axial_width)),
            fprime=pot_diff2_ax_numpy,
            xtol=1e-28,
            fprime2=pot_diff3_ax_numpy,
        ).root
        minimum = root_scalar(
            pot_diff_ax_numpy,
            x0=0,
            fprime=pot_diff2_ax_numpy,
            xtol=1e-28,
            fprime2=pot_diff3_ax_numpy,
        ).root
        # States that are below the potential barrier
        bound_states = energies < potential(barrier)

        # Density of states is larger on the left than on the right
        # Likely that the state in question is a true bound state
        true_bound_states = np.logical_and(
            bound_states,
            np.sum(states[:, coords[z] < barrier] ** 2, axis=1)
            > np.sum(states[:, coords[z] > barrier] ** 2, axis=1),
        )
        
        if t_spill == 0:
            atom_number[i] = np.sum(true_bound_states)
            continue

        transmission_probability = np.full_like(energies, np.nan, dtype=float)
        for j, energy in enumerate(energies):
            if not true_bound_states[j]:
                continue
            intersect_end = root_scalar(
                lambda x: potential(x) - energy,
                bracket=(barrier, 3 * float(trap.subs(axial_width))),
            ).root
            intersect_start = root_scalar(
                lambda x: potential(x) - energy,
                bracket=(minimum, barrier),
            ).root
            barrier_interval = np.logical_and(
                coords[z] > intersect_start, coords[z] < intersect_end
            )
            s = quad(
                lambda x: np.sqrt(
                    2
                    * float(trap.subs(trap.m))
                    * np.clip(potential(x) - energy, a_min=0, a_max=None)
                )
                / const.hbar,
                intersect_start,
                intersect_end,
            )
            transmission_probability[j] = sp.exp(-2 * s[0])
        tunneling_rate = (
            transmission_probability * np.abs(energies - potential(minimum)) / const.h
        )
        atom_number[i] = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))

        ax: plt.Axes
    fig: plt.Figure
    fig, ax = plt.subplots(figsize=(2.5, 2.5))

    ax.plot(spill_grad_factor, atom_number, marker="None")
    ax.fill_between(spill_grad_factor, atom_number, fc=colors_alpha["red"], alpha=0.5)
    ax.set_title(f"power:{initial_power/si.mW}mW, initial B':{initial_grad/si.G*si.cm}G/cm, w_0:{trap[trap.waist_tweezer]/si.um}um, wvl:{trap[trap.wvl]/si.nm}nm")
    ax.set_xlabel("rel. gradient (a.u.)")
    ax.set_ylabel("atom number")