#!/usr/bin/env python
# -*- coding: utf-8 -*-
from gi.repository import v_sim
import cProfile

render = v_sim.UiMainClass.getDefaultRendering()
data = render.getVisuData()

# Example with integrated iterators (still an issue with the
# stop criterion)
rad = {}
for it in data.__iterElements__():
  rad[it.element.getName()] = v_sim.RenderingAtomic.getRadius(it.element)
#print "Radii: ", rad

def hit((x1,y1,z1), r1, (x2,y2,z2), r2):
  return ( (x1-x2)**2 + (y1-y2)**2 + (z1-z2)**2 < (r1 + r2)**2 )

# With Python iterators.
collision = set()

def byPairs(data, collision):
  for (it1, it2) in data.__iterByPairs__():
    if hit(data.getNodeCoordinates(it1.node), rad[it1.element.getName()], \
           data.getNodeCoordinates(it2.node), rad[it2.element.getName()]):
      collision.add(int(it1.node.number))
      collision.add(int(it2.node.number))

def doubleLoop(data, collision):
  for it1 in data:
    for it2 in data:
      if it1.node.number > it2.node.number and \
         hit(data.getNodeCoordinates(it1.node), rad[it1.element.getName()], \
             data.getNodeCoordinates(it2.node), rad[it2.element.getName()]):
        collision.add(int(it1.node.number))
        collision.add(int(it2.node.number))

def byCopy(data, collision):
  nodes = data.getAllNodePositions()
  for (c1, i1, e1) in nodes:
    for (c2, i2, e2) in nodes:
      if i1 > i2 and hit(c1, rad[e1], c2, rad[e2]):
        collision.add(int(i1))
        collision.add(int(i2))

cProfile.run('doubleLoop(data, collision)')
cProfile.run('byPairs(data, collision)')
cProfile.run('byCopy(data, collision)')

marks = render.getMarks()
if marks.setHighlightedList(tuple(collision), v_sim.MarksStatus.SET):
  v_sim.Object.redraw("My script")

#v_sim.UiMainClass.getCurrentPanel().quit(True)