#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import scipy.linalg
from matplotlib import pyplot

import kwant
import kwant.continuum


# In[2]:


import matplotlib
import matplotlib.pyplot
from matplotlib_inline.backend_inline import set_matplotlib_formats

matplotlib.rcParams['figure.figsize'] = matplotlib.pyplot.figaspect(1) * 2
set_matplotlib_formats('svg')


# In[3]:


hamiltonian = """
   + C * identity(4) + M * kron(sigma_0, sigma_z)
   - F * (k_x**2 + k_y**2) * kron(sigma_0, sigma_z)
   - D * (k_x**2 + k_y**2) * kron(sigma_0, sigma_0)
   + A * k_x * kron(sigma_z, sigma_x)
   - A * k_y * kron(sigma_0, sigma_y)
"""

syst = kwant.continuum.discretize_landau(hamiltonian, N=10)
syst = syst.finalized()


# In[4]:


params = dict(A=3.645, F =-68.6, D=-51.2, M=-0.01, C=0)
b_values = np.linspace(0.0001, 0.0004, 200)

fig = kwant.plotter.spectrum(syst, ('B', b_values), params=params, show=False)
pyplot.ylim(-0.1, 0.2);


# In[5]:


lat = kwant.lattice.square(norbs=1)
syst = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))

def peierls(to_site, from_site, B):
    y = from_site.tag[1]
    return -1 * np.exp(-1j * B * y)

syst[(lat(0, j) for j in range(-19, 20))] = 4
syst[lat.neighbors()] = -1
syst[kwant.HoppingKind((1, 0), lat)] = peierls
syst = syst.finalized()

landau_syst = kwant.continuum.discretize_landau("k_x**2 + k_y**2", N=5)
landau_syst = landau_syst.finalized()


# In[6]:


fig, ax = pyplot.subplots(1, 1)
ax.set_xlabel("momentum")
ax.set_ylabel("energy")
ax.set_ylim(0, 1)

params = dict(B=0.1)

kwant.plotter.bands(syst, ax=ax, params=params)

h = landau_syst.hamiltonian_submatrix(params=params)
for ev in scipy.linalg.eigvalsh(h):
  ax.axhline(ev, linestyle='--')


# In[7]:


continuum_hamiltonian = """
    (k_x**2 + k_y**2 + k_z**2) * sigma_0
    + V(z) * sigma_0
    + mu * B * sigma_z / 2
    + alpha * (sigma_x * k_y - sigma_y * k_x)
"""

template = kwant.continuum.discretize_landau(continuum_hamiltonian, N=10)


# In[8]:


def hetero_structure(site):
    z, = site.pos
    return 0 <= z < 10

def hetero_potential(z):
    if z < 2:
      return 0
    elif z < 7:
      return 0.5
    else:
      return 0.7

syst = kwant.Builder()
syst.fill(template, hetero_structure, (0,))

syst = syst.finalized()

params = dict(
    B=0.5,
    mu=0.2,
    alpha=0.4,
    V=hetero_potential,
)

syst.hamiltonian_submatrix(params=params);