import kwant
import tbmodels

import numpy as np
import scipy.linalg as la
import matplotlib.pyplot as plt
%matplotlib inline
%config InlineBackend.figure_format = 'retina'

"""
Explanation: Interface from TBmodels to kwant
Prerequisites:
* TBmodels
* kwant
End of explanation
"""


model = tbmodels.Model.from_wannier_files(hr_file='data/wannier90_hr.dat')

"""
Explanation: Bulk Hamiltonian with wraparound
In this first example, we compare the bulk Hamiltonian from TBmodels with that of the model in kwant, using wraparound.
End of explanation
"""


lattice = model.to_kwant_lattice()

sym = kwant.TranslationalSymmetry(
    lattice.vec((1, 0, 0)),
    lattice.vec((0, 1, 0)),
    lattice.vec((0, 0, 1))
)

"""
Explanation: First we need to create the lattice from the tight-binding model and define the translation symmetries.
End of explanation
"""


kwant_sys = kwant.Builder(sym)

"""
Explanation: Now we define a Builder with these symmetries
End of explanation
"""


kwant_sys[lattice.shape(lambda p: True, (0, 0, 0))] = 0

"""
Explanation: We give the system an "infinite" shape. This needs to be done before adding the hoppings, because on-site energies and hoppings are added only to existing sites.
End of explanation
"""


model.add_hoppings_kwant(kwant_sys)

"""
Explanation: Now we can add the hoppings. This modifies the model in-place.
End of explanation
"""


kwant_model = kwant.wraparound.wraparound(kwant_sys).finalized()

"""
Explanation: Finally, use wraparound to finalize the bulk system:
End of explanation
"""


k_list = [(kx, 0, 0) for kx in np.linspace(0, 1, 100)]
x = range(100)

eigs_tbmodels = [model.eigenval(k) for k in k_list]
eigs_kwant = [la.eigvalsh(
    kwant_model.hamiltonian_submatrix(
        params={key: val for key, val in zip(['k_x', 'k_y', 'k_z'], 2 * np.pi * np.array(k))}
    )
) for k in k_list]

"""
Explanation: To see that the two models are the same, we plot the bands along some line. Note that the periodicity of the k-vector is $1$ in TBmodels, but $2\pi$ in kwant. The k-vector needs to be scaled accordingly.
End of explanation
"""


np.isclose(eigs_tbmodels, eigs_kwant).all()

fig, ax = plt.subplots()
for band in np.array(eigs_tbmodels).T:
    ax.plot(x, band, 'k')
for band in np.array(eigs_kwant).T:
    ax.plot(x, band, 'b')

"""
Explanation: Numerical and visual test for equivalence:
End of explanation
"""


wire = kwant.Builder()

"""
Explanation: Finite wire with leads
In the second example, we build a finite wire and attach two leads on either side.
Since the finite wire doesn't have translation symmetry, we can just create a bare Builder.
End of explanation
"""


def shape(p):
    x, y, z = p
    return -20 < x < 20 and -5 < y < 5 and -5 < z < 5

"""
Explanation: Now we define a shape for the wire - for simplicity we use a square.
End of explanation
"""


wire[lattice.shape(shape, (0, 0, 0))] = 0
model.add_hoppings_kwant(wire)

kwant.plot(wire);

"""
Explanation: Again, we explicitly create the lattice sites before populating the hoppings.
End of explanation
"""


sym_lead = kwant.TranslationalSymmetry(lattice.vec((-5, 0, 0)))
lead = kwant.Builder(sym_lead)
def lead_shape(p):
    x, y, z = p
    return -5 <= x <= 0 and -5 < y < 5 and -5 < z < 5
lead[lattice.shape(lead_shape, (0, 0, 0))] = 0
model.add_hoppings_kwant(lead)

wire.attach_lead(lead);
wire.attach_lead(lead.reversed());

kwant.plot(wire);

"""
Explanation: Now we create and attach two leads on either side. The lead must be long enough s.t. the most long-range hopping stays within the lead.
End of explanation
"""


for i, dir in enumerate(['x', 'y', 'z']):
    print(dir + ':', max([abs(R[i]) for R in model.hop.keys()]))

"""
Explanation: Here's how you can find out what's the longest-range hopping in a given direction:
End of explanation
"""


