import numpy as np


def tensor_product(*args):
    if len(args) < 2:
        return args[0]

    result = args[0]
    for i in range(1, len(args)):
        result = np.kron(result, args[i])
    return result

Sz = np.array([[1, 0], [0, -1]])
Sp = np.array([[0, 1], [0, 0]])
Sm = np.array([[0, 0], [1, 0]])
II = np.eye(2)

H = (
    1 / 2 * (tensor_product(Sp, Sm, II, II) + tensor_product(Sm, Sp, II, II))
    + tensor_product(Sz, Sz, II, II)
    + 1 / 2 * (tensor_product(II, Sp, Sm, II) + tensor_product(II, Sm, Sp, II))
    + tensor_product(II, Sz, Sz, II)
    + 1 / 2 * (tensor_product(II, II, Sp, Sm) + tensor_product(II, II, Sm, Sp))
    + tensor_product(II, II, Sz, Sz)
    + 1 / 2 * (tensor_product(Sm, II, II, Sp) + tensor_product(Sp, II, II, Sm))
    + tensor_product(Sz, II, II, Sz)
)
H111 = H

from scipy.sparse.linalg import eigsh

eigvals, eigvecs = eigsh(H, k=1, which="SA")
eigvals

def get_state_index(state: int, index: int) -> int:
    """获得一个整数某位置处的二进制值"""
    mask = 1 << index
    return (state & mask) >> index


def flip_state(state: int, index: int) -> int:
    """翻转一个整数某位置处的二进制值"""
    mask = 1 << index
    return state ^ mask

N = 4
H = np.zeros((2 ** N, 2 ** N))

# a is a basis
for a in range(2 ** N):
    # i is position of this basis
    for i in range(N):
        # j is the nearest neighbor, mod N for periodic boundary conditions
        j = (i + 1) % N
        ai = get_state_index(a, i)
        aj = get_state_index(a, j)

        # Sz
        b = a
        if ai == aj:
            H[a, b] += 1
        else:
            H[a, b] += -1

        # SxSx + SySy
        if ai != aj:
            b = flip_state(a, i)
            b = flip_state(b, j)
            H[a, b] += 1 / 2

H112 = H

eigvals, eigvecs = eigsh(H, k=1, which="SA")
eigvals, np.allclose(H112 - H111, 0)

N = 4
# 这里的总自旋写2是因为总的因子 1/2 被省略了
Nz = 2

state_list = []
for a in range(2 ** N):
    if bin(a).count("1") == Nz:
        state_list.append(a)

state_list

from itertools import combinations

state_list = [sum([1 << i for i in c]) for c in combinations(range(N), Nz)]
state_list.sort()
state_list

H = np.zeros((len(state_list),) * 2)

for new_a, a in enumerate(state_list):
    for i in range(N):
        j = (i + 1) % N
        ai = get_state_index(a, i)
        aj = get_state_index(a, j)

        # Sz
        new_b = new_a
        if ai == aj:
            H[new_a, new_b] += 1
        else:
            H[new_a, new_b] += -1

        # SxSx + SySy
        if ai != aj:
            b = flip_state(a, i)
            b = flip_state(b, j)
            new_b = state_list.index(b)
            H[new_a, new_b] += 1 / 2

eigvals, eigvecs = eigsh(H, k=1, which="SA")
eigvals

def state_cycle_shift(state: int, shift: int, total: int) -> int:
    """循环移位操作"""
    low = state >> (total - shift)
    mask = 2 ** (total - shift) - 1
    high = mask & state
    return (high << shift) + low

state_list_temp = state_list.copy()
state_k_all = []

while len(state_list_temp) != 0:
    state_k_list = []
    s_first = state_list_temp[0]
    s = s_first
    while True:
        state_k_list.append(s)
        state_list_temp.remove(s)
        s_next = state_cycle_shift(s, 1, N)
        if s_next == s_first:
            break
        else:
            s = s_next

    state_k_list.sort()
    state_k_all.append(state_k_list)

state_k_all

k = 0
Hk = np.zeros((len(state_k_all),) * 2).astype(complex)

for new_b, b_list in enumerate(state_k_all):

    # build state_b
    state_new_b = np.zeros((1, len(state_list))).astype(complex)
    for rb, b in enumerate(b_list):
        ib = state_list.index(b)
        state_new_b[0, ib] = 1
    state_new_b = state_new_b / np.linalg.norm(state_new_b)

    for new_a, a_list in enumerate(state_k_all):
        state_new_a = np.zeros((len(state_list), 1)).astype(complex)

        # build state_a
        state_new_a_temp = np.zeros((len(state_list), 1)).astype(complex)
        for ra, a in enumerate(a_list):
            ia = state_list.index(a)
            state_new_a_temp[ia, 0] = np.exp(1j * 2 * np.pi / N * k * ra)
        state_new_a_temp = state_new_a_temp / np.linalg.norm(state_new_a_temp)

        # H @ state_a
        for ra, a in enumerate(a_list):
            ia = state_list.index(a)
            for i in range(N):
                j = (i + 1) % N
                ai = get_state_index(a, i)
                aj = get_state_index(a, j)

                # Sz
                new_ia = ia
                if ai == aj:
                    state_new_a[new_ia, 0] += state_new_a_temp[ia, 0]

                else:
                    state_new_a[new_ia, 0] -= state_new_a_temp[ia, 0]

                # SxSx + SySy
                if ai != aj:
                    b = flip_state(a, i)
                    b = flip_state(b, j)
                    new_ib = state_list.index(b)
                    state_new_a[new_ib, 0] += 1 / 2 * state_new_a_temp[ia, 0]
        # Hk
        Hk[new_a, new_b] = state_new_b @ state_new_a

np.linalg.eigvalsh(Hk)

from typing import Optional, Tuple


def check_U(state_up: int, state_down: int, i: int) -> Optional[Tuple[int, int]]:
    state_up_i = get_state_index(state_up, i)
    state_down_i = get_state_index(state_down, i)
    if state_up_i == state_down_i == 1:
        return state_up, state_down
    else:
        return None


def check_t(state: int, i: int, N: int) -> Optional[Tuple[int, int]]:
    j = (i + 1) % N
    state_i = get_state_index(state, i)
    state_j = get_state_index(state, j)
    if (state_i, state_j) == (0, 1) or (state_i, state_j) == (1, 0):
        state_new = flip_state(state, i)
        state_new = flip_state(state_new, j)
        return state_new
    else:
        return None

N = 4
t = 1
U = 1
H = np.zeros((4 ** N, 4 ** N))
al = []
for a_up in range(2 ** N):
    for a_down in range(2 ** N):
        a = (a_up << N) + a_down
        al.append(a)
        for i in range(N):
            # t
            b_up = check_t(a_up, i, N)
            if b_up is not None:
                b = (b_up << N) + a_down
                H[a, b] += -t

            b_down = check_t(a_down, i, N)
            if b_down is not None:
                b = (a_up << N) + b_down
                H[a, b] += -t

            # U
            b_result = check_U(a_up, a_down, i)
            if b_result is not None:
                b = (b_result[0] << N) + b_result[1]
                H[a, b] += U

# eigvals, eigvecs = eigsh(H, k=1, which='SA')
# eigvals
eigvals_all = np.linalg.eigvalsh(H)

N = 4
Nup = 2
Ndown = 2
t = 1
U = 1

from itertools import combinations

state_up_list = [sum([1 << i for i in c]) for c in combinations(range(N), Nup)]
state_up_list.sort()
state_down_list = [sum([1 << i for i in c]) for c in combinations(range(N), Ndown)]
state_down_list.sort()
state_total_list = [(i, j) for i in state_up_list for j in state_down_list]
state_up_list, state_down_list

state_list_temp = state_total_list.copy()
state_k_all = []

while len(state_list_temp) != 0:
    state_k_list = []
    s_first = state_list_temp[0]
    s = s_first
    while True:
        state_k_list.append(s)
        state_list_temp.remove(s)

        s0_next = state_cycle_shift(s[0], 1, N)
        s1_next = state_cycle_shift(s[1], 1, N)
        s_next = (s0_next, s1_next)

        if s_next == s_first:
            break
        else:
            s = s_next

    state_k_list.sort()
    state_k_all.append(state_k_list)

# state_k_all

k = 0
Hk = np.zeros((len(state_k_all),) * 2).astype(complex)

for new_b, b_list in enumerate(state_k_all):

    # build state_b
    state_new_b = np.zeros((1, len(state_total_list))).astype(complex)
    for rb, b in enumerate(b_list):
        ib = state_total_list.index(b)
        state_new_b[0, ib] = 1
    state_new_b = state_new_b / np.linalg.norm(state_new_b)

    for new_a, a_list in enumerate(state_k_all):
        state_new_a = np.zeros((len(state_total_list), 1)).astype(complex)

        # build state_a
        state_new_a_temp = np.zeros((len(state_total_list), 1)).astype(complex)
        for ra, a in enumerate(a_list):
            ia = state_total_list.index(a)
            state_new_a_temp[ia, 0] = np.exp(1j * 2 * np.pi / N * k * ra)
        state_new_a_temp = state_new_a_temp / np.linalg.norm(state_new_a_temp)

        # H @ state_a
        for ra, a in enumerate(a_list):
            a_up = a[0]
            a_down = a[1]
            ia = state_total_list.index(a)
            for i in range(N):

                # t
                b_up = check_t(a_up, i, N)
                if b_up is not None:
                    b = (b_up, a_down)
                    new_ib = state_total_list.index(b)
                    state_new_a[new_ib, 0] += -t * state_new_a_temp[ia, 0]

                b_down = check_t(a_down, i, N)
                if b_down is not None:
                    b = (a_up, b_down)
                    new_ib = state_total_list.index(b)
                    state_new_a[new_ib, 0] += -t * state_new_a_temp[ia, 0]

                # U
                b_result = check_U(a_up, a_down, i)
                if b_result is not None:
                    b = b_result
                    new_ib = state_total_list.index(b)
                    state_new_a[new_ib, 0] += U * state_new_a_temp[ia, 0]
        # Hk
        Hk[new_a, new_b] = state_new_b @ state_new_a

# Hk

eigvals = np.linalg.eigvalsh(Hk)

for x in eigvals:
    index = np.abs(eigvals_all - x).argmin()
    print(np.isclose(eigvals_all[index], x), eigvals_all[index], x)

import numpy as np

def proj(u, v, do_norm=False):
    if do_norm:
        result = np.vdot(u,v) / np.linalg.norm(u,u) * u
    else:
        result = np.vdot(u,v) * u
    return result.flatten()

def gram_schmidt(A):
    b = A.copy()
    b[:,0] = b[:,0] / np.linalg.norm(b[:,0])
    for i in range(1, A.shape[1]):
        b[:,i] = (b[:,i] 
                  - np.sum(np.apply_along_axis(
                      lambda u: proj(u, b[:,i]), 
                      0, 
                      b[:,:i]),
                           axis=1).flatten())
        b[:,i] = b[:,i] / np.linalg.norm(b[:,i])
    return b

a = np.random.rand(4,4)
b = gram_schmidt(a)

from itertools import combinations
for i,j in combinations(range(4),2):
    print(np.isclose(b[:,i] @ b[:,j], 0))

import numpy as np

def power_iteration(A, num_simulations: int, v0=None):
    if v0 is None:
        v = np.random.rand(A.shape[1])
    else: 
        v = v0
        
    for _ in range(num_simulations):     
        v = A @ v / np.max(v)
    return v

a = np.random.rand(4,4)
e2,v2=np.linalg.eig(a)
ei2 = e2.argmax()
e2[ei2], v2[:,ei2] / np.linalg.norm(v2[:,ei2])

v1 = power_iteration(a, 1000)
(a @ v1 / v1)[0], v1/ np.linalg.norm(v1)

def inverse_power_iteration(A, num_simulations: int, v0=None):
    if v0 is None:
        v = np.random.rand(A.shape[1])
    else: 
        v = v0
        
    for _ in range(num_simulations):  
        v = np.linalg.solve(A, v) 
        # v = np.linalg.inv(A) @ v
        v = v / np.max(v)
    return v
    

a = np.random.rand(4,4)
e2,v2=np.linalg.eig(a)
ei2 = np.abs(e2).argmin()
e2[ei2], v2[:,ei2] / np.linalg.norm(v2[:,ei2])

v1 = inverse_power_iteration(a, 1000)
1/ (np.linalg.inv(a) @ v1 / v1)[0], v1/ np.linalg.norm(v1)