import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches

a0 = 1.42 # [A]
a1 = np.array([a0/2*3, a0/2*np.sqrt(3)])
a2 = np.array([a0/2*3, -a0/2*np.sqrt(3)])
b1 = np.array([2*np.pi/(3*a0), 2*np.pi/(3*a0)*np.sqrt(3)])
b2 = np.array([2*np.pi/(3*a0), -2*np.pi/(3*a0)*np.sqrt(3)])
K = np.array([2*np.pi/(3*a0), 2*np.pi/(3*a0)/np.sqrt(3)])
Kp = np.array([2*np.pi/(3*a0), -2*np.pi/(3*a0)/np.sqrt(3)])
K1 = np.array([0, 4*np.pi/(3*np.sqrt(3)*a0)])
K2 = np.array([0, -4*np.pi/(3*np.sqrt(3)*a0)])
K3 = np.array([-2*np.pi/(3*a0), 2*np.pi/(3*a0)/np.sqrt(3)])
K4 = np.array([-2*np.pi/(3*a0), -2*np.pi/(3*a0)/np.sqrt(3)])
t = 2.8 #[meV]


K = np.array([2*np.pi/(3*a0), 2*np.pi/(3*a0)/np.sqrt(3)])
M = np.array([2*np.pi/(3*a0), 0])
Γ = np.array([0, 0])

# Define the number of points for each segment
num_points_per_segment = int(2000 / 3)
t0 = 0.35
# Create the k-path segments

# K to M
k_path_KΓ = np.linspace(K, Γ, num_points_per_segment)

# M to Γ
k_path_ΓM = np.linspace(Γ, M, num_points_per_segment)

# Γ to K (avoiding duplicate point)
k_path_MK = np.linspace(M, K, num_points_per_segment)  # Exclude the last point of K

# Combine the segments
k_path = np.concatenate(( k_path_ΓM, k_path_MK))

# Print the k_path (optional)
#print(k_path)
energy_list = []
k_list = []
for k in k_path:
    kx = k[0]
    ky = k[1]
    f = np.exp(-1j*kx*a0)+ np.exp(1j*kx*a0/2)* np.exp(1j*ky*np.sqrt(3)*a0/2) + np.exp(1j*kx*a0/2)* np.exp(-1j*ky*np.sqrt(3)*a0/2)
    f_T = np.conj(f)
    Hk = np.array([ 
        [0,t0,-t*f,0],
        [t0,0,0,-t*f],
        [-t*f_T,0,0,t0],
        [0,-t*f_T,t0,0]
    ])
    energy_list.append(np.linalg.eigvalsh(Hk))
    
    k_list.append(np.linalg.norm([kx, ky]))
    eigenvalues = np.array(energy_list)

energy_list2 = []
k_list2 = []
for k in k_path_KΓ:
    kx = k[0]
    ky = k[1]
    f = np.exp(-1j*kx*a0)+ np.exp(1j*kx*a0/2)* np.exp(1j*ky*np.sqrt(3)*a0/2) + np.exp(1j*kx*a0/2)* np.exp(-1j*ky*np.sqrt(3)*a0/2)
    f_T = np.conj(f)
    Hk = np.array([ 
        [0,t0,-t*f,0],
        [t0,0,0,-t*f],
        [-t*f_T,0,0,t0],
        [0,-t*f_T,t0,0]
    ])
    energy_list2.append(np.linalg.eigvalsh(Hk))
    
    k_list2.append(-np.linalg.norm([kx, ky]))
    eigenvalues2 = np.array(energy_list2)
# Plot the band structure
plt.figure(figsize=(8, 6))
plt.plot(k_list, eigenvalues[:, 0], label='Band 1', color='b')
plt.plot(k_list, eigenvalues[:, 1], label='Band 2', color='r')
plt.plot(k_list, eigenvalues[:, 2], label='Band 3', color='g')
plt.plot(k_list, eigenvalues[:, 3], label='Band 4', color='orange')
plt.plot(k_list2, eigenvalues2[:, 0], label='Band 1', color='b')
plt.plot(k_list2, eigenvalues2[:, 1], label='Band 2', color='r')
plt.plot(k_list2, eigenvalues2[:, 2], label='Band 3', color='g')
plt.plot(k_list2, eigenvalues2[:, 3], label='Band 4', color='orange')