#!/usr/bin/env python 
import sys
from state import State
from matplotlib import pyplot as plt

def main(argv):
  path = argv[1]
  with open(path) as f:
    lines = f.readlines()
    count = int(lines[0])
    size  = float(lines[2])
    split = count + 3
    states = [lines[i:i+split] for i in range(0, len(lines), split)]
    states = [State.deserialize(state) for state in states]

  fig = plt.figure()
  ax = fig.add_subplot(111, projection='3d')
  current = 0
  while True:
    ax.cla()
    plt.xlim([0,size])
    plt.ylim([0,size])
    plt.title(f'Step {current}')
    ax.set_zlim([0,size])
    ax.scatter(
      states[current].molecules[:,0], 
      states[current].molecules[:,1],
      states[current].molecules[:,2],
      c=range(count)
    )
    plt.pause(0.02)

    current += 1
    current %= len(states)


if __name__ == '__main__':
  main(sys.argv)