from ctypes import *

#TODO make a numpy version for coordinates array
class frame:
    #types
    rvec=c_float*3
    matrix=c_float*3*3

    #variables
    natoms = c_int()
    step = c_int()
    time = c_float()
    prec = c_float()
    box = matrix()
    #x = rvec*natoms - created in __init__

    def __init__(self,n):
        self.natoms=n
        xt=self.rvec*self.natoms.value  
        self.x=xt() 



class xdrfile:
    exdrOK, exdrHEADER, exdrSTRING, exdrDOUBLE, exdrINT, exdrFLOAT, exdrUINT, exdr3DX, exdrCLOSE, exdrMAGIC, exdrNOMEM, exdrENDOFFILE, exdrNR = range(13)

    def __init__(self,fn):
        self.xdr=cdll.LoadLibrary("libxdrfile.so")
        self.xd = self.xdr.xdrfile_open(fn,"r")
        if not self.xd: raise IOError("Cannot open file: '%s'"%fn)
        natoms=c_int()
        r=self.xdr.read_xtc_natoms(fn,byref(natoms))
        if r!=self.exdrOK: raise IOError("Error reading: '%s'"%fn)
        self.f = frame(natoms)
    def __iter__(self):
        f = self.f
        while 1:
            result = self.xdr.read_xtc(self.xd,f.natoms,byref(f.step),byref(f.time),f.box,f.x,byref(f.prec))
            if result==self.exdrENDOFFILE: break
            if result!=self.exdrOK: IOError("Error reading xtc file")
            yield f