"""
By Kirby Urner, 4D Solutions

Last modified: February, 2001
      Converted to subclass of Polyhedron to facilitate
      manipulation of watermans as generic polys
Original version:  January, 2001
"""

import operator
import coords,povray,vrml,lg3d

# Note: polyhedon depends on additional custom modules
# and should be edited to have pathnames appropriate for your setup
# polyhedron also makes use of Qhull, a 3rd party program

from polyhedron import Polyhedron

class W(Polyhedron):

    root    = 0
    type    = 1
    sphradius = 0.08
    cylradius = 0.02
    qhull = 1
    scalefactor = 1/(2.0)**0.5
    radius = 0
    
    def __init__(self,root,type=1,format="pov"):
        """
        Generate a povray file for W(root) -- this is the main program
        If type is not 'w', then fixed radius option is invoked
        """
        self.format = format
        self.root = root
        self.radius = root**0.5
        
        allcoords = []
        poscoords = []
        
        if type==1:
            for i in range(10):
               if root>=(i+1):
                   poscoords += isoradial(root-i)
            self.filename = "wm"+str(root)
        else:
            poscoords = isoradial(root)            
            self.filename = "ir"+str(root)

        if len(poscoords)==0:
            print "Missing for root %s\n" % root
            return 0

        for v in poscoords:        
            allcoords = allcoords + negperm(v)

        self.verts = map(coords.Vector,allcoords)
        
        self.writepoints() 
        print "Raw data generated for root %s" % root

        self.genhull()
        self.getdata()
        
        if self.type==1:
            self.colorcode()        

        for i in range(len(self.verts)):
            self.verts[i] *= self.scalefactor
            
    def colorcode(self):  
        print "Color coding vertices"
        for i in range(len(self.verts)):
           vlen = self.verts[i].dot(self.verts[i])
           if  vlen == 2*self.root:
               self.vcolors[i] = "Gold"
           elif vlen == 2*(self.root-1):               
               self.vcolors[i] = "Silver"
           elif vlen == 2*(self.root-2):
               self.vcolors[i] = "Green"
           elif vlen == 2*(self.root-3):
               self.vcolors[i] = "Blue"
           elif vlen == 2*(self.root-4):
               self.vcolors[i] = "Yellow"
           elif vlen == 2*(self.root-5):
               self.vcolors[i] = "Magenta"
           elif vlen == 2*(self.root-6):
               self.vcolors[i] = "LightSteelBlue"
           elif vlen == 2*(self.root-7):
               self.vcolors[i] = "Pink"
           elif vlen == 2*(self.root-8):
               self.vcolors[i] = "DarkTan"
           elif vlen == 2*(self.root-9):
               self.vcolors[i] = "DustyRose"
           else:
               self.verts[i].color = "Black"

def mkwater(root,type=1,format="pov"):
    wpoly  = W(root,type,format)
    wpoly  *= 1./(root**0.5)
    if wpoly.format=="pov":    
        myfile = povray.Povray(wpoly.filename+"."+format,bgcolor="Black")
    if wpoly.format=="wrl":
       myfile  = vrml.Vrml(wpoly.filename+"."+format,bgcolor="Black")
    if wpoly.format=="m":
       myfile  = lg3d.Lg3d(wpoly.filename+"."+format,bgcolor="Black",box=wpoly.radius*1.5)
       wpoly.showedges = 0
       wpoly.cylradius = 0.012
       wpoly.sphradius = 0.03
    wpoly.display(myfile)
    myfile.close()

def mkcouple(roots,type=1,format="pov"):
    wpoly1  = W(roots[0],type,format)
    wpoly1.sphradius = 0.05
    if wpoly1.format=="pov":
       myfile  = povray.Povray("cp."+format,bgcolor="Black")
    if wpoly1.format=="wrl":
       myfile  = vrml.Vrml("cp."+format,bgcolor="Black")
    if wpoly1.format=="m":
       myfile  = lg3d.Lg3d("cp."+format,bgcolor="Black")        
    wpoly2  = W(roots[1],type)
    wpoly2  *= 1./(roots[1]**0.5)
    wpoly1  *= 1./(roots[1]**0.5)
    wpoly2.showfaces = 0
    wpoly2.cylcolor = "Green"
    wpoly2.sphradius = 0.05
    wpoly1.display(myfile)
    wpoly2.display(myfile)
    myfile.close()
    
def isoradial(root,perm=1):
    """
    List all-positive, integral (x,y,z) positions at the centers
    of CCP spheres that are distance pow(2*root,0.5) from the
    origin -- optionally suppress permutations with 2nd argument
    of 0.

    CCP spheres have diameter root(2) in this coordinate system,
    so 2nd powers of sphere-center distances from (0,0,0) will
    be multiples of 2, i.e. pow(N*root(2),2) = 2*(N**2).  N**2
    = root, the first argument to this method.  A scale factor
    of 1/root(2) may be applied later such that root = 3 results
    in CCP spheres a distance square root of 3 from (0,0,0).

    x + y + z = even number (i.e. (x+y+z)%2==0.  This stipulation
    ensures that our integer (x,y,z) coordinates are coincident
    with the center of a CCP sphere.

    See: http://www.inetarena.com/~pdx4d/ocn/wmalgorithm.html for
    additional background information
    """
    rad = root*2
    coords = []
    for x in range(int(pow(rad,.5)),-1,-1):  # descending order
        if x**2 < rad/3:
            break
        for y in range(x+1):
            if x**2 + y**2 > rad:
                break
            for z in range(y+1):
                if (x+y+z)%2 == 0:  # check for even sum
                    if x**2 + y**2 + z**2 > rad:
                        break                    
                    if x**2 + y**2 + z**2 == rad:
                        coords.append((x,y,z))
                        
    if len(coords)==0: return coords
    elif perm: return reduce(operator.add,map(permute,coords))
    else:      return coords

def permute(e):
    """
    return list of unique permutations of 3 elements
    """
    result = []    
    for i in range(3):
        for j in range(3):
            for k in range(3):
                if i<>j and j<>k and i<>k:
                   permutation = (e[i],e[j],e[k])
                   if not permutation in result:
                       result.append(permutation)
    return result

def negperm(i):
    """
    return list of tuples permuted with negative signs
    """
    def addperm(r,i):
        if not i in r:
            r.append(i)
        return r
    
    r = [i]
    r = addperm(r,( i[0],-i[1], i[2]))
    r = addperm(r,( i[0], i[1],-i[2]))
    r = addperm(r,( i[0],-i[1],-i[2]))    
    r = addperm(r,(-i[0], i[1], i[2]))
    r = addperm(r,(-i[0],-i[1], i[2]))
    r = addperm(r,(-i[0], i[1],-i[2]))
    r = addperm(r,(-i[0],-i[1],-i[2]))

    return r