r"""
The 38-Fold Way: Symmetry Classes of Non-Hermitian Bands
==============================================================

The tenfold way (``topology/plot_tenfold_way.py``) sorts Hermitian
Hamiltonians by time reversal, particle-hole and chiral symmetry. For a
non-Hermitian Hamiltonian, complex conjugation and transposition are no
longer the same operation, :math:`H^* \neq H^T`, so each of these
symmetries splits in two. Time reversal :math:`U_TH^*(\mathbf{k})U_T^\dagger
= H(-\mathbf{k})` gets a partner :math:`\mathrm{TRS}^\dagger`,
:math:`U_{C_+}H^T(\mathbf{k})U_{C_+}^\dagger = H(-\mathbf{k})`, which is
reciprocity. Particle-hole symmetry gets a partner
:math:`\mathrm{PHS}^\dagger`, and the chiral symmetry
:math:`U_\Gamma H^\dagger U_\Gamma^\dagger = -H` becomes distinct from the
sublattice symmetry :math:`U_SHU_S^\dagger = -H` and from
pseudo-Hermiticity :math:`U_\eta H^\dagger U_\eta^\dagger = H`.

Kawabata, Shiozaki, Ueda and Sato (2019), and Zhou and Lee (2019), counted
the classes these symmetries generate: 38, not 10. Several pairs of
symmetry sets are equivalent, because multiplying :math:`H` by
:math:`i` turns time reversal into :math:`\mathrm{PHS}^\dagger`. For each
class and gap type (a *point gap*, :math:`E \neq E_P` for every
:math:`\mathbf{k}`, or a real or imaginary *line gap*), they then gave the
topological invariant in every dimension.

:meth:`~tbkit.kspace.KSpace.non_hermitian_class` checks each operator with
:meth:`~tbkit.kspace.KSpace.symmetry_error` and names the class.
"""
import itertools

import numpy as np
import matplotlib.pyplot as plt

import tbkit.lattices as lattices
from tbkit.lattice import Lattice
from tbkit.kspace import KSpace, PAULI

I1, I2, SZ = np.eye(1), np.eye(2), PAULI['z']


def chain(t_right, t_left, onsite=0.):
    model = KSpace(lattices.chain())
    model.set_hopping([{'i': 0, 'j': 0, 'R': (1,), 't': t_right},
                                {'i': 0, 'j': 0, 'R': (-1,), 't': t_left}], hermitian=False)
    model.onsite[0] = onsite
    return model


def ssh(t1, t2, gamma=0.):
    lat = Lattice(unit_cell=[{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'b', 'r0': (0.5, 0.)}],
                         prim_vec=[(1., 0.)])
    model = KSpace(lat)
    model.set_hopping([{'i': 0, 'j': 1, 'R': (0,), 't': t1 + gamma / 2},
                                {'i': 1, 'j': 0, 'R': (0,), 't': t1 - gamma / 2},
                                {'i': 1, 'j': 0, 'R': (1,), 't': t2},
                                {'i': 0, 'j': 1, 'R': (-1,), 't': t2}], hermitian=False)
    return model


# %%
# Classifying non-Hermitian chains
# --------------------------------
# The Hatano-Nelson chain :math:`H(k) = t_Re^{ik} + t_Le^{-ik}` has real
# hoppings, hence time reversal :math:`U_T = 1`. Its non-reciprocity
# :math:`t_R \neq t_L` breaks :math:`\mathrm{TRS}^\dagger`
# (:math:`H^T(k) = H(k) \neq H(-k)`): class AI. A reciprocal chain with
# loss, :math:`H(k) = 2t\cos k - i\gamma`, has :math:`\mathrm{TRS}^\dagger`
# but not time reversal: class AI†. Multiplying the Hatano-Nelson chain by
# :math:`i` swaps time reversal for :math:`\mathrm{PHS}^\dagger`: class D†,
# the same class as AI with real and imaginary line gaps exchanged
# (``canonical=True``). The non-Hermitian SSH chain of Yao and Wang adds a
# sublattice symmetry :math:`\sigma_z` to time reversal, and their product
# is a :math:`\mathrm{PHS}^\dagger`. A Hermitian model has all the
# symmetries at once, since :math:`H^T = H^*` and :math:`\eta = 1`, so the
# Hermitian SSH chain (BDI) is BDI + S++.

hatano_nelson = chain(1., 0.5)
lossy = chain(1., 1., onsite=-0.5j)
i_hatano_nelson = chain(1j, 0.5j)
nh_ssh = ssh(1., 0.8, gamma=0.6)
herm_ssh = ssh(1., 0.6)

# the broken symmetries: TRS^dagger of Hatano-Nelson, TRS of the lossy chain
assert hatano_nelson.symmetry_error(I1, transpose=True) > 0.1
assert lossy.symmetry_error(I1, antiunitary=True) > 0.1

cases = [
    ('Hatano-Nelson chain', hatano_nelson, dict(time_reversal=I1), 'AI'),
    ('reciprocal chain with loss', lossy, dict(time_reversal_dagger=I1), 'AI†'),
    ('i x Hatano-Nelson chain', i_hatano_nelson, dict(particle_hole_dagger=I1), 'D†'),
    ('non-Hermitian SSH chain', nh_ssh, dict(time_reversal=I2, sublattice=SZ), 'AI + S+'),
    ('Hermitian SSH chain', herm_ssh,
     dict(time_reversal=I2, chiral=SZ, time_reversal_dagger=I2, pseudo_hermiticity=I2), 'BDI + S++'),
]
for name, model, symmetries, expected in cases:
    cls = model.non_hermitian_class(**symmetries)
    print('{:<28s} class {:<10s} (38-fold: {})'.format(
        name, cls, model.non_hermitian_class(**symmetries, canonical=True)))
    assert cls == expected
assert i_hatano_nelson.non_hermitian_class(particle_hole_dagger=I1, canonical=True) == 'AI'
# the tenfold way only sees the Hermitian class
assert herm_ssh.tenfold_class(time_reversal=I2, chiral=SZ) == 'BDI'

# %%
# The class decides the point-gap topology
# ----------------------------------------
# In one dimension a point gap at :math:`E_P` carries the spectral winding
# number :math:`W(E_P)` (see ``plot_skin_effect.py``). The periodic table
# of Kawabata et al. allows it, :math:`\mathbb{Z}`, in classes A and AI.
# In class AI† it is trivial, 0: reciprocity forbids the winding, and with
# it the skin effect. The Hatano-Nelson ellipse winds once around
# :math:`E = 0`. The lossy chain traces its arc twice, once in each
# direction. Twenty random reciprocal chains with four orbitals
# (:math:`H^T(k) = H(-k)`) never wind around any point, while random
# chains without the symmetry do.

rng = np.random.default_rng(0)


def matrix_chain(blocks):
    """The chain with hopping matrices blocks[R] (R = -1, 0, 1) between cells n and n + R."""
    norb = len(blocks[0])
    model = KSpace(Lattice(unit_cell=[{'tag': 'a', 'r0': (0., 0.)}] * norb, prim_vec=[(1., 0.)]))
    hops = []
    for R, block in blocks.items():
        for i, j in itertools.product(range(norb), repeat=2):
            if R == 0 and i == j:
                model.onsite[i] = block[i, i]
            else:
                hops.append({'i': i, 'j': j, 'R': (R,), 't': block[i, j]})
    model.set_hopping(hops, hermitian=False)
    return model


def random_blocks(norb):
    return {R: rng.normal(size=(norb, norb)) + 1j * rng.normal(size=(norb, norb)) for R in (-1, 0, 1)}


e_refs = rng.normal(size=10) + 1j * rng.normal(size=10)
windings_reciprocal, windings_generic = [], []
for _ in range(20):
    blocks = random_blocks(4)
    reciprocal = {R: (blocks[R] + blocks[-R].T) / 2 for R in blocks}  # t(R) = t(-R)^T
    windings_reciprocal.append([matrix_chain(reciprocal).spectral_winding(e) for e in e_refs])
    windings_generic.append([matrix_chain(blocks).spectral_winding(e) for e in e_refs])
assert abs(hatano_nelson.spectral_winding(0.) - 1.) < 1e-10
assert abs(lossy.spectral_winding(-0.5j + 0.1)) < 1e-10
assert np.max(np.abs(windings_reciprocal)) < 1e-8
assert np.max(np.abs(np.round(windings_generic))) >= 1
print('largest |W| of the reciprocal chains: {:.1e}; of the generic chains: {:.0f}'.format(
    np.max(np.abs(windings_reciprocal)), np.max(np.abs(np.round(windings_generic)))))

k = np.linspace(-np.pi, np.pi, 201)
fig, axes = plt.subplots(1, 4, figsize=(13, 3.4))
for ax, (name, model, symmetries, cls) in zip(axes, cases[:4]):
    energies = np.array([np.linalg.eigvals(model.get_ham([kk])) for kk in k])
    for band in energies.T:
        ax.plot(band.real, band.imag, lw=2)
    ax.plot(0, 0, 'k+', ms=10)
    ax.set_title('{}\nclass {}'.format(name, cls), fontsize=10)
    ax.set_xlabel('Re E')
    ax.set_aspect('equal', adjustable='datalim')
axes[0].set_ylabel('Im E')
fig.tight_layout()

# %%
# Thirty-eight classes
# --------------------
# Every class can be realized. Here the operators are random products of
# Pauli matrices on eight orbitals, and the hoppings of a random chain
# are averaged over the symmetries they generate (each one maps
# :math:`H(k)` to :math:`\pm U\,\mathrm{op}[H(\pm k)]\,U^\dagger`). Random
# draws reach all 54 names of Kawabata et al.'s tables: 10 AZ classes,
# 8 AZ† classes, 3 complex and 24 real classes with sublattice symmetry,
# and 9 with pseudo-Hermiticity. The :math:`H \to iH` equivalences fold
# them into the 38 classes. As a check, every model classified again as
# :math:`iH`, with the operators relabelled accordingly, lands in the same
# class.

KINDS = {'time_reversal': (1, 0, 1), 'particle_hole': (0, 1, -1), 'chiral': (1, 1, -1),
              'time_reversal_dagger': (0, 1, 1), 'particle_hole_dagger': (1, 0, -1),
              'sublattice': (0, 0, -1), 'pseudo_hermiticity': (1, 1, 1)}
I_TIMES = {'time_reversal': 'particle_hole_dagger', 'particle_hole_dagger': 'time_reversal',
                 'chiral': 'pseudo_hermiticity', 'pseudo_hermiticity': 'chiral'}
GENERATORS = [(), ('time_reversal',), ('particle_hole',), ('chiral',), ('time_reversal_dagger',),
                         ('particle_hole_dagger',), ('sublattice',), ('pseudo_hermiticity',),
                         ('time_reversal', 'particle_hole'), ('time_reversal_dagger', 'particle_hole_dagger'),
                         ('time_reversal', 'time_reversal_dagger'), ('particle_hole', 'particle_hole_dagger'),
                         ('time_reversal', 'sublattice'), ('particle_hole', 'sublattice'),
                         ('chiral', 'sublattice'), ('time_reversal', 'particle_hole', 'sublattice')]
paulis = [PAULI[a] for a in '0xyz']
strings = [np.kron(np.kron(a, b), c) for a, b, c in itertools.product(paulis, repeat=3)]


def symmetrize(blocks, symmetries):
    for name, u in symmetries.items():
        conj, transpose, sign = KINDS[name]
        image = {}
        for R in blocks:
            t = blocks[-R if transpose else R]
            t = t.conj() if conj else t
            image[R] = sign * u @ (t.T if transpose else t) @ u.conj().T
        blocks = {R: (blocks[R] + image[R]) / 2 for R in blocks}
    return blocks


names = {}
while len(names) < 54:
    generators = GENERATORS[rng.integers(len(GENERATORS))]
    symmetries = {g: strings[rng.integers(64)] for g in generators}
    blocks = symmetrize(random_blocks(8), symmetries)
    try:  # a product of the operators may be an extra unitary symmetry: draw again
        name = matrix_chain(blocks).non_hermitian_class(**symmetries)
    except ValueError:
        continue
    if name not in names:
        times_i = matrix_chain({R: 1j * b for R, b in blocks.items()})
        relabelled = {I_TIMES.get(g, g): u for g, u in symmetries.items()}
        names[name] = (matrix_chain(blocks).non_hermitian_class(**symmetries, canonical=True),
                              times_i.non_hermitian_class(**relabelled, canonical=True))
assert all(c == c_i for c, c_i in names.values())
classes = sorted({c for c, _ in names.values()})
print('{} names, {} classes'.format(len(names), len(classes)))
assert len(classes) == 38

# %%
# The 54 names, grouped into the 38 classes (black: the representative;
# grey: an equivalent name, with the representative it maps to).

families = [
    ('AZ', ['A', 'AIII', 'AI', 'BDI', 'D', 'DIII', 'AII', 'CII', 'C', 'CI']),
    ('AZ†', ['AI†', 'BDI†', 'D†', 'DIII†', 'AII†', 'CII†', 'C†', 'CI†']),
    ('+ SLS', ['A + S', 'AIII + S+', 'AIII + S-', 'AI + S+', 'AI + S-', 'AII + S+', 'AII + S-',
                  'D + S+', 'D + S-', 'C + S+', 'C + S-']),
    ('+ SLS (BDI ... CI)', [cls + ' + S' + s for cls in ('BDI', 'DIII', 'CII', 'CI')
                                    for s in ('++', '+-', '-+', '--')]),
    ('+ pH', ['A + η', 'AI + η+', 'AI + η-', 'AII + η+', 'AII + η-', 'D + η+', 'D + η-',
                 'C + η+', 'C + η-']),
]
assert sorted(n for _, family in families for n in family) == sorted(names)
fig, ax = plt.subplots(figsize=(11, 4.4))
for col, (title, family) in enumerate(families):
    ax.text(col, 0, title, ha='center', fontsize=11, weight='bold')
    for row, name in enumerate(family, start=1):
        canonical = names[name][0]
        label = name if canonical == name else '{}  = {}'.format(name, canonical)
        ax.text(col, -0.6 * row, label, ha='center', fontsize=9.5,
                    color='k' if canonical == name else '0.6')
ax.set_xlim(-0.5, len(families) - 0.5)
ax.set_ylim(-0.6 * 16.6, 0.4)
ax.set_title('{} names, {} symmetry classes'.format(len(names), len(classes)))
ax.axis('off')
plt.show()