import string
import numpy as np
import kwant

def onsite(site, salt):
    return kwant.digest.uniform(site.tag, salt=salt)


def hopping(a, b, salt):
    c = bytes(a.tag) + bytes(b.tag)
    return complex(*kwant.digest.uniform2(c, salt=salt))


lat = kwant.lattice.square(norbs=1)

syst = kwant.Builder()
syst[(lat(i, j) for i in range(10) for j in range(5))] = onsite
syst[lat.neighbors()] = hopping

lead = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))
lead[(lat(0, j) for j in range(5))] = 4
lead[lat.neighbors()] = -1

syst.attach_lead(lead)
syst.attach_lead(lead.reversed())

syst = syst.finalized()


# Loop over a bunch of realizations of the randomness

params = dict(salt='')
energy = 0.5

for params['salt'] in string.ascii_lowercase:
    wf = kwant.wave_function(syst, energy=energy, params=params)
    psis = np.array([psi for i in (0, 1) for psi in wf(i)])
    norms = np.sum(np.square(np.abs(psis)), axis=1)  # sum over all orbitals
    # some norm may be 1 by coincidence, but this is unlikely.
    assert not any(np.isclose(norms, 1))

    rho = kwant.operator.Density(syst).bind(params=params)
    op = sum(rho(psi) for psi in psis) / (2 * np.pi)

    ldos = kwant.ldos(syst, energy=energy, params=params)

    np.testing.assert_array_almost_equal(op, ldos)


print('all equal!')