"""Short-range forces in O(N): linked cells and Verlet lists for a 2D Lennard-Jones system.

Cleaned from the Scientific Computing 2 (TUM, 2020) worksheet 3 (Verlet-list_solution.ipynb) and the
linked-cell assignment (cell_2.py; the notebook itself is left with unresolved git merge conflicts).
Notes: ../../Molecular_dynamics.md, ../../Barnes-Hut_algorithm.md (the long-range counterpart)

  linked cells  bin particles into square cells of side >= r_c; a particle only looks at its own cell and
                the 8 neighbours. Cost O(N) per step, rebuilt every step.
  Verlet list   store, per particle, everyone within r_c + skin. Valid until some particle has moved
                more than skin/2, so it is rebuilt only every few steps (built here with linked cells).

Units: sigma = epsilon = 1, U(r) = 4 (r^-12 - r^-6) cut at r_c (not shifted). Open boundaries.
Changes against the course code: vectorised per cell, one potential instead of the course's
U = 1e-15/r^12 - 1e-5/r^6, no np.int / np.float.

Run:  python neighbour_lists.py
"""
import time

import numpy as np


def lj(r):
    s6 = r ** -6
    return 4 * (s6 * s6 - s6)


def energy_direct(x, rc):
    r = np.linalg.norm(x[:, None] - x[None], axis=2)
    iu = np.triu_indices(len(x), 1)
    r = r[iu]
    return lj(r[r < rc]).sum()


def cell_bins(x, size, box):
    n = max(1, int(box // size))
    c = np.minimum((x / (box / n)).astype(int), n - 1)
    cells = {}
    for k, key in enumerate(map(tuple, c)):
        cells.setdefault(key, []).append(k)
    return {key: np.array(v) for key, v in cells.items()}


def pairs_linked_cells(x, rc, box):
    """All pairs (i < j) closer than rc, found through cells of side >= rc."""
    cells = cell_bins(x, rc, box)
    I, J = [], []
    for (cx, cy), a in cells.items():
        for dx, dy in ((0, 0), (1, 0), (1, 1), (0, 1), (-1, 1)):    # half the neighbours: each pair once
            b = cells.get((cx + dx, cy + dy))
            if b is None:
                continue
            r = np.linalg.norm(x[a, None] - x[None, b], axis=2)
            mask = r < rc
            if dx == 0 and dy == 0:
                mask &= a[:, None] < b[None, :]
            ii, jj = np.nonzero(mask)
            I.append(a[ii]); J.append(b[jj])
    return np.concatenate(I), np.concatenate(J)


def energy_pairs(x, I, J, rc):
    r = np.linalg.norm(x[I] - x[J], axis=1)
    return lj(r[r < rc]).sum()


def main():
    rng = np.random.default_rng(3)
    rc, skin = 2.5, 0.3
    for n in (500, 1000, 2000, 4000):
        box = np.sqrt(n / 0.5)                            # number density 0.5
        g = int(np.ceil(np.sqrt(n)))                       # jittered lattice: no overlapping particles
        grid = np.stack(np.meshgrid(np.arange(g), np.arange(g)), -1).reshape(-1, 2)[:n]
        x = (grid + 0.5 + 0.2 * (rng.random((n, 2)) - 0.5)) * box / g
        t0 = time.perf_counter(); e_dir = energy_direct(x, rc); t1 = time.perf_counter()
        I, J = pairs_linked_cells(x, rc, box); e_lc = energy_pairs(x, I, J, rc); t2 = time.perf_counter()
        VI, VJ = pairs_linked_cells(x, rc + skin, box)     # Verlet list, reused while moves < skin/2
        x_moved = x + 0.1 * skin * (rng.random((n, 2)) - 0.5)
        t3 = time.perf_counter(); e_v = energy_pairs(x_moved, VI, VJ, rc); t4 = time.perf_counter()
        e_v_ref = energy_direct(x_moved, rc)
        print(f"N = {n:5d}: direct {t1 - t0:.3f} s | linked cells {t2 - t1:.3f} s, |dE| = {abs(e_lc - e_dir):.1e}"
              f" | Verlet reuse {t4 - t3:.4f} s, |dE| = {abs(e_v - e_v_ref):.1e}")


if __name__ == '__main__':
    main()