import numpy as np
import matplotlib.pyplot as plt
import scipy
import scipy.linalg
import tqdm

import bm

# 1.38 degree twist angle with 0.2% biaxial heterostrain.
ps = bm.Parameters(θ=np.radians(1.38), ϵ_bi_hetero=0.002)
L, H = bm.prepare(ps)

# L contains lattice information that we can use to traverse the BZ.
N = 101
ks = np.concatenate((
    np.linspace(L.Kt, L.Kb, N),
    np.linspace(L.Kb, L.Γ, N),
    np.linspace(L.Γ, L.M, N),
    np.linspace(L.M, L.Kt, N),
))

eigs = np.empty((len(ks), len(H(0))))
for i, k in enumerate(tqdm.tqdm(ks)):
    eigs[i, :] = np.sort(scipy.linalg.eigvalsh(H(k)))

plt.figure(figsize=(6, 4))
for i in range(5):
    plt.axvline(i * N, alpha=0.2, ls='--')
plt.axhline(0, alpha=0.2, ls='--')
for i in range(eigs.shape[1]):
    plt.plot(eigs[:, i], 'k')
plt.xticks([0, N, 2 * N, 3 * N, 4 * N], ['K\'', 'K', 'Γ', 'M', 'K\''])
plt.ylim(-220, 220)
plt.ylabel('E (meV)')
plt.tight_layout()
plt.show()