#!/usr/bin/env python
import tsase
from sys import argv, stdout
import numpy as np

BLUE  = '\033[34m'
RED   = '\033[31m'
CYAN  = '\033[36m'
BLACK = '\033[30m'
ENDC = '\033[0m'

atoms = tsase.io.read(argv[1], skip =0, every=1)
if type(atoms) == type([]):
    atoms = atoms[0]

print "ion      neighbors"
for i in range(len(atoms)):
    distances = []
    indices = []
    colors = []
    for j in range(len(atoms)):
        if i == j: continue
        r = atoms.get_distance(i,j,True)
        rc1 = tsase.data.elements[atoms[i].symbol]['radius']
        rc2 = tsase.data.elements[atoms[j].symbol]['radius']
        diameter = rc1 + rc2
        cutoff = 1.3 * diameter
        if r<cutoff:
            distances.append(r)
            indices.append(j)

            if r < 0.9*diameter:
                colors.append(RED)
            elif r > 1.2*diameter:
                colors.append(CYAN)
            else:
                colors.append('')

    print "%3i %2s: " % (i, atoms[i].symbol),
    permute = np.argsort(distances)
    distances = np.array(distances)[permute]
    indices = np.array(indices)[permute]
    colors = np.array(colors)[permute]
    for i in range(len(distances)):
        stdout.write(colors[i])
        stdout.write("%3i %.3f " % (indices[i], distances[i]))
        stdout.write(ENDC)
        if (i+1)%6 == 0 and i+1 != len(distances): stdout.write("\n        ")
    print

