# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_Winding_num.ipynb (unless otherwise specified).

__all__ = ['Fukui_Kitaev_LR_QP_wn', 'Fukui_Kitaev_LR_QP_inf_wn', 'Real_space_Kitaev_LR_QP_wn', 'w',
           'Chiral_Kitaev_QP_LR_wn', 'd_num_h_Kitaev_LR_QP_inf', 'd_num_h_Kitaev_LR_QP_inf', 'Fukui_Kitaev_AA_2D_chern',
           'Fukui_Kitaev_AA_2D_chern_inf', 'd_num_k_H_Kitaev_LR_QP', 'd_num_phase_H_Kitaev_LR_QP',
           'compute_Kitaev_AA_wn_TKNN', 'd_num_k_H_Kitaev_LR_QP_inf', 'd_num_phase_H_Kitaev_LR_QP_inf',
           'compute_Kitaev_AA_wn_TKNN_inf']

# Cell

from .Hamiltonians import *
import numpy as np
from mpmath import *

# Cell

def Fukui_Kitaev_LR_QP_wn(params, k_l, mu, L, rot = True, close_loop = True):

    '''Compute the winding number using Fukui for the finite AAH-LR Kitaev chain.'''

    Fn = params['Fn']
    U = np.zeros(len(k_l), dtype=complex)
    F = 0

    if close_loop == False:
        k_l_range = len(k_l)-1
    else:
        k_l_range = len(k_l)

    for i in range(k_l_range):

        evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP(params, k_l[i], mu, L, rot))
        evecs1 = evecs[:,:Fn]

        if i == range(len(k_l))[-1]:
            k = k_l[0]
        else:
            k = k_l[i+1]

        evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP(params, k, mu, L, rot))
        evecs2 = evecs[:,:Fn]

        Ui=np.dot(np.conjugate(np.transpose(evecs2)),evecs1)
        Ui=np.linalg.det(Ui)

        #Ui=Ui/np.abs(Ui)
        U[i]=Ui
    F = np.imag(np.log(np.prod(U)))

    return F, U


def Fukui_Kitaev_LR_QP_inf_wn(params, k_l, mu, rot = True, close_loop = True):

    '''Compute the winding number using Fukui for the infinite AAH-LR Kitaev chain.'''

    Fn = params['Fn']
    U = np.zeros(len(k_l), dtype=complex)
    F = 0

    if close_loop == False:
        k_l_range = len(k_l)-1
    else:
        k_l_range = len(k_l)

    for i in range(k_l_range):

        evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP_inf(params, k_l[i], mu, rot))
        evecs1 = evecs[:,:Fn]

        if i == range(len(k_l))[-1]:
            k = k_l[0]
        else:
            k = k_l[i+1]

        evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP_inf(params, k, mu, rot))
        evecs2 = evecs[:,:Fn]

        Ui=np.dot(np.conjugate(np.transpose(evecs2)),evecs1)
        Ui=np.linalg.det(Ui)

        #Ui=Ui/np.abs(Ui)
        U[i]=Ui
    F = np.imag(np.log(np.prod(U)))

    return F, U

# Cell

def Real_space_Kitaev_LR_QP_wn(params, mu, length, AA=False):

    '''Compute the real space winding number for the QP LR Kitaev chain.
    We use the real space Hamiltonian with APBC.'''

    Fn = params['Fn']
    m = length*Fn
    n = int((m-1)/2)

    mat = H_APBC_Kitaev_LR_QP(params, mu, length, rot=True, AA=AA)

    val, vec = np.linalg.eigh(mat)
    vec = vec[:,val<0]
    p = vec@np.conjugate(vec.T)
    q = np.identity(4*n+2) - 2 * p

    ga, gb = np.ones(4*n+2), np.ones(4*n+2)
    ga[1::2], gb[0::2] = 0, 0
    ga, gb = np.diag(ga), np.diag(gb)
    qab, qba = ga@q@gb, gb@q@ga

    ll = np.arange(-n, n+1)
    ll = ll[ll != 0]
    ll = np.exp(2*np.pi*1j/(2*n+1)*ll)
    cl = ll**(n+1)/(1-ll)

    x = np.repeat(np.arange(-n,n+1),2)
    mat= np.zeros((4*n+2,4*n+2))
    for i, l in enumerate(ll):
        mat = mat + cl[i] * np.diag(l**x) @ qab @ np.diag(l**(-x))

    w = qba@mat
    w = np.sum(np.diag(w))/(2*n+1)

    return -w.real

# Cell

def w(params, k, mu, L, num, inf, AA, H_pbc_delta, d_H_pbc_delta):

    '''Compute the relative phase for a given value of k using the definition.
    We can choose whether to use the numerical or analytical derivative for the infinite system.'''

    if inf == True:
        h = h_chiral_Kitaev_LR_QP_inf(params, k, mu, AA, H_pbc_delta)
    else:
        h = h_chiral_Kitaev_LR_QP(params, k, mu, L, AA)

    if num == True:
        if inf == True:
            d_h = d_num_h_Kitaev_LR_QP_inf(params, k, mu, AA)
        else:
            d_h = d_num_h_Kitaev_LR_QP(params, k, mu, L, AA)
    else:
        if inf == True:
            d_h = d_h_chiral_Kitaev_LR_QP_inf(params, k, mu, AA,  d_H_pbc_delta)
        else:
            print("There is no analytic derivative for the finite system.")

    w = np.trace(np.linalg.inv(h)@d_h)

    return w


def Chiral_Kitaev_QP_LR_wn(params, vec_mu, veck, inf=True, L=0, num=False, AA=True):

    ''' Compute the winding number using the definition for Hamiltonians
    with Chiral symmetry given a discretization. Both for the infinite and finite system.'''

    vec_w = np.zeros_like(vec_mu)
    Fn = params['Fn']

    arr_H_pbc_delta = arr_d_H_pbc_delta = []
    if inf == True:
        for k in veck:
            H_pbc_delta = H_pbc_sp(params, k)
            arr_H_pbc_delta = np.append(arr_H_pbc_delta, H_pbc_delta).reshape(len(arr_H_pbc_delta)+1, Fn, Fn)
            if num == False:
                d_H_pbc_delta = d_H_pbc_sp(params, k)
                arr_d_H_pbc_delta = np.append(arr_d_H_pbc_delta, d_H_pbc_delta).reshape(len(arr_d_H_pbc_delta)+1, Fn, Fn)

    for i, mu in enumerate(vec_mu):
        vecw = []
        for j, k in enumerate(veck):

            if inf == True and num == False:
                H_pbc_delta = arr_H_pbc_delta[j]
                d_H_pbc_delta = arr_d_H_pbc_delta[j]
                vecw = np.append(vecw, w(params, k, mu, num, L, inf, AA, H_pbc_delta, d_H_pbc_delta))

            elif inf == True and num == True:
                H_pbc_delta = arr_H_pbc_delta[j]
                vecw = np.append(vecw, w(params, k, mu, num, L, inf, AA, H_pbc_delta, 0))

            else:
                vecw = np.append(vecw, w(params, k, mu, num, L, inf, AA, 0, 0))

        vec_w[i] = (np.sum(vecw)/len(veck)).imag

    return vec_w

# Cell

def d_num_h_Kitaev_LR_QP_inf(params, k, mu, delta_k = 0.01, AA=False):

    '''Numerical derivative of one block of the off-diagonal Hamiltonian in momentum
    space for the infinite system. We use the forward method with a small delta.'''

    h_k = h_chiral_Kitaev_LR_QP_inf(params, k, mu, AA)
    h_k_delta_plus = h_chiral_Kitaev_LR_QP_inf(params, k+delta_k, mu, AA)

    return (h_k_delta_plus-h_k)/delta_k


def d_num_h_Kitaev_LR_QP_inf(params, k, mu, L, AA):

    '''Numerical derivative of one block of the off-diagonal Hamiltonian in momentum
    space for the finite system. '''

    delta_k = 2*np.pi/L
    h_k = h_chiral_Kitaev_LR_QP(params, k, mu, L, AA)
    h_k_delta_plus = h_chiral_Kitaev_LR_QP(params, k+delta_k, mu, L, AA)

    return (h_k_delta_plus-h_k)/delta_k


# Cell

def Fukui_Kitaev_AA_2D_chern(params, mu, L, list_kx, list_ky, energy_gap, rot=True):

    '''Compute the Chern number using Fukui for the AAH modes of the LR AAH Kitaev chain (finite).
    Here kx is the momentum and ky is the phase of the AAH chemical potential.'''

    Fn = params['Fn']

    delta_kx = 2*np.pi/L
    delta_ky = 2*np.pi/L
    deltas = [np.array([delta_kx, 0]), np.array([0, delta_ky]), np.array([delta_kx, delta_ky])]

    sum_F = []
    F_list = np.zeros((len(list_kx), len(list_ky)), dtype=complex)
    zero = np.array([0,0], dtype=complex)

    for i in range(len(list_kx)):
        kx = list_kx[i]

        for j in range(len(list_ky)):

            ky = list_ky[j]
            k = np.array([kx,ky])

            U = []
            for delta1, delta2 in zip([zero,deltas[0],deltas[2],deltas[1]],[deltas[0],deltas[2], deltas[1], zero]):

                k1 = k + delta1
                params['phase'] = k1[1]
                evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP(params, k1[0], mu, L, rot))
                evecs = evecs.T
                evecs1 = evecs[evals<energy_gap]

                k2 = k + delta2
                params['phase'] = k2[1]
                evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP(params, k2[0], mu, L, rot))
                evecs = evecs.T
                evecs2 = evecs[evals<energy_gap]

                Ui = np.dot(evecs1,np.conjugate(evecs2.T))
                if len(U) == 0:
                    U = np.identity(len(Ui),  dtype=complex)
                    U = U @ Ui
                else:
                    U = U @ Ui

            F_k = np.linalg.det(U)
            F_list[i][j] = np.log(F_k).imag
            sum_F = np.append(sum_F, np.log(F_k).imag)

    F = np.sum(sum_F)/(2.*np.pi)

    return F, F_list


def Fukui_Kitaev_AA_2D_chern_inf(params, mu, delta_kx, delta_ky, energy_gap, rot = True):

    '''Compute the Chern number using Fukui for the AAH modes of the LR AAH Kitaev chain (infinite).
    Here kx is the momentum and ky is the phase of the AAH chemical potential.'''

    Fn = params['Fn']

    list_ky = np.arange(delta_ky, 2*np.pi, delta_ky)
    list_kx = np.arange(delta_kx, 2*np.pi, delta_kx)
    deltas = [np.array([delta_kx, 0]), np.array([0, delta_ky]), np.array([delta_kx, delta_ky])]

    sum_F = []
    F_list = np.zeros((len(list_kx), len(list_ky)), dtype=complex)
    zero = np.array([0,0], dtype=complex)

    for i in range(len(list_kx)):
        kx = list_kx[i]

        for j in range(len(list_ky)):

            ky = list_ky[j]
            k = np.array([kx,ky])

            U = []
            for delta1, delta2 in zip([zero,deltas[0],deltas[2],deltas[1]],[deltas[0],deltas[2], deltas[1], zero]):

                k1 = k + delta1
                params['phase'] = k1[1]
                evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP_inf(params, k1[0], mu, rot))
                evecs = evecs.T
                evecs1 = evecs[evals<energy_gap]

                k2 = k + delta2
                params['phase'] = k2[1]
                evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP_inf(params, k2[0], mu, rot))
                evecs = evecs.T
                evecs2 = evecs[evals<energy_gap]

                Ui = np.dot(evecs1,np.conjugate(evecs2.T))
                if len(U) == 0:
                    U = np.identity(len(Ui),  dtype=complex)
                    U = U @ Ui
                else:
                    U = U @ Ui

            F_k = np.linalg.det(U)
            F_list[i][j] = np.log(F_k).imag
            sum_F = np.append(sum_F, np.log(F_k).imag)

    F = np.sum(sum_F)/(2.*np.pi)

    return F, F_list

# Cell

def d_num_k_H_Kitaev_LR_QP(params, kx, mu, L, AA):

    ''' Numerical forward derivative of the Hamiltonian in k in momentum space with APBC. '''

    H_k_delta = H_Kitaev_LR_QP(params, kx, mu, L, AA)
    H_k_delta_plus = H_Kitaev_LR_QP(params, kx+2*np.pi/L, mu, L, AA)

    return (H_k_delta_plus-H_k_delta)/(2*np.pi/L)


def d_num_phase_H_Kitaev_LR_QP(params, kx, mu, L, AA):

    ''' Numerical forward derivative of the Hamiltonian in the phase phi in momentum space with APBC. '''

    H_k_delta = H_Kitaev_LR_QP(params, kx, mu, L, AA)
    params['phase'] += 2*np.pi/L
    H_k_delta_plus = H_Kitaev_LR_QP(params, kx, mu, L, AA)

    return (H_k_delta_plus-H_k_delta)/(2*np.pi/L)


def compute_Kitaev_AA_wn_TKNN(params, mu, L, list_kx, list_ky, energy_gap, rot=True, AA=False):

    '''Compute the winding number for the AAH modes of the LR Kitaev chain'''

    Fn = params['Fn']

    delta_kx = 2*np.pi/L
    delta_ky = 2*np.pi/L
    deltas = [np.array([delta_kx, 0]), np.array([0, delta_ky]), np.array([delta_kx, delta_ky])]

    sum_F = 0
    F_list = np.zeros((len(list_kx), len(list_ky)), dtype=complex)
    zero = np.array([0,0], dtype=complex)

    for i in range(len(list_kx)):
        kx = list_kx[i]

        for j in range(len(list_ky)):
            ky = list_ky[j]

            params['phase'] = ky
            evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP(params, kx, mu, L, rot, AA=AA))
            evecs = evecs.T
            evecs_below = evecs[evals<energy_gap]
            evals_below = evals[evals<energy_gap]
            evecs_above = evecs[evals>energy_gap]
            evals_above = evals[evals>energy_gap]

            d_kx = d_num_k_H_Kitaev_LR_QP(params, kx, mu, L, AA)
            d_ky = d_num_phase_H_Kitaev_LR_QP(params, kx, mu, L, AA)

            F = 0
            for m in range(len(evals_below)):
                for n in range(len(evals_above)):
                    F += -(np.conjugate(evecs_below[m]).T@d_kx@evecs_above[n]*
                        np.conjugate(evecs_above[n]).T@d_ky@evecs_below[m]-
                        np.conjugate(evecs_below[m]).T@d_ky@evecs_above[n]*
                        np.conjugate(evecs_above[n]).T@d_kx@evecs_below[m])/(evals_below[m]-evals_above[n])**2

            F_list[i][j] = np.imag(F)/(L/2/np.pi)**2
            sum_F += np.imag(F)/(L/2/np.pi)**2

    return sum_F, F_list

# Cell

def d_num_k_H_Kitaev_LR_QP_inf(params, kx, mu, delta_k = 0.1, AA=False):

    ''' Numerical forward derivative of the Hamiltonian in k in momentum space with APBC. '''

    H_k_delta = H_Kitaev_LR_QP_inf(params, kx, mu, AA=AA)
    H_k_delta_plus = H_Kitaev_LR_QP_inf(params, kx+delta_k, mu, AA=AA)

    return (H_k_delta_plus-H_k_delta)/delta_k


def d_num_phase_H_Kitaev_LR_QP_inf(params, kx, mu, delta_ph = 0.1, AA=False):

    ''' Numerical forward derivative of the Hamiltonian in the phase phi in momentum space with APBC. '''

    H_k_delta = H_Kitaev_LR_QP_inf(params, kx, mu, AA=AA)
    params['phase'] += delta_ph
    H_k_delta_plus = H_Kitaev_LR_QP_inf(params, kx, mu, AA=AA)

    return (H_k_delta_plus-H_k_delta)/delta_ph


def compute_Kitaev_AA_wn_TKNN_inf(params, mu, list_kx, list_ky, energy_gap, num=False, rot=False, AA=True):

    '''Compute the winding number for the AAH modes of the infinite LR Kitaev chain'''

    sum_F = 0
    F_list = np.zeros((len(list_kx), len(list_ky)), dtype=complex)
    zero = np.array([0,0], dtype=complex)

    delta_kx = abs(list_kx[1]-list_kx[0])
    delta_ky = abs(list_ky[1]-list_ky[0])

    for i in range(len(list_kx)):
        kx = list_kx[i]

        for j in range(len(list_ky)):
            ky = list_ky[j]

            params['phase'] = ky
            evals, evecs = np.linalg.eigh(H_Kitaev_LR_QP_inf(params, kx, mu, rot, AA=AA))
            evecs = evecs.T
            evecs_below = evecs[evals<=energy_gap]
            evals_below = evals[evals<=energy_gap]
            evecs_above = evecs[evals>energy_gap]
            evals_above = evals[evals>energy_gap]

            if num == True:
                d_kx = d_num_k_H_Kitaev_LR_QP_inf(params, kx, mu, AA=AA)
                d_ky = d_num_phase_H_Kitaev_LR_QP_inf(params, kx, mu, AA=AA)
            else:
                d_kx = d_k_H_Kitaev_LR_QP_inf(params, kx, mu, rot=rot, AA=AA)
                d_ky = d_phase_H_Kitaev_LR_QP_inf(params, kx, mu, rot=rot, AA=AA)

            F = 0
            for m in range(len(evals_below)):
                for n in range(len(evals_above)):
                    F += -(np.conjugate(evecs_below[m]).T@d_kx@evecs_above[n]*
                        np.conjugate(evecs_above[n]).T@d_ky@evecs_below[m]-
                        np.conjugate(evecs_below[m]).T@d_ky@evecs_above[n]*
                        np.conjugate(evecs_above[n]).T@d_kx@evecs_below[m])/(evals_below[m]-evals_above[n])**2

            F_list[i][j] = np.imag(F)*delta_kx*delta_ky
            sum_F += np.imag(F)*delta_kx*delta_ky

    return sum_F, F_list