import numpy as np
import pickle
from scipy import constants as const
from scipy.optimize import root_scalar
from scipy.fft import fftn, ifftn, fftfreq

from trap_numerics import functional_hamiltonian


def import_results(line_number):
    """
    Imports all stored results and variables from the calculation
    in line "line_number" of data/overview.txt
        out:
            trap (object): trap object used for diagonalisation
            func_ham (LinearOperator): hamiltonian used for diagonalisation
    """
    with open("data/overview.txt", "r") as file:
        for current_line_number, line in enumerate(file, start=1):
            if current_line_number == line_number:
                timecode = line[:19]  # Get the first 19 characters
                break  # Stop reading once we get the desired line

    data = np.load(f"data/{timecode}_results.npz",allow_pickle=True)
    results = {key: data[key] for key in data.keys()}

    with open(f"data/{timecode}_trap.npz", 'rb') as file:
        trap = pickle.load(file)
    
    func_ham = functional_hamiltonian(results["dvol"],results["pot"],
                                      float(trap.subs(trap.m)))

    return trap, func_ham, results


def get_localised_GS(state0, state1,
                     degenerate=True):
    """
    Returns the localised ground states of left and right tweezer.
        in: state0, state1 (3darray): first two states from diagonalisation
        out: GS_left, GS_right (3d array): ground states of both tweezer
    """
    size = state0.shape
    
    #if states are not degenerate, they don't have to be rotated into each other
    if not degenerate:
        #then transform to localised wavefunctions
        GS_LR1 = state0.real / np.sum(np.abs(state0))**2
        GS_LR2 = state1.real / np.sum(np.abs(state1))**2
    else:
        #rotate resulting wavefunctions into symmetric and anti-symm. states 
        def f(alpha):

            state_S1 = np.cos(alpha) * state0.real + np.sin(alpha)* state1.real
            state_S1 /= np.sqrt(np.sum(np.abs(state_S1)**2))

            state_A1 = np.sin(alpha) * state0.real - np.cos(alpha)* state1.real
            state_A1 /= np.sqrt(np.sum(np.abs(state_A1)**2))

            return np.sum(state_S1[int(size[0]/2):]) - np.sum(state_S1[:int(size[0]/2+1)])

        #find optimal angle
        # alpha = minimize_scalar(f,bounds=(-2*np.pi, 2*np.pi+0.1)).x
        alpha = root_scalar(f, x0=0.0001).root

        state_S1 = np.cos(alpha) * state0.real + np.sin(alpha)* state1.real
        state_S1 /= np.sqrt(np.sum(np.abs(state_S1)**2))

        state_A1 = np.sin(alpha) * state0.real - np.cos(alpha)* state1.real
        state_A1 /= np.sqrt(np.sum(np.abs(state_A1)**2))

        #then transform to localised wavefunctions
        GS_LR1 = (state_S1 + state_A1)/np.sqrt(2)
        GS_LR2 = (state_S1 - state_A1)/np.sqrt(2)

    #asign left and right
    if abs(np.sum(GS_LR1[int(size[0]/2):])) < abs(np.sum(GS_LR1[:int(size[0]/2)])):
        GS_left = GS_LR1
        GS_right = GS_LR2
    else:
        GS_left = GS_LR2
        GS_right = GS_LR1

    #state vectors are already real valued
    #Choose phase such that they are also both of same sign
    sign_left = np.sign(np.mean(np.sign(GS_left)))
    sign_right = np.sign(np.mean(np.sign(GS_right)))
    GS_left *= sign_left
    GS_right *= sign_right

    return GS_left, GS_right


def get_J(GS_left,GS_right, func_hamiltonian, dvol):
    """
    Returns the hopping amplitude between GS_left and GS_right.
        in: GS_left,GS_right (3d arrays): ground states of both tweezer
            func_hamiltonian (scipy.sparse.linalg.LinearOperator): Hamiltonian of the used potential
            dvol ((3,) array): grid spacing in all 3 directions
        out: J (float): hopping amplitude in SI units
    """
    #calculate volume element
    dV = np.prod(dvol)
    #make sure we're working with state vector and not wave function
    GS_left /= np.sum(np.abs(GS_left)**2)
    GS_right /= np.sum(np.abs(GS_right)**2)

    #find shape of grid
    size = GS_right.shape
    #call Hamiltonian on one side
    H_psi = np.reshape(func_hamiltonian(GS_right.flatten()), size)

    #calculate integral
    J = -(np.sum(GS_left.conj()/np.sqrt(dV)* H_psi/np.sqrt(dV))*dV)

    #check imaginary part of result (should vanish)
    if abs(J.real/J.imag) < 10:
        raise Exception(f"imaginary part of J too big: J = {J}")
    else:
        J = J.real

    return J

    
def get_U_s(trap, state, dvol):
    """
    Returns the contact interaction energy for a tweezer state.
        in: trap (object): trap object for parameters
            state (3d array): ground state of one tweezer
            dvol ((3,) array): grid spacing in all 3 directions
        out: U_s (float): contact interaction energy in SI units
    """
    #get paramters
    a_s = float(trap.subs(trap.a_s))
    mass = float(trap.subs(trap.m))
    #calculate volume element
    dV = np.prod(dvol)
    #make sure we're working with state vector
    state /= np.sum(np.abs(state)**2)

    #calculate integral over psi^4 (divide by sqrt(dV) to convert to wavefunction)
    U_s = np.sum(np.abs(state/np.sqrt(dV))**4)*dV
    U_s *= 4*np.pi*const.hbar**2*a_s/mass

    return U_s
    

def get_U_dd(trap, state, dvol):
    """
    Returns the dipole-dipole interaction energy for a tweezer state.
        in: trap (object): trap object for parameters
            state (3d array): ground state of one tweezer
            dvol ((3,) array): grid spacing in all 3 directions
        out: U_dd (float): dipole-dipole interaction energy in SI units
    """
    #get paramters
    mu = float(trap.subs(trap.mu_b))
    #get and normalise polarisation
    B_x = float(trap.subs(trap.B_x))
    B_y = float(trap.subs(trap.B_y))
    B_z = float(trap.subs(trap.B_z))
    polarisation = np.array([B_x, B_y, B_z])
    polarisation /= np.sum(np.abs(polarisation)**2)
    #calculate volume element
    dV = np.prod(dvol)

    #compute dipolar interactions using fourier transform

    #density
    rho = np.abs(state/np.sqrt(dV))**2
    # Compute Fourier transform of density
    rho_k = fftn(rho)

    #find shape of grid
    size = state.shape
    #set up fourier space
    kx = 2*np.pi*fftfreq(size[0], dvol[0])
    ky = 2*np.pi*fftfreq(size[1], dvol[1])
    kz = 2*np.pi*fftfreq(size[2], dvol[2])
    KX, KY, KZ = np.meshgrid(kx, ky, kz, indexing='ij')
    K2 = KX**2 + KY**2 + KZ**2
    K2[K2 == 0] = np.inf  # Avoid division by zero

    # Compute dipolar interaction kernel in Fourier space
    V_k = (-1 + 3* (polarisation[0]*KX + polarisation[1]*KY + polarisation[2]*KZ)**2 / K2) /3

    # Compute convolution in Fourier space
    U_dd_k = V_k * rho_k

    # Transform back to real space
    U_dd_real = ifftn(U_dd_k)

    # Integrate
    U_dd = np.sum(U_dd_real*rho) *dV
    #multiply by prefactor C_dd (4pi already in kernel)
    U_dd *= const.mu_0* mu**2

    #check imaginary part of result (should vanish)
    if abs(U_dd.real/U_dd.imag) < 10:
        raise Exception(f"imaginary part of U_dd too big: U_dd = {U_dd}")
    else:
        U_dd = U_dd.real

    return U_dd


def get_U(trap, state, dvol):
    """
    Returns the total interaction energy for a tweezer state.
        in: trap (object): trap object for parameters
            state (3d array): ground state of one tweezer
            dvol ((3,) array): grid spacing in all 3 directions
        out: U_dd (float): contact interaction energy in SI units
    """
    U_s = get_U_s(trap, state, dvol)
    U_dd = get_U_dd(trap, state, dvol)

    return U_s + U_dd

def get_NNI(trap, GS_left, GS_right, dvol):
    """
    Returns the dipole-dipole nearest-neighbour interaction energy of the GSs.
        in: trap (object): trap object for parameters
            GS_left,GS_right (3d array): ground states of both tweezers
            dvol ((3,) array): grid spacing in all 3 directions
        out: DeltaJ_lr (float): NNI energy in SI units
    """
    #get paramters
    mu = float(trap.subs(trap.mu_b))
    #get and normalise polarisation
    B_x = float(trap.subs(trap.B_x))
    B_y = float(trap.subs(trap.B_y))
    B_z = float(trap.subs(trap.B_z))
    polarisation = np.array([B_x, B_y, B_z])
    polarisation /= np.sum(np.abs(polarisation)**2)
    #calculate volume element
    dV = np.prod(dvol)

    #make sure we're working with state vector
    GS_left /= np.sum(np.abs(GS_left)**2)
    GS_right /= np.sum(np.abs(GS_right)**2)

    #compute dipolar interactions using fourier transform

    #density
    rho_r = np.abs(GS_right/np.sqrt(dV))**2
    rho_l = np.abs(GS_left/np.sqrt(dV))**2
    # Compute Fourier transform of density
    rho_k = fftn(rho_l)

    #find shape of grid
    size = GS_left.shape
    #set up fourier space
    kx = 2*np.pi*fftfreq(size[0], dvol[0])
    ky = 2*np.pi*fftfreq(size[1], dvol[1])
    kz = 2*np.pi*fftfreq(size[2], dvol[2])
    KX, KY, KZ = np.meshgrid(kx, ky, kz, indexing='ij')
    K2 = KX**2 + KY**2 + KZ**2
    K2[K2 == 0] = np.inf  # Avoid division by zero

    # Compute dipolar interaction kernel in Fourier space
    V_k = (-1 + 3* (polarisation[0]*KX + polarisation[1]*KY + polarisation[2]*KZ)**2 / K2) /3

    # Compute convolution in Fourier space
    V_l_k = V_k * rho_k

    # Transform back to real space
    V_l_real = ifftn(V_l_k)

    # Integrate
    V_lr = (np.sum(V_l_real*rho_r) * dV) 
    #multiply by prefactor C_dd (4pi already in kernel)
    V_lr *= const.mu_0*mu**2

    #check imaginary part of result (should vanish)
    if abs(V_lr.real/V_lr.imag) < 10:
        raise Exception(f"imaginary part of V_lr too big: V_lr = {V_lr}")
    else:
        V_lr = V_lr.real

    return V_lr

def get_DIT_s(trap, GS_left, GS_right, dvol):
    """
    Returns the density-induced tunneling from contact interactions between
    the two tweezer groundstates.
        in: trap (object): trap object for parameters
            GS_left,GS_right (3d array): ground states of both tweezers
            dvol ((3,) array): grid spacing in all 3 directions
        out: DeltaJ_s (float): DIT energy in SI units
    """
    #get paramters
    a_s = float(trap.subs(trap.a_s))
    mass = float(trap.subs(trap.m))
    #calculate volume element
    dV = np.prod(dvol)
    #make sure we're working with state vector
    GS_left /= np.sum(np.abs(GS_left)**2)
    GS_right /= np.sum(np.abs(GS_right)**2)

    #calculate integral (divide by sqrt(dV) to convert to wavefunction)
    DeltaJ_s = np.sum(np.abs(GS_left)**2 *GS_left.conj() *GS_right /dV**2)*dV
    DeltaJ_s *= -4*np.pi*const.hbar**2*a_s/mass

    return DeltaJ_s

def get_DIT_dd(trap, GS_left, GS_right, dvol):
    """
    Returns the density-induced tunneling due to dipole-dipole interactions
    between the two GSs.
        in: trap (object): trap object for parameters
            GS_left,GS_right (3d array): ground states of both tweezers
            dvol ((3,) array): grid spacing in all 3 directions
        out: DeltaJ_lr (float): DIT energy in SI units
    """
    #get paramters
    mu = float(trap.subs(trap.mu_b))
    #get and normalise polarisation
    B_x = float(trap.subs(trap.B_x))
    B_y = float(trap.subs(trap.B_y))
    B_z = float(trap.subs(trap.B_z))
    polarisation = np.array([B_x, B_y, B_z])
    polarisation /= np.sum(np.abs(polarisation)**2)
    #calculate volume element
    dV = np.prod(dvol)

    #make sure we're working with state vector
    GS_left /= np.sum(np.abs(GS_left)**2)
    GS_right /= np.sum(np.abs(GS_right)**2)

    #compute dipolar interactions using fourier transform

    #density
    rho_l = np.abs(GS_left/np.sqrt(dV))**2
    # Compute Fourier transform of density
    rho_k = fftn(rho_l)

    #find shape of grid
    size = GS_left.shape
    #set up fourier space
    kx = 2*np.pi*fftfreq(size[0], dvol[0])
    ky = 2*np.pi*fftfreq(size[1], dvol[1])
    kz = 2*np.pi*fftfreq(size[2], dvol[2])
    KX, KY, KZ = np.meshgrid(kx, ky, kz, indexing='ij')
    K2 = KX**2 + KY**2 + KZ**2
    K2[K2 == 0] = np.inf  # Avoid division by zero

    # Compute dipolar interaction kernel in Fourier space
    V_k = (-1 + 3* (polarisation[0]*KX + polarisation[1]*KY + polarisation[2]*KZ)**2 / K2) /3

    # Compute convolution in Fourier space
    DeltaJ_l_k = V_k * rho_k

    # Transform back to real space
    DeltaJ_l_real = ifftn(DeltaJ_l_k)

    # Integrate
    DeltaJ_lr = (np.sum(DeltaJ_l_real* GS_left.conj()*GS_right /dV) * dV) 
    #multiply by prefactor C_dd (4pi already in kernel)
    DeltaJ_lr *= -const.mu_0*mu**2

    #check imaginary part of result (should vanish)
    if abs(DeltaJ_lr.real/DeltaJ_lr.imag) < 10:
        raise Exception(f"imaginary part of V_lr too big: V_lr = {DeltaJ_lr}")
    else:
        DeltaJ_lr = DeltaJ_lr.real

    return DeltaJ_lr

def get_DIT(trap, GS_left, GS_right, dvol):
    """
    Returns the total DIT term for two groundstates.
        in: trap (object): trap object for parameters
            GS_left,GS_right (3d array): ground states of both tweezers
            dvol ((3,) array): grid spacing in all 3 directions
        out: DeltaJ (float): DIT energy in SI units
    """
    DeltaJ_s = get_DIT_s(trap, GS_left, GS_right, dvol)
    DeltaJ_dd = get_DIT_dd(trap, GS_left, GS_right, dvol)

    return DeltaJ_s + DeltaJ_dd


def analyse_diagonalisation(line_number,
                            n_angles=50,
                            state_nr0=0, state_nr1=1,
                            degenerate=True):
    """
    For a given line number in overview.txt,
    returns J, U_s, U_dd of the two groundstates of both tweezers
    out:
        J, U_s (float)
        U_dds, angles ((n_angles,) arrays): arrays of angles(rad) of B-field and cor. U_dd
    """
    #import results
    trap, ham, res = import_results(line_number)

    #find groundstates
    GS_left, GS_right = get_localised_GS(res["states"][state_nr0]
                                         ,res["states"][state_nr1],
                                         degenerate=degenerate)

    #calculate hubbard parameters
    J = get_J(GS_left, GS_right, ham, res["dvol"])
    U_s = get_U_s(trap, GS_left, res["dvol"])

    angles= np.deg2rad(np.linspace(0,90))
    U_dds = np.zeros_like(angles)
    V_lrs = np.zeros_like(angles)
    DeltaJs = np.zeros_like(angles)
    polarisations = np.zeros((len(angles),3))
    polarisations[:,0] = np.sin(angles)
    polarisations[:,2] = np.cos(angles)

    for i, pol in enumerate(polarisations):
        trap[trap.B_x] = pol[0]
        trap[trap.B_y] = pol[1]
        trap[trap.B_z] = pol[2]

        U_dds[i] = get_U_dd(trap, GS_left, res["dvol"])
        V_lrs[i] = get_NNI(trap, GS_left, GS_right, res["dvol"])
        DeltaJs[i] = get_DIT_dd(trap, GS_left, GS_right, res["dvol"])

    return J, U_s, U_dds, angles, V_lrs, DeltaJs