r"""
Batched and Threaded k-Space Diagonalization: the Haldane Model on 90,000 k-Points
===================================================================================

The band and topology tools of :class:`~tbkit.kspace.KSpace` diagonalize
:math:`H(\mathbf{k})` at every point of a mesh. Calling the eigensolver
once per k-point spends most of its time in Python, not in LAPACK. tbkit
builds all the Bloch matrices with one vectorized Bloch sum, then
diagonalizes them with stacked LAPACK calls (``numpy.linalg.eigh`` on an
array of shape ``(nk, norb, norb)``), in chunks of bounded memory.
:meth:`~tbkit.kspace.KSpace.set_workers` spreads the chunks over threads
(LAPACK releases the GIL).

Here: the Berry curvature of the Haldane model on a 300 x 300 mesh. The
batched energies match the one-call-per-k-point loop, the Chern number
does not depend on the number of threads, and the timings show the cost
per k-point of each approach.
"""
import os
import time

import numpy as np
import matplotlib.pyplot as plt
import scipy.linalg as LA

from tbkit.bridges import finite_system
from tbkit.kspace import KSpace
from tbkit.lattice import Lattice
from tbkit.plot import Plot

DX, DY = 0.5 * np.sqrt(3), 0.5
lat = Lattice(unit_cell=[{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'b', 'r0': (DX, DY)}],
              prim_vec=[(2 * DX, 0.), (DX, 1.5)])
hal = KSpace(lat)
hal.set_hopping([{'i': 0, 'j': 1, 'R': R, 't': 1.} for R in [(0, 0), (-1, 0), (0, -1)]])
for R in [(1, 0), (0, 1), (1, -1)]:
    hal.set_hopping([{'i': 0, 'j': 0, 'R': R, 't': 0.2j}, {'i': 1, 'j': 1, 'R': R, 't': -0.2j}])
hal.set_onsite({'a': 0.3, 'b': -0.3})

# %%
# The lattice
# -----------
# The Haldane model on a patch of :math:`5\times4` cells
# (:func:`~tbkit.bridges.finite_system`), bond widths proportional to the
# amplitudes: the honeycomb lattice's two sublattices ('a' and 'b', in two
# colours), the nearest-neighbour hopping (thick) and the complex
# second-neighbour hoppings :math:`\pm 0.2i` (thin).

fig_lat = Plot(finite_system(hal, (5, 4))).lattice(plt_hop=True, ms=12, figsize=(5.5, 4.5))

# %%
# The batched solver against one scipy call per k-point
# ------------------------------------------------------
_, ks = hal.mesh_grid(60)
t0 = time.perf_counter()
loop = np.array([LA.eigvalsh(hal.get_ham(k)) for k in ks])
t_loop = (time.perf_counter() - t0) / len(ks)
t0 = time.perf_counter()
batched = hal.mesh_bands(60)
t_batched = (time.perf_counter() - t0) / len(ks)
assert np.allclose(batched, loop)

# %%
# Berry curvature on 90,000 k-points, with one thread and with several
# ---------------------------------------------------------------------
nk = 300
t0 = time.perf_counter()
curv = hal.berry_curvature(0, nk=nk)
t_serial = (time.perf_counter() - t0) / nk ** 2
workers = min(4, os.cpu_count() or 1)
hal.set_workers(workers)
t0 = time.perf_counter()
curv_threads = hal.berry_curvature(0, nk=nk)
t_threads = (time.perf_counter() - t0) / nk ** 2
assert np.allclose(curv_threads, curv)
chern = curv.sum() / (2 * np.pi)
assert abs(chern - round(chern)) < 1e-6 and abs(round(chern)) == 1
print(f'C = {chern:.6f}')

# sphinx_gallery_thumbnail_number = 2
fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(10, 4))
im = ax0.imshow(curv.T * nk ** 2 / (2 * np.pi), origin='lower', extent=(0, 1, 0, 1), cmap='RdBu_r')
fig.colorbar(im, ax=ax0, label=r'$\Omega\, n_k^2 / 2\pi$')
ax0.set_xlabel(r'$k_1$ (units of $b_1$)')
ax0.set_ylabel(r'$k_2$ (units of $b_2$)')
ax0.set_title(f'Haldane model, lower band, C = {round(chern)}')
labels = ['eigvalsh\nper k-point', 'mesh_bands', 'berry_curvature\n1 thread', f'berry_curvature\n{workers} threads']
ax1.bar(labels, 1e6 * np.array([t_loop, t_batched, t_serial, t_threads]), color=['0.6', 'C0', 'C1', 'C2'])
ax1.set_ylabel(r'time per k-point ($\mu$s)')
ax1.set_title('Cost per k-point')
ax1.tick_params(axis='x', labelsize=8)
fig.tight_layout()
plt.show()