"""Generate HDF5 optical structures for FDTD simulations"""

import Numeric,math,types,tables

_array=Numeric.array
_sqrt=math.sqrt
_resize=Numeric.resize

def _alwaystrue(x,y,z):
    return True

def _spherefunc(x0,y0,z0,r,x,y,z):
    return ((x0-x)**2+(y0-y)**2+(z0-z)**2<r**2)

def _ellipsefunc(x0,y0,rx,ry,x,y):
    if(rx<ry):
        # vertical ellipse
        f1x=x0
        f2x=x0
        f1y=y0+_sqrt(ry**2-rx**2)
        f2y=y0-_sqrt(ry**2-rx**2)
        r=ry
    else:
        # horizontal ellipse
        f1x=x0+_sqrt(rx**2-ry**2)
        f2x=x0-_sqrt(rx**2-ry**2)
        f1y=y0
        f2y=y0
        r=rx
    return (_sqrt((x-f1x)**2+(y-f1y)**2)+_sqrt((x-f2x)**2+(y-f2y)**2) < 2*r)

class IndexFile:
    """Workspace class to build HDF5 index files

This class defines a workspace in which you can put geometric
structures, and finally you can write them to a HDF5 file."""
    def __init__(self,lx,ly,lz,index=1):
        """Initialize an IndexFile structure

lx, ly, lz: dimensions of the workspace, in metric units
index: the default index in this workspace (default: 1)"""
        
        self.lx=float(lx)
        self.ly=float(ly)
        self.lz=float(lz)
        self.index=float(index)
        self.structlist=[]
        
    def write(self,filename,delta_x=1.,oversampling=4):
        """Actually write the data to a HDF5 file

filename: name of the HDF5 file to be created
delta_x: space step of the final grid (default=1.)
oversampling: the grid is subdivized in a higher resolution grid
              to smooth the borders. This parameter sets how much
              better it should be (default: 4)"""

        delta_x=float(delta_x)
        nx=int(math.ceil(self.lx/delta_x))
        ny=int(math.ceil(self.ly/delta_x))
        nz=int(math.ceil(self.lz/delta_x))
        indexarr=_resize(float(self.index),(nx,ny,nz))
        absorarr=None
        do_absor=False
        for struct in self.structlist:
            if(struct.absor):
                do_absor=True
                absorarr=_resize(0.,(nx,ny,nz))
                break
	
        for i in range(nx):
            x=i*delta_x
            try:
                for struct in self.structlist:
                    assert x+delta_x<struct.x1 or x>struct.x2
            except AssertionError:            
                for j in range(ny):
                    y=j*delta_x
                    try:
                        for struct in self.structlist:
                            assert y+delta_x<struct.y1 or y>struct.y2
                    except AssertionError:
                        for k in range(nz):
                            z=k*delta_x
                            try:
                                for struct in self.structlist:
                                    assert z+delta_x<struct.z1 or z>struct.z2
                            except AssertionError:
                                tmpdata=_resize(float(self.index),
                                   (oversampling,oversampling,oversampling))
                                if(do_absor):
                                    tmpdata2=_resize(0.,
                                       (oversampling,oversampling,oversampling))
                                    for struct in self.structlist:
                                        struct.local_index_absor_fill(x,y,z,delta_x,
                                                                      oversampling,
                                                                      tmpdata,
                                                                      tmpdata2)
                                    absorarr[i][j][k]=sum(sum(sum(tmpdata2)))/oversampling**3
                                else:
                                    for struct in self.structlist:
                                        struct.local_index_fill(x,y,z,delta_x,
                                                                oversampling,
                                                                tmpdata)
                                indexarr[i][j][k]=sum(sum(sum(tmpdata)))/oversampling**3
	
        if(filename[-3:]!='.h5'): filename+='.h5'
        f=tables.openFile(filename,mode='w')
	f.createArray(f.root, 'index', indexarr)
        if(do_absor):
            f.createArray(f.root, 'absor', absorarr)
        f.close()

    def addobj(self,struct):
        """Add a geometric object to the structure

struct: the object to add (must be a GeomStruct instance)"""
        assert isinstance(struct,GeomStruct)
        self.structlist.append(struct)



class GeomStruct:
    """Generic geometric object class"""
    def __init__(self,func,x1,x2,y1,y2,z1,z2,index=1,absor=0,check_bounds=False):
    	"""Create a generic geometric object

func: a 3-argument function describing the object. func(x,y,z)
      must be True when inside the object, and False outside
x1,x2,y1,y2,z1,z2: bounds outside which func(x,y,z) is always
                   False, i.e. a bounding box surrounding the
                   object
index: the refractive index inside the object (default=1.)
absor: the absorption coefficient inside the object (default=0.)
check_bounds: set it to True when func(x,y,z) is not guaranteed
              to be False outside the bounding box. E.g. if
              func(x,y,z) is always True you can use it to
              build a block (default=False)"""
        self.index=float(index)
        self.absor=float(absor)
        assert type(func) is types.FunctionType
        self.func=func
        self.x1=float(x1)
        self.x2=float(x2)
        self.y1=float(y1)
        self.y2=float(y2)
        self.z1=float(z1)
        self.z2=float(z2)
        self.check_bounds=check_bounds

    def isinside(self,x,y,z):
    	"""Check whether a given coordinate is inside the geometric object

x,y,z: coordinates of where to check"""
        if(self.check_bounds):
            if(x<self.x1 or x>=self.x2
               or y<self.y1 or y>=self.y2
               or z<self.z1 or z>=self.z2):
                return False
        return self.func(x,y,z)

    def local_index_fill(self,x0,y0,z0,delta_x,nb,ar):
    	"""Fill an cubic array with the index of the object at the
coordinates where we are inside.

x0,y0,z0: coordinates of the starting corner of the array
delta_x: the space step between two consecutive positions in the
         array
nb: number of cells in each direction
ar: the array to be filled, of size nb^3"""
        if(x0+1<self.x1 or x0>self.x2 or
           y0+1<self.y1 or y0>self.y2 or
           z0+1<self.z1 or z0>self.z2):
            return
        for i in range(nb):
            x=x0+i*delta_x/nb
            for j in range(nb):
                y=y0+j*delta_x/nb
                for k in range(nb):
                    z=z0+k*delta_x/nb
                    if self.isinside(x,y,z):
                        ar[i,j,k]=self.index
    
    def local_index_absor_fill(self,x0,y0,z0,delta_x,nb,ar,ar2):
        """Fill an cubic array with the index of the object at the
coordinates where we are inside, and another array with the absorption
coefficient.

x0,y0,z0: coordinates of the starting corner of the array
delta_x: the space step between two consecutive positions in the
         array
nb: number of cells in each direction
ar: the array to be filled with the index, of size nb^3
ar2: the array to be filled with the absorption coefficient, same size"""
        if(x0+1<self.x1 or x0>self.x2 or
           y0+1<self.y1 or y0>self.y2 or
           z0+1<self.z1 or z0>self.z2):
            return
        for i in range(nb):
            x=x0+i*delta_x/nb
            for j in range(nb):
                y=y0+j*delta_x/nb
                for k in range(nb):
                    z=z0+k*delta_x/nb
                    if self.isinside(x,y,z):
                        ar[i,j,k]=self.index
                        ar2[i,j,k]=self.absor


class Block(GeomStruct):
    """Block object class"""
    def __init__(self,x1,x2,y1,y2,z1,z2,index=1,absor=0):
    	"""Create and initialize a block object

x1,x2,y1,y2,z1,z2: bounding coordinates for the block
index: refractive index inside the block (default=1.)
absor: the absorption coefficient inside the object (default=0.)"""
        GeomStruct.__init__(self,_alwaystrue,
                            x1,x2,y1,y2,z1,z2,index=index,absor=absor,check_bounds=True)

class Sphere(GeomStruct):
    """Sphere object class"""
    def __init__(self,x0,y0,z0,r,index=1,absor=0):
    	"""Create and initialize a sphere object
    	
x0,y0,z0: center of the sphere
r: radius of the sphere
index: refractive index inside the sphere (default=1.)
absor: the absorption coefficient inside the object (default=0.)"""
        self.center=(x0,y0,z0)
        self.radius=r
        GeomStruct.__init__(self,(lambda x,y,z:_spherefunc(x0,y0,z0,r,x,y,z)),
                            x0-r,x0+r,y0-r,y0+r,z0-r,z0+r,index=index,absor=absor)

# We define a prism by:
# * a projection plane: x, y or z (c=0)
# * an axis - by default it is orthogonal to the plane
# * a 2D function of the remaining variables in this plane
# * the limits of this function *in the plane* a1<a<a2, b1<b<b2
#   with (a,b) = (y,z) or (x,z) or (x,y) depending on the projection
# * the height of the prism, h = c2-c1

class Prism(GeomStruct):
    """Generic prism object class"""
    def __init__(self,func,a1,a2,b1,b2,c1,c2,
                 projplane='z',index=1,absor=0,axis=None):
        """Create and initialize a generic prism object

func: a 2-variable function returning True when inside the
      object, with the two arguments being the coordinates
      inside the projection plane
a1,a2,b1,b2: the bounding box outside which func(a,b)
             is always False.
c1,c2: the limits of our prism for the coordinate orthogonal
       to the projection plane
projplane: the projection plane, defined by its normal
           direction. Must be either of 'x', 'y', or 'z'
           (default='z')
index: the refractive index inside the structure (default=1.)
absor: the absorption coefficient inside the object (default=0.)
axis: a tuple of length 3, defining a vector outside the
      projection plane, which is the projection direction
      (default= normal to the projection plane)"""
        assert type(func) is types.FunctionType
        self.func2d=func
	xcoef,ycoef,zcoef=0.,0.,0.

        # check the axis is valid
        assert projplane in ('x','y','z')
        foo={'x':(1,0,0),'y':(0,1,0),'z':(0,0,1)}
        if(axis==None): axis=foo[projplane]
        assert type(axis) is types.TupleType and len(axis)==3
        assert map((lambda a,b:a*b),axis,foo[projplane]) != [0,0,0]
        axis=map(float,axis)

        if(projplane=='x'):
            ycoef=axis[1]/axis[0]
            zcoef=axis[2]/axis[0]
            x1,x2=c1,c2
            y1=a1+min(x1*ycoef,x2*ycoef)
            y2=a2+max(x1*ycoef,x2*ycoef)
            z1=b1+min(x1*zcoef,x2*zcoef)
            z2=b2+max(x1*zcoef,x2*zcoef)
            f=lambda x,y,z: self.func2d(y-x*ycoef,z-x*zcoef)
        elif(projplane=='y'):
            xcoef=axis[0]/axis[1]
            zcoef=axis[2]/axis[1]
            y1,y2=c1,c2
            x1=a1+min(y1*xcoef,y2*xcoef)
            x2=a2+max(y1*xcoef,y2*xcoef)
            z1=b1+min(y1*zcoef,y2*zcoef)
            z2=b2+max(y1*zcoef,y2*zcoef)
            f=lambda x,y,z: self.func2d(x-y*xcoef,z-y*zcoef)
        else: # projplane == 'z'
            xcoef=axis[0]/axis[2]
            ycoef=axis[1]/axis[2]
            z1,z2=c1,c2
            x1=a1+min(z1*xcoef,z2*xcoef)
            x2=a2+max(z1*xcoef,z2*xcoef)
            y1=b1+min(z1*ycoef,z2*ycoef)
            y2=b2+max(z1*ycoef,z2*ycoef)
            f=lambda x,y,z: self.func2d(x-z*xcoef,y-z*ycoef)
        GeomStruct.__init__(self,f,x1,x2,y1,y2,z1,z2,index=index,absor=absor,
                            check_bounds=True)

class Cylinder(Prism):
    """Cylinder object class"""
    def __init__(self,a0,b0,ra,rb,c1,c2,
                 projplane='z',index=1,absor=0,axis=None):
        """Create and initialize a cylinder object

a0,b0: the coordinates of the center of the cylinder in the
       projection plane
ra,rb: the two radii of the ellipse centered at (a0,b0). If
       rb=None, it is taken equal to ra and we have a circle
c1,c2: the limits of our cylinder for the coordinate orthogonal
       to the projection plane
projplane: the projection plane, defined by its normal
           direction. Must be either of 'x', 'y', or 'z'
           (default='z')
index: the refractive index inside the structure (default=1.)
absor: the absorption coefficient inside the object (default=0.)
axis: a tuple of length 3, defining a vector outside the
      projection plane, which is the projection direction
      (default= normal to the projection plane)"""
        if rb==None: rb=ra
        Prism.__init__(self,(lambda a,b: _ellipsefunc(a0,b0,ra,rb,a,b)),
                       a0-ra,a0+ra,b0-rb,b0+rb,c1,c2,projplane=projplane,
                       index=index,absor=absor,axis=axis)

def _coneprojfunc(a,b,c,a0,b0,c0,c1,func2d):
    if c==c0:
        return True
    return func2d(a0+(a-a0)*(c1-c0)/(c-c0),b0+(b-b0)*(c1-c0)/(c-c0))

class GenericCone(GeomStruct):
    """Generic cone object class"""
    def __init__(self,func,a1,a2,b1,b2,c1,c2,x0,y0,z0,
                 projplane='z',index=1,absor=0):
        """Create and initialize a generic cone object

func: a 2-variable function returning True when inside the
      object, with the two arguments being the coordinates
      inside the projection plane (the c=c1 plane)
a1,a2,b1,b2: the bounding box outside which func(a,b)
             is always False
c1,c2: the limits of our prism for the coordinate orthogonal
       to the projection plane; c1 is taken as the plane
       where we calculate the 2D function
x0,y0,z0: the position of the vertex
projplane: the projection plane, defined by its normal
           direction. Must be either of 'x', 'y', or 'z'
           (default='z')
index: the refractive index inside the structure (default=1.)
absor: the absorption coefficient inside the object (default=0.)"""
        assert type(func) is types.FunctionType
        self.func2d=func
        assert projplane in ('x','y','z')
        c1=float(c1)
        c2=float(c2)
        if(projplane=='x'):
            x1=min(c1,c2)
            x2=max(c1,c2)
            lis=(a1,a2,y0+(a1-y0)*(c2-x0)/(c1-x0),y0+(a2-y0)*(c2-x0)/(c1-x0))
            y1=min(lis)
            y2=max(lis)
            lis=(b1,b2,z0+(b1-z0)*(c2-x0)/(c1-x0),z0+(b2-z0)*(c2-x0)/(c1-x0))
            z1=min(lis)
            z2=max(lis)
            f=lambda x,y,z: _coneprojfunc(y,z,x,y0,z0,x0,c1,self.func2d)
        elif(projplane=='y'):
            y1=min(c1,c2)
            y2=max(c1,c2)
            lis=(a1,a2,x0+(a1-x0)*(c2-y0)/(c1-y0),x0+(a2-x0)*(c2-y0)/(c1-y0))
            x1=min(lis)
            x2=max(lis)
            lis=(b1,b2,z0+(b1-z0)*(c2-y0)/(c1-y0),z0+(b2-z0)*(c2-y0)/(c1-y0))
            z1=min(lis)
            z2=max(lis)
            f=lambda x,y,z: _coneprojfunc(x,z,y,x0,z0,y0,c1,self.func2d)
        else: # projplane == 'z'
            z1=min(c1,c2)
            z2=max(c1,c2)
            lis=(a1,a2,x0+(a1-x0)*(c2-z0)/(c1-z0),x0+(a2-x0)*(c2-z0)/(c1-z0))
            x1=min(lis)
            x2=max(lis)
            lis=(b1,b2,y0+(b1-y0)*(c2-z0)/(c1-z0),y0+(b2-y0)*(c2-z0)/(c1-z0))
            y1=min(lis)
            y2=max(lis)
            f=lambda x,y,z: _coneprojfunc(x,y,z,x0,y0,z0,c1,self.func2d)
        GeomStruct.__init__(self,f,x1,x2,y1,y2,z1,z2,index=index,absor=absor,
                            check_bounds=True)

class Cone(GenericCone):
    """Ellipsoidal cone object class"""
    def __init__(self,a0,b0,ra,rb,c1,c2,cvertex,
                 avertex=None,bvertex=None,
                 projplane='z',index=1,absor=0):
        """Create and initialize an ellipsoidal cone object

a0,b0: the coordinates of the center of the ellipse in the
       projection plane
ra,rb: the two radii of the ellipse centered at (a0,b0). If
       rb=None, it is taken equal to ra and we have a circle
c1,c2: the limits where to truncate the cone for the coordinate
       orthogonal to the projection plane; c1 is the coordinate
       where we calculate the ellipse
cvertex: the coordinate of the vertex in the direction orthogonal
         to the projection plane
avertex,bvertex: same in the other directions (default to a0,b0)
projplane: the projection plane, defined by its normal
           direction. Must be either of 'x', 'y', or 'z'
           (default='z')
index: the refractive index inside the structure (default=1.)
absor: the absorption coefficient inside the object (default=0.)"""
        if rb==None: rb=ra
        if avertex==None: avertex=a0
        if bvertex==None: bvertex=b0
        assert projplane in ('x','y','z')
        if(projplane=='x'):
            x0,y0,z0=cvertex,avertex,bvertex
        elif(projplane=='y'):
            x0,y0,z0=avertex,cvertex,bvertex
        else:
            x0,y0,z0=avertex,bvertex,cvertex
        GenericCone.__init__(self,(lambda a,b: _ellipsefunc(a0,b0,ra,rb,a,b)),
                             a0-ra,a0+ra,b0-rb,b0+rb,c1,c2,x0,y0,z0,
                             projplane=projplane,index=index,absor=absor)