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=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,
            fc=colors_alpha["red"],
            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