r"""
Automatic Lead Attachment: Narrow Contacts on a Wide Cavity
=====================================================================

A lead is a semi-infinite repetition of a unit cell. To attach it, one
needs the device sites it touches and the hoppings :math:`\tau` between
them and the lead's first cell. For a straight strip, these are the edge
column and the identity times :math:`t`; for anything else, writing them
by hand is error-prone.

:meth:`~tbkit.transport.Transport.attach_lead` finds them from geometry.
The lead is a 1D *KSpace* model (here a strip cut with
:func:`~tbkit.kspace.ribbon`) whose orbitals sit at their positions in
the device's coordinates. The device sites at those positions (with the
same tags), translated by the lead's primitive vector, are copies of lead
orbitals; the outermost such cell along the lead's direction is the
interface, and the lead's own hopping to the next cell couples it to the
lead.

Here two narrow leads, at different heights, contact a wide square
cavity: only the sites under each lead's footprint couple, and the
transmission is bounded by the number of open modes of the narrow leads.
"""
import numpy as np
import matplotlib.pyplot as plt

import tbkit.lattices as lattices
from tbkit.kspace import ribbon
from tbkit.lattice import Lattice
from tbkit.system import System
from tbkit.transport import Transport, lead_from_kspace


t = -1.
square = [{'i': 0, 'j': 0, 'R': (1, 0), 't': t}, {'i': 0, 'j': 0, 'R': (0, 1), 't': t}]


def strip(width, y0=0.):
    '''A lead along x: a strip of *width* sites whose lowest site is at height y0.'''
    lat = Lattice(unit_cell=[{'tag': 'a', 'r0': (0., y0)}], prim_vec=[(1., 0.), (0., 1.)])
    return ribbon(lat, square, width=width, direction=1)


def system(length, width):
    lat = lattices.square()
    lat.get_lattice(length, width)
    sys = System(lat)
    sys.set_hopping([{'n': 1, 't': t}])
    sys.set_onsite({'a': 0.})
    sys.get_ham()
    return sys


# %%
# The same leads as by hand
# --------------------------
# On a straight strip, the automatic coupling must be the hand-written
# one: the edge columns, each site coupled to the lead orbital at its height.

width, length = 8, 12
sys = system(length, width)
lead = strip(width)
x, y = sys.lat.coor['x'], sys.lat.coor['y']

by_hand = Transport(sys.ham)
for x_edge, direction in ((0., -1), (length - 1., 1)):
    h0, v = lead_from_kspace(lead, direction)
    by_hand.add_lead(h0, v, t * np.eye(width), list(np.flatnonzero(np.isclose(x, x_edge))))

auto = Transport(sys.ham)
auto.attach_lead(sys, lead, -1)
auto.attach_lead(sys, lead, 1)

energies = np.linspace(-3.93, 3.93, 41)
assert np.allclose(auto.transmission(energies), by_hand.transmission(energies), atol=1e-8)
print('Straight strip: attach_lead reproduces the hand-built leads.')

# %%
# Narrow leads on a wide cavity
# ------------------------------
# Two leads of 6 sites, the left one near the bottom of the cavity and the
# right one near the top. Each attaches to the 6 sites of the outermost
# column under its footprint.

w_lead, width, length = 6, 20, 30
y_left, y_right = 2., 12.
sys = system(length, width)
x, y = sys.lat.coor['x'], sys.lat.coor['y']
tr = Transport(sys.ham)
tr.attach_lead(sys, strip(w_lead, y_left), -1)
tr.attach_lead(sys, strip(w_lead, y_right), 1)

for (_, _, _, sites), x_edge, y0 in zip(tr.leads, (0., length - 1.), (y_left, y_right)):
    assert np.allclose(x[sites], x_edge)
    assert np.allclose(np.sort(y[sites]), y0 + np.arange(w_lead))
print('Contacts found: left x = 0, y = 2..7; right x = 29, y = 12..17.')

# %%
# The transmission is bounded by the open modes of the leads, and the
# scattering matrix is unitary (current is conserved through the
# automatically built interfaces).

energies = np.linspace(-3.93, -0.07, 60)
T = tr.transmission(energies)
modes = np.array([sum(abs(e - 2*t*np.cos(n*np.pi/(w_lead + 1))) < 2*abs(t)
                          for n in range(1, w_lead + 1)) for e in energies])
assert np.all(T <= modes + 1e-8) and np.all(T >= -1e-10)
s = tr.smatrix(-1.5)
assert np.allclose(s.data.conj().T @ s.data, np.eye(s.data.shape[1]), atol=1e-8)
print('T(E) <= open modes of the 6-site leads; S is unitary.')

fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(10, 4), width_ratios=(1.3, 1))
ax0.scatter(x, y, s=6, c='0.7')
for (_, _, _, sites), x_edge, y0, sign, color in zip(
        tr.leads, (0., length - 1.), (y_left, y_right), (-1, 1), ('C0', 'C3')):
    ax0.scatter(x[sites], y[sites], s=20, c=color)
    cells = x_edge + sign * np.arange(1, 5)
    gx, gy = np.meshgrid(cells, y0 + np.arange(w_lead))
    ax0.scatter(gx, gy, s=20, facecolors='none', edgecolors=color)
ax0.set_aspect('equal')
ax0.set_title('cavity (grey), interface sites, lead cells')
ax0.set_xlabel('$x$')
ax0.set_ylabel('$y$')
ax1.plot(energies, modes, '--k', lw=1, label='open modes of a lead')
ax1.plot(energies, T, '-b', label='$T(E)$')
ax1.set_xlabel('$E$')
ax1.set_ylabel('$T$')
ax1.legend()
ax1.set_title('narrow leads, wide cavity')
fig.set_layout_engine('tight')