#!/usr/bin/env python
# coding: utf-8

# (local-chern-nb)=
# # Bianco-Resta Local Chern marker
# 
# We compute the Bianco–Resta local Chern marker for a finite Haldane patch: first verify the bulk Chern number in momentum space, then build an open-boundary supercell and map the real-space marker, trimming the boundary to recover the bulk value.
# 
# :::{admonition} What you will learn
# :class: tip
# 
# - Instantiate the Haldane model and verify its bulk Chern number on a Brillouin-zone mesh.
# - Generate a large open-boundary patch with `make_finite` and evaluate the local Chern marker.
# - Aggregate site-resolved markers into unit-cell values, trim edges, and estimate the bulk topological index.
# - Visualize the spatial distribution of the marker.
# 
# :::

# In[ ]:


from pythtb import Mesh, WFArray

import matplotlib.pyplot as plt
from pythtb.models import haldane


# ## Periodic benchmark
# 
# Construct the topological Haldane model ($\Delta=0$, $t=1$, $t_2=0.15$), check its built-in Chern number, and reproduce it via `WFArray` on a $50\times 50$ Monkhorst–Pack mesh.

# In[ ]:


model = haldane(0, 1, 0.15)
print(model)


# In[ ]:


mesh = Mesh(["k", "k"])
mesh.build_grid(shape=(50, 50))
wfa = WFArray(model.lattice, mesh)
wfa.solve_model(model)

chern_band0 = wfa.chern_number(state_idx=[0], plane=(0, 1)).item()
print(f"Chern number (valence band): {chern_band0:+.3f}")


# ## Open-boundary patch
# 
# Cut a large rectangular patch by repeating the unit cell $(L_x, L_y)$ times with `model.make_finite`. This returns a purely real-space `TBModel` (no k-axes) ready for local-marker evaluation.

# In[ ]:


# Finite OBC patch: (Lx, Ly) supercell
Lx, Ly = 42, 24
finite_model = model.make_finite(periodic_dirs=[0, 1], num_cells=[Lx, Ly])


# ## Site-resolved marker
# 
# `TBModel.local_chern_marker()` yields the Bianco–Resta marker per orbital. Optionally, we can also return the bulk averaged Chern number. This requires setting `trim_cells` to remove the edge regions where the marker rapidly varies to compensate for the bulk value. This will trim that many unit cells from each edge of the patch. The finite model must then have at least `2*trim_cells` unit cells in each direction.

# :::{tip}
# :class: dropdown
# 
# Increase `Lx`, `Ly`, or adjust `trim_cells` if the bulk average hasn’t converged.
# 
# :::

# In[ ]:


C_local, C_bulk_avg = finite_model.local_chern_marker(
    return_bulk_avg=True, trim_cells=4
)
print(f"Trimmed bulk marker: {C_bulk_avg:+.3f}")


# ## Visualize edge vs bulk
# 
# Color the orbitals by their marker value to show bulk local marker ($\approx$ Chern number / number of sites per primitive cell) and boundary oscillations.

# In[ ]:


fig, ax = plt.subplots(figsize=(10, 6), dpi=500)
pos = finite_model.get_orb_vecs(cartesian=True)
sc = ax.scatter(
    pos[:, 0],
    pos[:, 1],
    c=C_local,
    s=10,
    cmap="coolwarm",
    vmin=-1.0,
    vmax=1.0,
)
ax.set_aspect("equal")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_title(f"Local Chern marker \n Bulk average = {C_bulk_avg:+.3f}")
fig.colorbar(sc, ax=ax, label=r"$C(r)$")


# :::{admonition} Next steps
# :class: seealso
# 
# - Tune $(\Delta, t_2)$ to compare how the local marker behaves in the topological and trivial phases.
# - Compare `local_chern_marker` with hybrid Wannier flow for the same finite region.
# 
# :::