#!/usr/bin/env python
## Automatically adapted for numpy.oldnumeric Sep 06, 2007 by alter_code1.py

#!

"""
from example by Tarn Weisner Burton <twburton@users.sourceforge.net> in pyopengl
"""
# This is statement is required by the build system to query build info
if __name__ == '__build__':
    raise Exception


import string
__version__ = string.split('$Revision: 2867 $')[1]
__date__ = string.join(string.split('$Date: 2008-05-22 16:42:08 +0200 (Thu, 22 May 2008) $')[1:3], ' ')
__author__ = 'Jon Wright <jpwright@users.sourceforge.net> from example by Tarn Weisner Burton <twburton@users.sourceforge.net>'

try:
    import numpy
except:
    import sys
    print "This demo requires the numpy extension, sorry."
    sys.exit()


import OpenGL.GL as GL
import OpenGL.Tk as Tk


import sys


class plot3d(Tk.Toplevel):
    def __init__(self,parent,data=None,lines=None,
                 ubis=None,image=None,pars=None,spline=None):
        """
        Data would be your observed g-vectors. Lines will
        be a computed lattice
        """
        Tk.Toplevel.__init__(self,parent)
        self.parent=parent
        if data!=None:
            xyz=data.copy()
        else:
            xyz=numpy.array([0,0,0])
        self.ps=Tk.StringVar()
        self.ps.set('1.')
        self.pointsize=1.
        self.npeaks=xyz.shape[0]

        self.o = Tk.Opengl(self, width = 400, height = 400, double = 1)
        self.o.redraw = self.redraw
        self.o.autospin_allowed = 1
        self.o.fovy=5
        self.o.near=1e6
        self.o.far=1e-6
        import math
        self.o.distance=numpy.maximum.reduce(numpy.ravel(xyz))*4 / \
            math.tan(self.o.fovy*math.pi/180)
        print type(xyz),xyz.dtype.char,xyz.shape
        self.xyz=xyz
        f=Tk.Frame(self)
        Tk.Button(f,text="Help",command=self.o.help).pack(side=Tk.LEFT)
        Tk.Button(f,text="Reset",command=self.o.reset).pack(side=Tk.LEFT)
        Tk.Button(f,text="Pointsize",command=self.setps).pack(side=Tk.LEFT)
        Tk.Entry(f,textvariable=self.ps).pack(side=Tk.LEFT)
        Tk.Button(f,text="Quit",command=self.goaway).pack(side=Tk.RIGHT)
        self.dataoff=0
        self.o.pack(side = 'top', expand = 1, fill = 'both')
        f.pack(side=Tk.BOTTOM,expand=Tk.NO,fill=Tk.X)
        Tk.Label(self,text="Red=[1,0,0] Green=[0,1,0] Blue=[0,0,1]").pack(
            side=Tk.BOTTOM,expand=Tk.NO,fill=Tk.X)
        self.ubis=ubis
        self.color=numpy.ones((xyz.shape[0],3),numpy.float)
        print self.color.shape
        self.tex=False
        if ubis is not None:
           self.ubis = self.readubis(ubis)
           self.scorecolor(0)
        if pars is not None:
           self.tex=True
           self.readspline(spline)
           self.readprms(pars)
           self.readimage(image)
        self.changedata(xyz)

    def readspline(self,spline):
        from ImageD11 import blobcorrector
        self.corrector = blobcorrector.correctorclass(spline)

    def readubis(self,ubis):
        from ImageD11 import indexing
        return indexing.readubis(ubis)

    def readprms(self,prms):
        from ImageD11 import parameters
        o = parameters.parameters()
        o.loadparameters(prms)
        self.pars=o.get_parameters()

    def readimage(self,image):
        from ImageD11 import opendata, transform
        self.imageobj=opendata.opendata(image)
        # map from 2048x2048 to 1024x1024
        mi=1000
        mx=1500
        shape=self.imageobj.data.shape
        d=numpy.reshape(numpy.clip(self.imageobj.data,mi,mx),shape) # makes a clipped copy
        d=(255.*(d-mi)/(mx-mi)) # scale intensity
        self.image=numpy.zeros((1024,1024),numpy.UInt8)
        if d.shape==(2048,2048):
            # rebin 2x2
            im=(d[::2,::2]+d[::2,1::2]+d[1::2,::2]+d[1::2,1::2])/4
            self.image=(255-im).astype(numpy.UInt8).tostring()
        self.imageWidth=1024
        self.imageHeight=1024
        # make a 2D array of x,y
        p=[]
        pk=[]
        step = 64
        r=[ [ 0,0 ], [0,step], [step,step], [step,0] ]
        for i in range(0,1024,step):
            for j in range(0,1024,step):
                # i,j 1024x1024 texture coords
                # x,y spatially corrected
                for v in r:
                    pk.append([i+v[0],j+v[1]])
                    x,y = self.corrector.correct((i+v[0])*2 , (j+v[1])*2) # corrected
                    p.append([x,y])
        p=numpy.array(p).T 
        pk=numpy.array(pk).T
        omega=float(self.imageobj.header['Omega'])
        self.pars['distance']=float(self.pars['distance'])*1000
        tth,eta=transform.compute_tth_eta(p,**self.pars)
        gve = transform.compute_g_vectors(tth,eta,omega*self.pars['omegasign'],self.pars['wavelength'])
        self.pts = []
        print "Setting up image mapping",p.shape,gve.shape
        for i in range(pk.shape[1]):
            self.pts.append([pk[1,i]/1024.,pk[0,i]/1024.,gve[0,i],gve[1,i],gve[2,i]])
        #for p in self.pts:
        #    print p
        self.setupTexture()

    def setupTexture(self):
        GL.glDisable(GL.GL_TEXTURE_2D)
        GL.glPixelStorei(GL.GL_UNPACK_ALIGNMENT, 1)
        GL.glTexImage2D(GL.GL_TEXTURE_2D,#target
                    0,#level
                    3,#internalformat
                    self.imageWidth, self.imageHeight,
                    0,#border
                    GL.GL_LUMINANCE,#format
                    GL.GL_UNSIGNED_BYTE,# type
                    self.image)
        GL.glTexParameterf(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_WRAP_S, GL.GL_CLAMP)
        GL.glTexParameterf(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_WRAP_T, GL.GL_CLAMP)
        GL.glTexParameterf(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_WRAP_S, GL.GL_REPEAT)
        GL.glTexParameterf(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_WRAP_T, GL.GL_REPEAT)
        GL.glTexParameterf(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_MAG_FILTER, GL.GL_NEAREST)
        GL.glTexParameterf(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_MIN_FILTER, GL.GL_NEAREST)
        GL.glTexEnvf(GL.GL_TEXTURE_ENV, GL.GL_TEXTURE_ENV_MODE, GL.GL_DECAL)
        GL.glEnable(GL.GL_TEXTURE_2D)
        GL.glEnable(GL.GL_NORMALIZE)
        GL.glShadeModel(GL.GL_FLAT)



    def scorecolor(self,i=0):
        cc = [ [ 1,0,0] , [0,1,0] , [0,0,1], [1,1,0], [1,0,1], [0,1,1]]
        if self.ubis is not None:
            from ImageD11 import indexing
            for u,i in zip(self.ubis,range(len(self.ubis))):
                scores=indexing.calc_drlv2(self.ubis[i],self.xyz)
                print self.xyz.shape,scores.shape
                ind = numpy.compress( numpy.less(scores,0.02) , 
                                      numpy.arange(self.xyz.shape[0]) )
                print "Grain",i,scores.shape,ind.shape
                for j in range(3):
                    c=numpy.ones(self.color.shape[0])
                    numpy.put(c,ind,cc[i%len(cc)][j])
                    self.color[:,j]*=c

    def go(self):
        """
        Allow the toplevel to return a handle for changing data
        """
        self.o.mainloop()

    def goaway(self):
        print "Called goaway"
        self.o.destroy()
        self.destroy()
        if self.parent is None: sys.exit()
        print "Ought to be gone now..."

    def changedata(self,xyz):
        self.xyz=xyz.copy()
        self.npeaks=xyz.shape[0]
        GL.glDisableClientState(GL.GL_VERTEX_ARRAY)
        GL.glDisableClientState(GL.GL_COLOR_ARRAY)
        GL.glVertexPointerd(self.xyz)
        GL.glColorPointerd(self.color)
        GL.glEnableClientState(GL.GL_VERTEX_ARRAY)
        GL.glEnableClientState(GL.GL_COLOR_ARRAY)
        self.o.tkRedraw()

    def setps(self):
        self.pointsize=float(self.ps.get())
        self.o.tkRedraw()


    def redraw(self,o):
        
        GL.glDisable(GL.GL_LIGHTING)
        GL.glClearColor(0., 0., 0., 0)
        GL.glClear(GL.GL_COLOR_BUFFER_BIT | GL.GL_DEPTH_BUFFER_BIT)
        GL.glOrtho(-1,1,-1,1,-1,1)
        GL.glColor3f(1.0, 1.0, 1.0) # white
        GL.glPointSize(self.pointsize)
        GL.glDrawArrays(GL.GL_POINTS, 0, self.npeaks )

        if self.ubis is not None and len(self.ubis)==1:
            hkl = numpy.dot(numpy.linalg.inv(self.ubis[0]), 
                            numpy.identity(3,numpy.float)).T
            # print hkl
        else:
            hkl = numpy.identity(3,numpy.float)
            # print hkl
                            
        GL.glBegin(GL.GL_LINE_LOOP)
        GL.glColor3f(1.0, 0.0, 0.0) # red
        GL.glVertex3f(0.,0.,0.)
        GL.glVertex3f(hkl[0][0],hkl[0][1],hkl[0][2])
        GL.glEnd()
        GL.glBegin(GL.GL_LINE_LOOP)
        GL.glColor3f(0.0, 1.0, 0.0) # green
        GL.glVertex3f(0.,0.,0.)
        GL.glVertex3f(hkl[1][0],hkl[1][1],hkl[1][2])
        GL.glEnd()
        GL.glBegin(GL.GL_LINE_LOOP)
        GL.glColor3f(0.0, 0.0, 1.0) # blue
        GL.glVertex3f(0.,0.,0.)
        GL.glVertex3f(hkl[2][0],hkl[2][1],hkl[2][2])

        GL.glEnd()
        if self.tex:
#            print "drawing images"
            GL.glEnable(GL.GL_TEXTURE_2D)
            GL.glColor4f(.0, 1.0, .0, 1.0) # red        
            GL.glBegin(GL.GL_QUADS)
            # generate a grid of squares to map the texture in 3D
            # opengl has better "map" methods to do this
            for i,j,g1,g2,g3 in self.pts:
#                print i,j,g1,g2,g3
                GL.glTexCoord2f(i,j)
                GL.glVertex3f(g1, g2, g3) 
            GL.glEnd()
            GL.glDisable(GL.GL_TEXTURE_2D)
            
        GL.glFlush()
        GL.glEnable(GL.GL_LIGHTING)







if __name__=="__main__":

    try:
        lines=open(sys.argv[1],"r").readlines()
    except:
        print "Usage %s gvector_file [ubifile] [image parfile]"%(sys.argv[0])
        sys.exit()
   
    on=0
    xyz=[]
    for line in lines:
        if on==1:
            try:
                vals=[float(x) for x in line.split()]
                xyz.append( [ vals[0],vals[1],vals[2] ])
            except:
                pass
        if line.find("xr yr zr")>0 or line.find("gx  gy  gz  ")>0:
            on=1
   
    npeaks = len(xyz)
    xyz=numpy.array(xyz)
    if len(sys.argv)==3:
       o=plot3d(None,data=xyz,ubis=sys.argv[2])
    elif len(sys.argv)==6:
       o=plot3d(None,data=xyz,ubis=sys.argv[2],image=sys.argv[3],pars=sys.argv[4],spline=sys.argv[5])
    else:
       o=plot3d(None,data=xyz,ubis=None)
    o.mainloop()
