import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from scipy import constants as const
from scipy.integrate import quad
from scipy.optimize import root_scalar
from tqdm import tqdm

import trap_units as si


def plot_solutions(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4,plot=True,ret_num=False):
    """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),
    )

    if plot:
        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))
        ax.fill_between(
            z_np / si.um,
            potential(z_np) / const.h / si.kHz,
            abs_min / const.h / si.kHz,
            alpha=0.5,
        )

        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)")

    if ret_num:
        return np.sum(true_bound_states)
    
def solutions_power(trap,power,n_levels,left_cutoff,right_cutoff):
    trap[trap.power_tweezer] = power
    return plot_solutions(trap,n_levels,left_cutoff,right_cutoff,plot=False, ret_num=True)

def root_power(trap,num_atoms,bracket,n_levels,left_cutoff,right_cutoff):
    f = lambda x: solutions_power(trap,x,n_levels,left_cutoff,right_cutoff) - num_atoms
    return root_scalar(f,bracket=bracket).root

def sweep_power(trap, gradient,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,n_pot_steps=1000,max_atoms=10,n_plotpoints=100):
    """
    Sweeps over power and gets the power for 0 atom boundary and the precision level
    """

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

    trap[trap.grad_z] = gradient
    initial_power = float(4/3/np.sqrt(3)*trap.subs(zr) * np.pi* trap.subs(trap.m * trap.g + trap.mu_b * trap.grad_z) * trap.subs(trap.waist_tweezer**2/trap.a))
    final_power = root_power(trap,max_atoms,[initial_power,initial_power*5],n_levels,left_cutoff,right_cutoff)
    powers = np.linspace(initial_power,final_power,n_plotpoints)
    atom_number = np.zeros_like(powers,dtype=float)


    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(zr)),
            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_num = np.sum(true_bound_states)
            atom_number = np.append(atom_number,atom_num)
            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(zr))),
            ).root
            intersect_start = root_scalar(
                lambda x: potential(x) - energy,
                bracket=(minimum, barrier),
            ).root

            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]))

    return powers, atom_number

def sweep_power_old(trap, gradient,dp,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,max_spill_steps=200,n_pot_steps=1000,max_atoms=10):
    """
    Sweeps over power and gets the power for 0 atom boundary and the precision level
    """

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

    trap[trap.grad_z] = gradient
    initial_power = 4/3/np.sqrt(3)*trap.subs(zr) * np.pi* trap.subs(trap.m * trap.g + trap.mu_b * trap.grad_z) * trap.subs(trap.waist_tweezer**2/trap.a)
    trap[trap.power_tweezer] = initial_power
    powers = np.array([initial_power],dtype=float)
    atom_number = np.array([0.0],dtype=float)

    i = 0
    pbar = tqdm(disable=True)
    while atom_number[i] <max_atoms and i<max_spill_steps:
        #print(i)
        trap[trap.power_tweezer] = initial_power + (i+1)*dp
        powers = np.append(powers, initial_power + (i+1)*dp)

        # 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(zr)),
            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_num = np.sum(true_bound_states)
            atom_number = np.append(atom_number,atom_num)
            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(zr))),
            ).root
            intersect_start = root_scalar(
                lambda x: potential(x) - energy,
                bracket=(minimum, barrier),
            ).root

            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_num = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))
        atom_number = np.append(atom_number,atom_num)
        i += 1
        pbar.update(1)

    pbar.close()
    if i == max_spill_steps:
        print(f"{max_atoms} atoms not reached")
        return powers,atom_number
        raise Exception(f"{max_atoms} atoms not reached")

    return powers, atom_number

def solutions_gradient(trap,gradient,n_levels,left_cutoff,right_cutoff):
    trap[trap.grad_z] = gradient
    return plot_solutions(trap,n_levels,left_cutoff,right_cutoff,plot=False, ret_num=True)

def root_gradient(trap,num_atoms,bracket,n_levels,left_cutoff,right_cutoff):
    f = lambda x: solutions_gradient(trap,x,n_levels,left_cutoff,right_cutoff) - num_atoms
    return root_scalar(f,bracket=bracket).root

def sweep_gradient(trap, power,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,n_pot_steps=1000,max_atoms=10,n_plotpoints=100):
    """
    Sweeps over gradient and gets the gradient for 0 atom boundary and the precision level
    """

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

    trap[trap.power_tweezer] = power
    initial_grad = float(trap.subs(1/trap.mu_b * (3*np.sqrt(3)/4 * power * trap.a/np.pi/trap.waist_tweezer**2/zr - trap.m * trap.g)))
    final_grad = root_gradient(trap,max_atoms,[-2.89*si.G/si.cm,initial_grad],n_levels,left_cutoff,right_cutoff)
    gradients = np.linspace(initial_grad,final_grad,n_plotpoints,dtype=float)
    atom_number = np.zeros_like(gradients,dtype=float)

    for i, gradient in enumerate(tqdm(gradients)):
        #print(i)
        trap[trap.grad_z] = gradient

        # 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(zr)),
            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_num = np.sum(true_bound_states)
            atom_number = np.append(atom_number,atom_num)
            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(zr))),
            ).root
            intersect_start = root_scalar(
                lambda x: potential(x) - energy,
                bracket=(minimum, barrier),
            ).root

            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]))

    return gradients, atom_number

def sweep_gradient_old(trap, power,dgrad,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,max_spill_steps=200,n_pot_steps=1000,max_atoms=10):
    """
    Sweeps over gradient and gets the gradient for 0 atom boundary and the precision level
    """

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

    trap[trap.power_tweezer] = power
    initial_grad = float(trap.subs(1/trap.mu_b * (3*np.sqrt(3)/4 * power * trap.a/np.pi/trap.waist_tweezer**2/zr - trap.m * trap.g)))
    trap[trap.grad_z] = initial_grad
    gradients = np.array([initial_grad],dtype=float)
    atom_number = np.array([0.0],dtype=float)

    i = 0
    pbar = tqdm(disable=True)
    while atom_number[i] <max_atoms and i<max_spill_steps:
        #print(i)
        trap[trap.grad_z] = initial_grad - (i+1)*dgrad
        gradients = np.append(gradients, initial_grad - (i+1)*dgrad)

        # 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(zr)),
            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_num = np.sum(true_bound_states)
            atom_number = np.append(atom_number,atom_num)
            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(zr))),
            ).root
            intersect_start = root_scalar(
                lambda x: potential(x) - energy,
                bracket=(minimum, barrier),
            ).root

            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_num = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))
        atom_number = np.append(atom_number,atom_num)
        i += 1
        pbar.update(1)

    pbar.close()
    if i == max_spill_steps:
        print(f"{max_atoms} atoms not reached")
        return gradients,atom_number
        raise Exception(f"{max_atoms} atoms not reached")

    return gradients, atom_number

def calculate_stepsize(powers, atom_number,plot=False,max_atoms=10,cutoff=0.1):
    """
    Given an array of powers (or gradients) and atom_numbers, returns the average length of steps.
    The step length is the area where atom number is within cutoff percent of an integer.
    """

    bool_array = np.logical_and(np.abs(atom_number - np.round(atom_number,0)) < cutoff, atom_number > cutoff) #find points where atomnumber is close to integer
    bool_array = np.logical_and(bool_array, atom_number < (max_atoms-0.5))
    if plot:
        plt.plot(powers,atom_number)
        plt.plot(powers[bool_array],atom_number[bool_array],"b.")

    diff = np.diff(bool_array.astype(int))  # Convert to int to use np.diff
    start_indices = np.where(diff == 1)[0] + 1  # Indices where blobs start
    end_indices = np.where(diff == -1)[0] + 1  # Indices where blobs end

    # Special case: handle if the array starts or ends with a True
    if bool_array[0]:
        start_indices = np.insert(start_indices, 0, 0)  # Add the start of the array
    if bool_array[-1]:
        end_indices = np.append(end_indices, len(bool_array))  # Add the end of the array

    #step length
    step_len = np.abs(np.mean(np.diff(powers)))

    # Blob lengths
    blob_lengths = float(np.mean((end_indices - start_indices)*step_len))

    # Results
    return blob_lengths