#Module for reading FreeSurfer Surface files

#This module developed by Will Woods
# Last edit 24 November 2008

from struct import *
from numpy import *
import vtk


def read_fs_surf_orig(filename, numfields=4):
    fid = open(filename, 'rb')
    b1,b2,b3 = unpack('BBB', fid.read(calcsize('BBB')))
    magic = (b1<<16) + (b2<<8) + b3
    if magic==16777214:
        fid.readline()
        fid.readline()
        v,t = unpack('ii', fid.read(calcsize('ii')))
        vertices = reshape(fromfile(fid, dtype=float32, count=3*v),(-1,3))
        if numfields == 4:
            tmp = reshape(fromfile(fid, dtype=int, count=3*t),(-1,3))
            triangles = ones((tmp.shape[0],4),dtype=int)*3
            triangles[:,1:] = tmp
        fid.close()
    else:
        raise IOError('Unable to read file "%s"' % filename)

    return vertices, triangles




def read_fs_ascii(fname):

    f = open(fname)
    d = f.readlines()
    f.close()
    nVertices = int(d[1].split()[0])
    p = d[2:nVertices+2]
    pv = [ [float(y) for y in z.split()] for z in p]
    vertices = array(pv, dtype=float32)
    t = d[nVertices+2:]
    tv = [ [int(y) for y in z.split()] for z in t]
    triangles = array(tv, dtype=int)
    nTriangles = len(tv)
    
    vertices = vertices[:,:-1]
    
    t2 = empty((triangles.shape[0],triangles.shape[1]), dtype='int')
    t2[:,1:4] = triangles[:,0:3]
    
    return vertices, t2 



def read_fs_surf(filename, numfields=4):
    #TODO this routine assumes input file are big endian as the example files are
    
    fid = open(filename, 'rb')
    b1,b2,b3 = unpack('>BBB', fid.read(calcsize('>BBB')))
    magic = (b1<<16) + (b2<<8) + b3
    
    if magic==16777214:
        fid.readline()
        fid.readline()
        v,t = unpack('>ii', fid.read(calcsize('>ii')))
        v2, t2 = v*3, t*3 
        s = unpack('>%df' %v2, fid.read(calcsize('>%df' %v2)))
        
        vertices = reshape(s,(-1,3))
        
        if numfields == 4:            
            tmp = unpack('>%di' %t2, fid.read(calcsize('>%di' %t2)))
            tmp2 = reshape(tmp, (-1,3))
            
            triangles = ones((tmp2.shape[0],4),dtype=int)*3
            triangles[:,1:] = tmp2
        fid.close()
        
    else:
        raise IOError('Unable to read file "%s"' % filename)

    return vertices, triangles



def get_surf_polyData(vertices, triangles):

    # Process vertices
    points = vtk.vtkPoints()
    for i in range(vertices.shape[0]):
        #points.InsertPoint(i,-vertices[i,0]+88,(vertices[i,1]+128),(vertices[i,2]+128)) #others
        #points.InsertPoint(i,-vertices[i,0]+88,(vertices[i,1]+128)*1.13,(vertices[i,2]+128)*1.13) #ag
        points.InsertPoint(i,vertices[i,0],vertices[i,1],vertices[i,2]) #ar - YNIC_FS

    # Process faces (triangles)
    polys = vtk.vtkCellArray()
    nTriangles = triangles.shape[0]
    for i in range(nTriangles):
        polys.InsertNextCell(3)
        for j in range(1,4):
            polys.InsertCellPoint(triangles[i,j])

    # Assemble as PolyData
    polyData = vtk.vtkPolyData()
    polyData.SetPoints(points)
    polyData.SetPolys(polys)

    return polyData
