#!/usr/bin/python

# Speedtest, Copyright (c) 2000, Richard P. Muller
# Runs simple speed tests in python.

import time,random,math
from Speedtest import *
from Numeric import *

def calcElectroStaticEnergy(nat,charge,x,y,z,distfunc):
    es = 0.
    for i in range(nat):
        qi,xi,yi,zi = charge[i],x[i],y[i],z[i]
        for j in range(i):
            qj,xj,yj,zj = charge[j],x[j],y[j],z[j]
            rij = distfunc(xi,yi,zi,xj,yj,zj)
            es = es + qi*qj/rij
    return es

def distance(xi,yi,zi,xj,yj,zj):
    return math.sqrt((xi-xj)*(xi-xj)+(yi-yj)*(yi-yj)+(zi-zj)*(zi-zj))

def getBoxOfCharge(n,xmin,xmax,ymin,ymax,zmin,zmax):
    charge = []; x = []; y = []; z = []
    dx,dy,dz = xmax-xmin,ymax-ymin,zmax-zmin
    for i in range(n):
        charge.append(random.randint(-1,1))
        x.append(dx*random.random()-xmin)
        y.append(dy*random.random()-ymin)
        z.append(dz*random.random()-zmin)
    charge = array(charge,Int)
    x = array(x,Float)
    y = array(y,Float)
    z = array(z,Float)
    return charge,x,y,z

def main():
    print "Python Version"

    t0 = time.clock()
    n = 1000
    charge,x,y,z = getBoxOfCharge(n,-10.,10.,-10.,10.,-10.,10.)
    t1 = time.clock()
    print "Setup time: ",t1-t0

    es = calcElectroStaticEnergy(n,charge,x,y,z,distance)
    t2 = time.clock()
    print "ES Energy, time =  ", es,t2-t1

    print "\nCalling C distance Function:"
    es = calcElectroStaticEnergy(n,charge,x,y,z,cdistance)
    t3 = time.clock()
    print "ES energy, time = ", es,t3-t2

    print "\nCalling pure C function:"
    es = cCalcESEnergy(n,charge,x,y,z)
    t4 = time.clock()
    print "ES Energy, time = ", es, t4-t3

if __name__=='__main__':
    main()
    #import profile,pstats
    #profile.run('main()','es_prof')
    #p = pstats.Stats('es_prof')
    #p.strip_dirs().sort_stats('time').print_stats(15)