#!/usr/bin/env python

import sys

if len(sys.argv) != 2:
  print "Computes velocity power spectrum (autocorrelation Fourier-transform"
  print "  magnitude) as a function of frequency"
  print "Usage: %s infile" % sys.argv[0]
  sys.exit(1)

import numpy as np
from quippy import *

ar = AtomsReader(sys.argv[1])
n_frames = len(ar)
n_atoms = ar[1].n
vels = fzeros((3*n_frames, 3*n_atoms))
i_frame = 1
prev_at = None
for at in ar:
  if i_frame%10 == 0:
     print >> sys.stderr, "config %d" % i_frame
  if 'velo' in at.properties.keys():
    if not prev_at is None:
      dt = at.params['time'] - prev_at.params['time']
    for i in frange(n_atoms):
      vels[i_frame+n_frames, 1:3*at.n:3]   = at.velo[1,1:at.n]
      vels[i_frame+n_frames, 2:1+3*at.n:3] = at.velo[2,1:at.n]
      vels[i_frame+n_frames, 3:2+3*at.n:3] = at.velo[3,1:at.n]
    prev_at = at.copy()
  else:
    if not prev_at is None:
      dt = at.params['time'] - prev_at.params['time']
      for i in frange(n_atoms):
	vels[i_frame+n_frames, 1:3*at.n:3]   = (at.pos[1,1:at.n] - prev_at.pos[1,1:at.n])*dt
	vels[i_frame+n_frames, 2:1+3*at.n:3] = (at.pos[2,1:at.n] - prev_at.pos[2,1:at.n])*dt
	vels[i_frame+n_frames, 3:2+3*at.n:3] = (at.pos[3,1:at.n] - prev_at.pos[3,1:at.n])*dt
    prev_at = at.copy()
  i_frame += 1
print >> sys.stderr, "dt %f n_frames %d" % (dt, n_frames)

if (len(vels[:,1]) % 2) == 0:
  mean_vel_fft = np.zeros(len(vels[:,1])/2+1)
else:
  mean_vel_fft = np.zeros(len(vels[:,1]-1)/2+1)

for i in frange(n_atoms): 
  print >> sys.stderr, "fft atom %d" % i
  for j in frange(3):
    vel_fft = np.fft.rfft(vels[:,3*(i-1)+j])
    mean_vel_fft[:] += abs(vel_fft[:]*vel_fft[:])

print "# nu(1/fs)  |v(nu)^2|"
mean_vel_fft[:] /= 3*n_atoms
ii = 0
for vel_nu in mean_vel_fft:
  print "%f %.10f" % (ii/(3*n_frames*dt), vel_nu)
  ii += 1
