# Copyright (c) 2011, 2014 Bartosz Szczesny <bszcz@bszcz.org>
# This program is free software under the MIT license.

import pylab as pl

def main():
	shifts = [
		[-1,  1], [0,  1], [1,  1],
		[-1,  0],          [1,  0],
		[-1, -1], [0, -1], [1, -1]
	]

	num_atoms = 100
	num_dims = 2 # dimensions
	coords = pl.random((num_atoms, num_dims))
	chosen = pl.random_integers(num_atoms) # from 1 to num_atoms
	chosen -= 1 # from 0 to num_atoms - 1

	for i in range(len(shifts)):
		coords = pl.vstack((coords, coords[:num_atoms] + shifts[i]))
	num_atoms *= 9 # after 8 shifts added

	max_distance = 0.9
	for i in range(num_atoms):
		if i != chosen:
			dx = coords[chosen, 0] - coords[i, 0]
			dy = coords[chosen, 1] - coords[i, 1]
			distance = pl.sqrt(dx*dx + dy*dy)
			if distance < max_distance:
				pl.plot([coords[i, 0]], [coords[i, 1]], "bo")
			else:
				pl.plot([coords[i, 0]], [coords[i, 1]], "ko")

	# plot last for visibility
	pl.plot([coords[chosen, 0]], [coords[chosen, 1]], "ro")
	pl.grid(True)
	pl.show()

if __name__ == "__main__":
	main()
