r"""
Vacancy-Induced Zero Modes in Graphene
============================================

Pereira, Guinea, Lopes dos Santos, Peres and Castro Neto (2006) showed
that removing a single carbon atom from graphene binds a state at
exactly the Dirac energy :math:`E = 0`. It lives only on the sublattice
*opposite* to the vacancy and is quasi-localized: its amplitude falls as
:math:`1/r`, so slowly that it is not normalizable in an infinite sheet.

The count follows from the bipartite structure (the sublattice
imbalance behind Lieb's theorem): a nearest-neighbour Hamiltonian on
sublattices A and B has at least :math:`\bigl||A| - |B|\bigr|` zero
modes, on the majority sublattice. One vacancy gives one; two vacancies
on the same sublattice give two; two on opposite sublattices give none,
in general.

The vacancy comes from :meth:`~tbkit.lattice.Lattice.remove_sites`, and
its limit from :meth:`~tbkit.system.System.set_onsite_def`: an onsite
energy :math:`U \to \infty` on one site. The states near :math:`E = 0`
come from the sparse shift-invert solver
:meth:`~tbkit.system.System.get_eig_sparse`.
"""
import numpy as np
import matplotlib.pyplot as plt

import tbkit.lattices as lattices
from tbkit.plot import Plot
from tbkit.system import System


L = 32  # not a multiple of 3: the clean torus has a gap at E = 0


def graphene_torus(vacancies=(), impurity=None):
    lat = lattices.honeycomb()
    lat.get_lattice(n1=L, n2=L)
    if vacancies:
        lat.remove_sites(sorted(vacancies))
    sys = System(lat, periodic=True)
    sys.set_hopping([{'n': 1, 't': -1.}])
    if impurity is not None:
        sys.set_onsite({'a': 0., 'b': 0.})
        sys.set_onsite_def({impurity[0]: impurity[1]})
    sys.get_ham()
    sys.get_eig_sparse(n_eig=6, sigma=0., eigenvec=True)  # the states nearest E = 0
    return sys


lat0 = lattices.honeycomb()
lat0.get_lattice(n1=L, n2=L)
coor = lat0.coor
centre = int(np.argmin(np.hypot(coor['x'] - coor['x'].mean(), coor['y'] - coor['y'].mean())))
vac_tag = coor['tag'][centre]
other = 'b' if vac_tag == 'a' else 'a'

# %%
# The lattice
# -----------
# The honeycomb lattice, sublattices 'a' and 'b' in two colours, here an
# :math:`8\times8` patch with its central site removed: the vacancy
# leaves three dangling bonds, all ending on the other sublattice.

lat_vis = lattices.honeycomb()
lat_vis.get_lattice(n1=8, n2=8)
c_vis = lat_vis.coor
lat_vis.remove_sites([int(np.argmin(np.hypot(c_vis['x'] - c_vis['x'].mean(),
                                                         c_vis['y'] - c_vis['y'].mean())))])
vis = System(lat_vis)
vis.set_hopping([{'n': 1, 't': -1.}])
fig_lat = Plot(vis).lattice(plt_hop=True, ms=9, figsize=(6, 4.5))

# %%
# One vacancy, one zero mode on the other sublattice
# ------------------------------------------------------
# The clean :math:`32 \times 32` torus has no state near :math:`E = 0`
# (its gap is a finite-size one). Removing one site leaves exactly one
# zero mode, entirely on the sublattice it did not belong to.

gap = np.abs(graphene_torus().en).min()
sys = graphene_torus([centre])
lat = sys.lat
zero = np.abs(sys.en) < 1e-9
assert gap > 0.1 and zero.sum() == 1
psi2 = sys.intensity[:, zero][:, 0]
assert np.isclose(psi2[lat.coor['tag'] == other].sum(), 1.)
print('one zero mode (clean gap {:.3f}), weight 1 on sublattice {}: OK'.format(gap, other))

# %%
# Quasi-localized: :math:`|\psi|^2 \sim 1/r^2`
# ------------------------------------------------
# The mean weight in rings around the vacancy (minimum-image distances on
# the torus) falls as a power law with exponent 2, not exponentially.

a = np.array(lat0.prim_vec)
d = np.stack([lat.coor['x'] - coor['x'][centre], lat.coor['y'] - coor['y'][centre]], axis=1)
frac = np.linalg.solve(a.T, d.T).T
images = L * np.array([(i, j) for i in (-1, 0, 1) for j in (-1, 0, 1)])
r = np.min(np.linalg.norm((frac[:, None, :] + images) @ a, axis=2), axis=1)
on = psi2 > 1e-6 * psi2.max()  # the sites where the mode does not vanish
shells = np.arange(1., 14.)
r_shell = np.array([r[on & (r >= s) & (r < s + 1)].mean() for s in shells])
w_shell = np.array([psi2[on & (r >= s) & (r < s + 1)].mean() for s in shells])
slope = np.polyfit(np.log(r_shell), np.log(w_shell), 1)[0]
assert abs(slope + 2.) < 0.15
print('|psi|^2 ~ r^{:.2f} around the vacancy (expected r^-2): OK'.format(slope))

# %%
# The vacancy as an infinitely strong impurity
# ------------------------------------------------
# An onsite energy :math:`U` on the same site pulls one state towards
# :math:`E = 0` as :math:`U` grows, and onto the opposite sublattice:
# the vacancy is its :math:`U \to \infty` limit.

us = np.array([2., 10., 50., 250., 1000.])
e_res, w_other = [], []
for u in us:
    sys_u = graphene_torus(impurity=(centre, u))
    n = np.argmin(np.abs(sys_u.en))
    e_res.append(abs(sys_u.en[n]))
    w_other.append(sys_u.intensity[sys_u.lat.coor['tag'] == other, n].sum())
e_res = np.array(e_res)
assert np.all(np.diff(e_res) < 0) and e_res[-1] < 1e-3
assert w_other[-1] > 0.99
print('|E| of the resonant state: {} -> 0 as U grows: OK'.format(np.round(e_res, 4)))

# %%
# Counting by sublattice imbalance
# ------------------------------------
# A second vacancy about 8 bond lengths away: on the same sublattice the
# imbalance is 2 and so is the number of zero modes; on the other one
# the imbalance vanishes and the two quasi-localized modes hybridize away
# from :math:`E = 0`.

dist = np.hypot(coor['x'] - coor['x'][centre] - 8., coor['y'] - coor['y'][centre])
same = int(np.argmin(np.where(coor['tag'] == vac_tag, dist, np.inf)))
opposite = int(np.argmin(np.where(coor['tag'] == other, dist, np.inf)))
n_same = np.sum(np.abs(graphene_torus([centre, same]).en) < 1e-9)
n_opp = np.sum(np.abs(graphene_torus([centre, opposite]).en) < 1e-9)
assert n_same == 2 and n_opp == 0
print('two vacancies: {} zero modes on the same sublattice, {} on opposite ones: OK'.format(n_same, n_opp))

# sphinx_gallery_thumbnail_number = 2
fig, axes = plt.subplots(1, 2, figsize=(11, 4.6))
sc = axes[0].scatter(lat.coor['x'], lat.coor['y'], c=np.log10(psi2 + 1e-12), s=6, cmap='magma', vmin=-6)
axes[0].plot(coor['x'][centre], coor['y'][centre], 'cx', ms=10)
axes[0].set_aspect('equal')
axes[0].axis('off')
axes[0].set_title(r'zero mode: $\log_{10}|\psi|^2$')
fig.colorbar(sc, ax=axes[0])
axes[1].loglog(r[on], psi2[on], '.', c='C7', ms=2)
axes[1].loglog(r_shell, w_shell, 'o', c='C3', label='ring means')
axes[1].loglog(r_shell, w_shell[0] * (r_shell / r_shell[0])**-2, 'k--', label=r'$1/r^2$')
axes[1].set_xlabel('distance to the vacancy $r$')
axes[1].set_ylabel(r'$|\psi|^2$')
axes[1].legend()
fig.tight_layout()