r"""
Mirror Chern Numbers: a Topological Crystalline Insulator
================================================================

Two copies of a strong topological insulator make a trivial one, as far
as time reversal can tell: the :math:`\mathbb{Z}_2` indices add modulo 2,
so all four vanish, and the two surface Dirac cones can in principle
annihilate. Teo, Fu and Kane (2008) pointed out that a crystal symmetry
can still protect them. On a plane of the Brillouin zone left invariant
by a mirror :math:`M` (here :math:`z\to-z`, the plane :math:`k_z = 0`),
:math:`M` commutes with :math:`H(\mathbf{k})`, the bands split into the
mirror sectors :math:`M = \pm i`, and each sector has its own Chern number.
Their half-difference is the *mirror Chern number*

.. math::

    n_M = \frac{C_{+i} - C_{-i}}{2}\, ,

an integer, of which the :math:`\mathbb{Z}_2` index of the plane is only
the parity. Hsieh, Lin, Liu, Duan, Lin and Fu (2012) predicted that SnTe is
such a *topological crystalline insulator*, with :math:`n_M = 2` and
trivial :math:`\mathbb{Z}_2` indices, and photoemission found its surface
Dirac cones the same year (Tanaka et al.; Dziawa et al.; Xu et al.).

The model: two copies of the cubic topological insulator of
``plot_3d_topological_insulator.py``
(:math:`H = m(\mathbf{k})\tau_z + \sum_i\sin k_i\,\tau_x\sigma_i`), coupled by
a hopping :math:`\delta` between the copies that respects time reversal and
the mirror :math:`M = i\tau_z\sigma_z`.
:meth:`~tbkit.kspace.KSpace.mirror_chern_number` computes :math:`n_M`.
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Line3DCollection

from tbkit.lattice import Lattice
from tbkit.kspace import KSpace, PAULI
from tbkit.system import System

cubic = [(1., 0., 0.), (0., 1., 0.), (0., 0., 1.)]
SIGMA = [PAULI['x'], PAULI['y'], PAULI['z']]
MIRROR = 1j * np.kron(np.diag([1., -1.]), PAULI['z'])  # z -> -z: s even, p_z odd, i sigma_z on spin
TIME_REVERSAL = np.kron(np.eye(2), 1j * PAULI['y'])


def cubic_ti_hoppings(s, p):
    '''The topological insulator on orbitals s, p (each with spin).'''
    hops = []
    for d in range(3):
        R = tuple(int(x) for x in np.eye(3, dtype=int)[d])
        hops += [{'i': s, 'j': s, 'R': R, 't': PAULI['0']}, {'i': p, 'j': p, 'R': R, 't': -PAULI['0']},
                     {'i': s, 'j': p, 'R': R, 't': -0.5j * SIGMA[d]},
                     {'i': p, 'j': s, 'R': R, 't': -0.5j * SIGMA[d]}]
    return hops


def single(M):
    lat = Lattice(unit_cell=[{'tag': t, 'r0': (0., 0., 0.)} for t in 'sp'], prim_vec=cubic)
    ti = KSpace(lat, spin=True)
    ti.set_hopping(cubic_ti_hoppings(0, 1))
    ti.set_onsite({'s': (M - 6) * PAULI['0'], 'p': -(M - 6) * PAULI['0']})
    return ti


def doubled(M, delta=0.3):
    lat = Lattice(unit_cell=[{'tag': t, 'r0': (0., 0., 0.)} for t in 'spSP'], prim_vec=cubic)
    tci = KSpace(lat, spin=True)
    tci.set_hopping(cubic_ti_hoppings(0, 1) + cubic_ti_hoppings(2, 3)
                           + [{'i': 0, 'j': 2, 'R': (0, 0, 0), 't': delta * PAULI['0']},
                              {'i': 1, 'j': 3, 'R': (0, 0, 0), 't': delta * PAULI['0']}])
    tci.set_onsite({tag: sign * (M - 6) * PAULI['0'] for tag, sign in zip('spSP', (1, -1, 1, -1))})
    return tci


mirror_2 = np.kron(np.eye(2), MIRROR)
tci = doubled(2.)
assert tci.symmetry_error(np.kron(np.eye(2), TIME_REVERSAL), antiunitary=True) < 1e-12
for kz in (0., np.pi):
    ham = tci.get_ham([0.4, -1.1, kz])
    assert np.allclose(mirror_2 @ ham, ham @ mirror_2)  # M commutes with H on k_z = 0, pi
ham = tci.get_ham([0.4, -1.1, 0.7])
assert not np.allclose(mirror_2 @ ham, ham @ mirror_2)  # but not elsewhere

# %%
# The lattice
# -----------
# The simple cubic lattice, here :math:`4\times4\times4` sites, with its
# nearest-neighbour bonds. Each site holds an s and a p orbital with spin
# (two such pairs in the doubled model), and every bond carries a hopping
# matrix that couples them (not drawn).

lat_vis = Lattice(unit_cell=[{'tag': 'a', 'r0': (0., 0., 0.)}], prim_vec=[(1., 0., 0.), (0., 1., 0.), (0., 0., 1.)])
lat_vis.get_lattice(4, 4, 4)
vis = System(lat_vis)
vis.set_hopping([{'n': 1, 't': 1.}])
pos = np.column_stack([lat_vis.coor[f] for f in 'xyz'])
bonds = np.stack([pos[vis.hop['i'].astype(int)], pos[vis.hop['j'].astype(int)]], axis=1)
fig_lat = plt.figure(figsize=(4.5, 4.5))
ax_lat = fig_lat.add_subplot(projection='3d')
ax_lat.add_collection3d(Line3DCollection(bonds, colors='k', linewidths=1))
ax_lat.scatter(*pos.T, s=30, c='b', depthshade=False)
ax_lat.set_aspect('equal')
ax_lat.set_axis_off()

# %%
# A single topological insulator: nu is the parity of n_M
# -------------------------------------------------------------
# On each mirror plane, :math:`\nu = n_M` modulo 2. For :math:`0 < M < 4`
# the band inversion at :math:`\Gamma` gives :math:`n_M = 1` on
# :math:`k_z = 0`, and nothing on :math:`k_z = \pi`.

for M in (-2., 2., 6.):
    ti = single(M)
    nu0, nu1, nu2, nu3 = ti.z2_indices_3d([0, 1], nk=60, nk_perp=31)
    planes = {0.: (nu0 + nu3) % 2, 0.5: nu3}  # nu(k_z = 0), nu(k_z = pi)
    for k_fixed, nu in planes.items():
        n_m = ti.mirror_chern_number(MIRROR, [0, 1], k_fixed=k_fixed)
        assert np.isclose(n_m, np.round(n_m)) and int(np.round(n_m)) % 2 == nu
    print('M = {:4.1f}: indices ({}; {}{}{}), n_M(k_z = 0) = {:+d}'.format(
        M, nu0, nu1, nu2, nu3, int(np.round(ti.mirror_chern_number(MIRROR, [0, 1])))))

# %%
# Two copies: Z2 says trivial, the mirror Chern number does not
# ------------------------------------------------------------------

occupied = [0, 1, 2, 3]
indices = tci.z2_indices_3d(occupied, nk=60, nk_perp=31)
n_m = tci.mirror_chern_number(mirror_2, occupied)
n_m_pi = tci.mirror_chern_number(mirror_2, occupied, k_fixed=0.5)
print('Two coupled copies: Z2 indices {}, n_M = {:.3f} on k_z = 0, {:.3f} on k_z = pi.'
          .format(indices, n_m, n_m_pi))
assert indices == (0, 0, 0, 0)
assert np.isclose(n_m, 2.) and np.isclose(n_m_pi, 0., atol=1e-8)
assert np.isclose(doubled(-2.).mirror_chern_number(mirror_2, occupied), 0., atol=1e-8)

# %%
# The surface states protected by the mirror
# ---------------------------------------------
# The (100) surface keeps the mirror :math:`z\to-z`, and its mirror line
# :math:`k_z = 0` carries :math:`|n_M| = 2` pairs of surface states across the
# bulk gap: two Dirac cones at :math:`\bar\Gamma`, at the energies
# :math:`\pm\delta` set by the coupling, whose branches cross at :math:`E = 0`
# at :math:`k_y = \pm\delta`. The bulk gap, :math:`|E| < 0.7`, is empty, and
# the doubled trivial insulator has no surface state.

kys = np.linspace(-np.pi, np.pi, 121)
path = np.array([[0., ky, 0.] for ky in kys])
energies = np.linspace(-2.5, 2.5, 201)
eta = 0.02
in_gap = np.abs(energies) < 0.5
surface = tci.surface_spectral_function(path, energies, direction=0, eta=eta)
bulk = tci.surface_spectral_function(path, energies, direction=0, eta=eta, bulk=True)
surface_triv = doubled(-2.).surface_spectral_function(path, energies, direction=0, eta=eta)
assert bulk[:, in_gap].max() < 0.1 and surface_triv[:, in_gap].max() < 0.1
col = surface[:, np.argmin(np.abs(energies))]
peaks = [i for i in range(1, len(col) - 1) if col[i] > max(col[i - 1], col[i + 1]) and col[i] > 1.]
print('E = 0 on the (100) surface: surface states at k_y =', np.round(kys[peaks], 3))
assert len(peaks) == 2 and np.allclose(np.abs(kys[peaks]), 0.3, atol=0.03)

# sphinx_gallery_thumbnail_number = 2
fig, axes = plt.subplots(1, 2, figsize=(10, 4.2), sharey=True)
for ax, spec, title in zip(axes, (surface, surface_triv),
                                     (r'$M = 2$: $n_M = 2$, $\mathbb{Z}_2$ trivial', r'$M = -2$: $n_M = 0$')):
    ax.pcolormesh(kys, energies, np.log10(spec.T + 1e-3), cmap='magma', vmin=-2, vmax=1.5,
                         shading='auto')
    ax.set_xlabel(r'$k_y$ (mirror line $k_z = 0$)')
    ax.set_title(title)
axes[0].set_ylabel('$E$')
fig.suptitle(r'(100) surface spectral function of two coupled topological insulators')
fig.set_layout_engine('tight')