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,find_peaks
from tqdm import tqdm

import fewfermions.analysis.units as si
from fewfermions.simulate.traps.twod.trap import DoubleTweezer

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

    # 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(left_cutoff)

    if plot:
        z_np = np.linspace(left_cutoff, right_cutoff, num=n_pot_steps)
        ax: plt.Axes
        fig, ax = plt.subplots(figsize=(5, 5))
        ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz,color="cornflowerblue" ,marker="None")
        ax.set_title(f"{np.sum(bound_states)} bound solutions, tweezer distance: {trap.subs(trap.distance_tweezers)/si.um}um")
        ax.set_xlabel(r"x ($\mathrm{\mu m}$)")
        ax.set_ylabel(r"E / h (kHz)")
        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,
            color="cornflowerblue"
        )

        count = 0
        for i, bound in enumerate(bound_states):
            if not bound:
                continue
            energy = energies[i]
            state = states[i]
            ax.plot(
                z_np / si.um,
                np.where(
                    (energy > potential(z_np)),
                    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 *state_mult, marker="None",label=f"state {count}")#, c="k")
            count += 1

        plt.legend()


    if ret_results:
        return np.sum(bound_states), energies[bound_states], states[bound_states], potential
    

def loop_distances(trap,distances,n_levels=10,n_pot_steps=2000):
    """
    returns unsorted arrays for energies, states and the potential for a given array of tweezer distances
    """
    left_cutoff = -0.5*np.max(distances)-3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])
    right_cutoff = 0.5*np.max(distances)+3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])
    x_np = np.linspace(left_cutoff,right_cutoff,n_pot_steps)

    energies = np.zeros((len(distances),n_levels))
    states = np.zeros((len(distances),n_levels,n_pot_steps))
    potential = np.zeros((len(distances),n_pot_steps))

    #diagonalise hamiltonian for all distances
    for i, dist in enumerate(distances):
        trap[trap.distance_tweezers] = dist
        num, ener, state, pot = plot_solutions(trap,n_levels,np.min(x_np),np.max(x_np),display_plot=0,state_mult=450,n_pot_steps=n_pot_steps,plot=False,ret_results=True)

        # fill in arrays
        energies[i,:len(ener)] = ener
        states[i,:len(ener)] = state
        potential[i] = pot(x_np)

    return energies, states, potential

def find_crossovers_topstate(energies):
    """
    returns the indices(distance) of all crossovers (with calculated and not calculated states) of the largest calculated state
    """
    arr = np.abs(np.gradient(np.gradient(energies[:,-1])))
    index = find_peaks(arr,width=(1,4))[0]
    return index

def find_crossovers(energies,mult_min_energy=0.01):
    """
    returns the indices(distance) of all crossovers between two calculated states
    (very dependent on mult_min_energy)
    """
    n_levels = len(energies[0,:])
    index = np.array([],dtype=int)
    swap_index =np.empty((0,2),dtype=int)
                        
    for i in range(n_levels):
        for j in range(i+1,n_levels):
            diff = np.abs(energies[:,i] - energies[:,j])
            mask = np.where(diff < mult_min_energy*np.max(diff),diff,mult_min_energy*np.max(diff))
            peaks = find_peaks(-mask)[0]
            index = np.append(index,peaks)
            for k in range(len(peaks)):
                swap_index = np.append(swap_index,[[i,j]],axis=0)

    return index, swap_index

def swapped_loop_distance(distances, energies, states, potentials):
    """
    returns the results of loop_distance() with energies and states matching to individual states
    """

    index_top = find_crossovers_topstate(energies)
    index, swap_index = find_crossovers(energies)

    new_energies = np.full((energies.shape[0],energies.shape[1]+len(index_top)),np.nan)
    new_states = np.full((states.shape[0],states.shape[1]+len(index_top),states.shape[2]),np.nan)

    swapped_index = np.arange(0,energies.shape[1])
    for i, dist in enumerate(distances): 
        if np.any(i == index):
            j = np.where(index == i)
            swapped_index[swap_index[j]] = swapped_index[np.roll(swap_index[j],1)]
            #print(f"crossover of states {swap_index[j]} at {dist/si.um:.2f}um")
            
        elif np.any(i == index_top):
            #print(f"crossover of top state at {dist/si.um:.2f}um")
            swapped_index[-1] = np.max(swapped_index)+1
            
                
        new_energies[i,swapped_index] = energies[i]
        new_states[i,swapped_index] = states[i]

    return new_energies, new_states, potentials, index_top, index, swap_index

def find_ass_tweezer(new_energies,new_states,return_deltaE=True):
    """
    Sorts energies and states into those coming from left and right tweezer.

    Returns the minimum energy difference of the groundstate of the shallow tweezer to any of the other tweezers eigenstates.
    Also returns the occupation number when the groundstates of both tweezers were initially occupied (when return_deltaE=True)
    """
    #seperate energies and states corresponding to the initial tweezer they belonged to
    mask_right = np.where(np.sum(new_states[0,:,:int(new_states.shape[2]/2)]**2,axis=1) < 0.5,True, False)
    mask_left = np.logical_not(mask_right)

    energies_right = new_energies[:,mask_right]
    energies_left = new_energies[:,mask_left]

    states_right = new_states[:,mask_right,:]
    states_left = new_states[:,mask_left,:]

    #identify GS of shallower tweezer and calculated minimal energy difference to other states
    if mask_left[np.nanargmin(new_energies[0])]:
        energy_GS = energies_right[:,0]
        state_GS = states_right[:,0]
        #smallest energy difference
        delta_E_min = np.nanmin(np.abs(energy_GS[:,np.newaxis] - energies_left))
        #final occupation if two GS are involved
        occ_numbers = np.array([0,np.searchsorted(energies_left[-1],energy_GS[-1])])
    else:
        energy_GS = energies_left[:,0]
        state_GS = states_left[:,0]
        #smallest energy difference
        delta_E_min = np.nanmin(np.abs(energy_GS[:,np.newaxis] - energies_right))
        #final occupation if two GS are involved
        occ_numbers = np.array([0,np.searchsorted(energies_right[-1],energy_GS[-1])])

    if return_deltaE:
        return energies_left, energies_right, states_left, states_right, delta_E_min, occ_numbers
    else:
        return energies_left, energies_right, states_left, states_right