r"""
The Spin Chern Number: Quantum Spin Hall Without Spin Conservation
=========================================================================

In Kane and Mele's model each spin is a Haldane model: spin up has Chern
number :math:`+1`, spin down :math:`-1`, and the *spin Chern number*
:math:`C_s = (C_\uparrow - C_\downarrow)/2 = 1` counts the helical edge
pairs. Rashba coupling ruins this picture, since :math:`s_z` is no longer
conserved and the bands are no longer spin up or spin down.

Sheng, Weng, Sheng and Haldane (2006) found numerically that a spin Chern
number survives anyway, and E. Prodan (2009) made it exact: project the
spin operator onto the occupied bands, :math:`P s_z P`. As long as its
spectrum keeps a gap around zero (a *spin gap*), its positive and
negative eigenvectors split the occupied bands into two sectors, each with
a well-defined Chern number :math:`C_\pm`, and

.. math::

    C_s = \frac{C_+ - C_-}{2}

is quantized. With time reversal, :math:`C_s` modulo 2 is the
:math:`\mathbb{Z}_2` invariant; :math:`C_s` survives even when time
reversal is broken, as long as the bulk gap and the spin gap stay open.

:meth:`~tbkit.kspace.KSpace.spin_chern_number` computes it
(:meth:`~tbkit.kspace.KSpace.sector_chern_numbers` gives :math:`C_\pm` for any
operator).
"""
import numpy as np
import matplotlib.pyplot as plt

import tbkit.lattices as lattices
from tbkit.kspace import KSpace, PAULI
from tbkit.plot import Plot
from tbkit.system import System


lat = lattices.honeycomb()
lam = 0.06
S_Z = np.kron(np.eye(2), PAULI['z'])  # sigma_z on both sites


def kane_mele(M=0.1, rashba=0.):
    km = KSpace(lat, spin=True)
    km.set_hopping([{'i': 0, 'j': 1, 'R': R, 't': 1.} for R in [(0, 0), (-1, 0), (0, -1)]])
    for R in [(0, 1), (-1, 0), (1, -1)]:
        km.set_hopping([{'i': 0, 'j': 0, 'R': R, 't': 1j*lam*PAULI['z']},
                                {'i': 1, 'j': 1, 'R': R, 't': -1j*lam*PAULI['z']}])
    if rashba:
        tau = np.array([d['r0'] for d in lat.unit_cell])
        a = np.array(lat.prim_vec)
        for R in [(0, 0), (-1, 0), (0, -1)]:
            d = tau[1] + R[0]*a[0] + R[1]*a[1] - tau[0]
            d = d / np.linalg.norm(d)
            km.set_hopping([{'i': 0, 'j': 1, 'R': R,
                                     't': 1j*rashba*(PAULI['x']*d[1] - PAULI['y']*d[0])}])
    km.set_onsite({'a': M, 'b': -M})
    return km


def spin_spectrum(km, ks):
    '''Eigenvalues of P s_z P within the two occupied bands.'''
    _, vec = km.get_bands(ks, eigenvec=True)
    occ = vec[:, :, :2]
    return np.linalg.eigvalsh(occ.conj().transpose(0, 2, 1) @ S_Z @ occ)


# %%
# The lattice
# -----------
# The Kane-Mele model lives on the honeycomb lattice: sublattices 'a'
# (onsite :math:`+M`) and 'b' (:math:`-M`) in two colours, each site
# holding a spin-1/2 orbital, with the nearest-neighbour hopping (thick)
# and the spin-orbit hoppings :math:`\pm i\lambda\sigma_z` between second
# neighbours (thin, not to scale; their spin matrices are not drawn).

patch = lattices.honeycomb()
patch.get_lattice(n1=5, n2=4)
vis = System(patch)
vis.set_hopping([{'n': 1, 't': 1.}, {'n': 2, 't': 0.3}])
fig_lat = Plot(vis).lattice(plt_hop=True, ms=12, figsize=(5.5, 4.5))

# %%
# Without Rashba coupling: the Chern numbers of the two spins
# ----------------------------------------------------------------
# :math:`P s_z P` has eigenvalues exactly :math:`\pm1`: the sectors are spin
# up and spin down, with Chern numbers :math:`\pm1`.

c_plus, c_minus = kane_mele().sector_chern_numbers(S_Z, [0, 1])
print('Rashba 0: C_+ = {:.3f}, C_- = {:.3f}'.format(c_plus, c_minus))
assert np.allclose([c_plus, c_minus], [1., -1.])

# %%
# With Rashba coupling: the spin gap, and the invariant
# ---------------------------------------------------------
# Rashba coupling does not commute with :math:`s_z`, and the eigenvalues of
# :math:`P s_z P` move away from :math:`\pm1` -- but not to zero. The spin
# Chern number stays 1 until the *bulk* gap closes (between
# :math:`\lambda_R = 0.15` and 0.2), where the insulator turns trivial; it
# agrees with the :math:`\mathbb{Z}_2` invariant of the Wannier-centre flow
# throughout.

rashbas = np.array([0., 0.05, 0.1, 0.15, 0.2, 0.3])
fracs = np.stack(np.meshgrid(np.arange(40) / 40, np.arange(40) / 40, indexing='ij'), -1).reshape(-1, 2)
c_s, nu, spin_gap, bulk_gap = [], [], [], []
for r in rashbas:
    km = kane_mele(rashba=r)
    c_s.append(km.spin_chern_number([0, 1]))
    nu.append(km.z2_invariant([0, 1], nk=60, nk_perp=41))
    spin_gap.append(np.min(np.abs(spin_spectrum(km, fracs @ km.rec_vec_k))))
    bands = km.mesh_bands(60)
    bulk_gap.append(np.min(bands[:, 2] - bands[:, 1]))
c_s = np.array(c_s)
for row in zip(rashbas, c_s, nu, spin_gap, bulk_gap):
    print('lambda_R = {:.2f}: C_s = {:+.3f}, nu = {}, spin gap {:.3f}, bulk gap {:.3f}'.format(*row))
km = kane_mele(rashba=0.1)
ham = km.get_ham([0.3, 0.7])
assert np.max(np.abs(ham @ S_Z - S_Z @ ham)) > 0.05  # s_z is not conserved
assert np.allclose(c_s, [1, 1, 1, 1, 0, 0], atol=1e-6)
assert np.allclose(np.round(c_s) % 2, nu)
assert min(spin_gap) > 0.1

# %%
# Without time reversal
# -------------------------
# An in-plane Zeeman field :math:`B s_x` breaks time reversal (so there is
# no :math:`\mathbb{Z}_2` invariant any more) and conserves no spin
# component, yet the spin Chern number is still 1, with a total Chern
# number of zero.

time_reversal = np.kron(np.eye(2), 1j*PAULI['y'])
for B in (0.02, 0.05):
    km_b = kane_mele(rashba=0.1)
    km_b.set_onsite({'a': 0.1*PAULI['0'] + B*PAULI['x'], 'b': -0.1*PAULI['0'] + B*PAULI['x']})
    assert km_b.symmetry_error(time_reversal, antiunitary=True) > 0.01
    assert np.isclose(km_b.spin_chern_number([0, 1]), 1.) and abs(km_b.chern_number([0, 1])) < 1e-8
    print('Zeeman B = {}: time reversal broken, C_s = {:.3f}.'.format(B, km_b.spin_chern_number([0, 1])))

# %%
# The spectrum of P s_z P along a path
# ----------------------------------------

b1, b2 = km.rec_vec_k
corners = [np.zeros(2), (b1 - b2) / 3, b1 / 2, np.zeros(2)]  # Gamma, K, M, Gamma
kpts = np.concatenate([np.linspace(corners[i], corners[i + 1], 60, endpoint=False) for i in range(3)])

# sphinx_gallery_thumbnail_number = 2
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
for r, color in ((0., 'k'), (0.1, 'b'), (0.15, 'r')):
    w = spin_spectrum(kane_mele(rashba=r), kpts)
    axes[0].plot(w, color=color, label=[r'$\lambda_R = {}$'.format(r), None])
axes[0].axhline(0., color='gray', lw=0.5)
axes[0].set_xticks([0, 60, 120, 179], [r'$\Gamma$', '$K$', '$M$', r'$\Gamma$'])
axes[0].set_ylabel(r'eigenvalues of $P s_z P$')
axes[0].set_title('the spin gap stays open')
axes[0].legend()
axes[1].plot(rashbas, c_s, 'o-b', label='$C_s$')
axes[1].plot(rashbas, nu, 's--r', label=r'$\nu$ (Wannier flow)')
axes[1].plot(rashbas, np.array(bulk_gap) / max(bulk_gap), ':k', label='bulk gap (scaled)')
axes[1].set_xlabel(r'Rashba coupling $\lambda_R$')
axes[1].legend()
fig.set_layout_engine('tight')