from Tkinter import *

from visual import *
from random import random
from time import clock

import time,sys,thread

class ProgClass:
    def __init__(self,parent):
        self.parent = parent
        self.run = "run"

        parent.wm_geometry(newGeometry="+10+10")
        parent.wm_title("Controls")

        scene.x = 200
        scene.y = 200
        scene.width = 800
        scene.height = 700
        scene.autoscale = 0
        scene.up = (0,0,1)
        scene.forward = (0,1,-.5)
        scene.lights = [vector(-.5,.5,.5), vector(.5,-.5,.5)]

        torque_toggle = 1

        thread.start_new_thread(self.lattice_simulate,())


    def lattice(self, Ntotx=8, Ntoty=8, deltax=1.0, deltay=1.0):
        spins = []
        
        xmin = -Ntotx*deltax/2.
        ymin = -Ntoty*deltay/2.
        nz = 0
    
        for nx in range(Ntotx):
          x = xmin + nx*deltax
          for ny in range(Ntoty):
              y = ymin + ny*deltay

              spins.append(frame())
              self.spins[-1].pos = vector(x,y,0)
              self.spins[-1].spin = vector()
              self.spins[-1].spinvec = arrow(pos=(x,y,0), axis=(0,0,1), color=(.9,.9,.9), shaftwidth=0.2)
          
              self.spins[-1].torque = vector()
              self.spins[-1].torquevec = arrow(pos=(x,y,0), axis=(0,0,1), color=(.3,.3,1), shaftwidth=0.05)
          
              self.spins[-1].nearx = range(2)
              self.spins[-1].neary = range(2)
              self.spins[-1].indices = (nx,ny,nz)
          
        for s in self.spins:
          nx, ny, nz = s.indices
          if nx == 0:                       # leftmost spin in a row
              nspinl = Ntotx*ny + Ntotx-1       # wrap around to spin on right side
              s.nearx[0] = nspinl               # reference this by its list element number, given by the order in which the spins are added to the list...
          else:
              nspinl = Ntotx*ny + nx-1
              s.nearx[0] = nspinl
          if nx == Ntotx-1:                 # rightmost spin in a row
              nspinr = Ntotx*ny
              s.nearx[1] = nspinr
          else:
              nspinr = Ntotx*ny + nx+1 
              s.nearx[1] = nspinr

          if ny == 0:                       # bottom spin in a column
              nspind = Ntotx*(Ntoty-1) + nx 
              s.neary[0] = nspind
          else:
              nspind = Ntotx*(ny-1) + nx
              s.neary[0] = nspind
          if ny == Ntoty-1:                 # top spin in a column
              nspinu = nx
              s.neary[1] = nspinu
          else:
              nspinu = Ntotx*(ny+1) + nx
              s.neary[1] = nspinu
        return self.spins

    def set_red(self,rstr): # set red component of color
        print ""
        
    def set_green(self,gstr): # set green component of color
        print ""
        
    def set_blue(self,bstr): # set blue component of color
        print ""

        

    def toggletorques(self):
        if torque_toggle == 1:
            torque_toggle = 0
            for s in spins:
                s.torquevec.visible = 0
                print "Torques made invisible"
        if torque_toggle == 0:
            torque_toggle = 1
            for s in spins:
                s.torquevec.visible = 1
                print "Torques made visible"


    def lattice_simulate(self):
        while(self.run == "run"):

            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
            #~~                       LATTICE CONFIGURATION                      ~~#
            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
            
            Ntotx = 5
            Ntoty = 1
            Ntotz = 1
            Ntotal = Ntotx*Ntoty*Ntotz

            if Ntotal == Ntotx or Ntotal == Ntoty:
                colormode = 1

            dx = 1
            dy = 1

            Lx = dx*Ntotx
            Ly = dy*Ntoty

            spins = self.lattice(Ntotx, Ntoty, dx, dy)

            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#



            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
            #~~                        INITIAL SPIN SETUP                        ~~#
            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#

            J = -1
            k = 2*pi*vector(10./Lx, 0./Ly, 0)

            print k.x*dx
            spin_mag_ratio = (1-sin(k.x*dx))/cos(k.x*dx)
            print spin_mag_ratio

            spin_length = 1.0

            spinA_sigma = .001
            spinB_sigma = spin_mag_ratio*spinA_sigma
            spinA_zmag = (spin_length**2-spinA_sigma**2)**(.5)
            spinB_zmag = (spin_length**2-spinB_sigma**2)**(.5)

            for s in spins:
                if ((J/abs(J))**(s.indices[0] + s.indices[1]) == 1):
                    spinx = spinA_sigma*sin(dot(k, s.pos))
                    spiny = spinA_sigma*cos(dot(k, s.pos))
                    spinz = spinA_zmag
        
                if ((J/abs(J))**(s.indices[0] + s.indices[1]) == -1):
                    spinx = -spinB_sigma*sin(dot(k, s.pos))
                    spiny = -spinB_sigma*cos(dot(k, s.pos))
                    spinz = -spinB_zmag
  
                s.spin = vector(spinx,spiny,spinz)
                s.spin = norm(s.spin)

            #    s.pos.z = s.pos.z + (-1)**(s.indices[0] + s.indices[1])

            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#



            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
            #~~             SETUP FOR VISUAL REPRESENTATION OF SPINS             ~~#
            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#

            spinvec_length = 2
            spinvec_sigma = 1
            spinvec_shaftwidth = .2
            spinvec_zmag = (spinvec_length**2-spinvec_sigma**2)**(.5)
            
            for s in spins:
                s.spinvec.axis = s.spin*spinvec_length
                s.spinvec.pos = s.pos - s.spinvec.axis/2.
                s.spinvec.shaftwidth = spinvec_shaftwidth
    
            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
    


            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
            #~~                          MAIN SIMULATION                         ~~#
            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
    
            dt = .05
            tt = clock()
            Nsteps = 0

            while 1:
                for s in spins:
                    s.torque = vector(0,0,0)

                    for nnx in range(2):
                        nspinx = s.nearx[nnx]
                        s.torque = s.torque + J*cross(s.spin,spins[nspinx].spin)
                    for nny in range(2):
                        nspiny = s.neary[nny]
                        s.torque = s.torque + J*cross(s.spin,spins[nspiny].spin)
        
                    s.torque.z = 0                                                  # fixed torque model
        
                for s in spins:
                    s.spin = s.spin + s.torque*dt

                    if ((J/abs(J))**(s.indices[0] + s.indices[1]) == 1):            # UP spin
                        s.spin = spinA_sigma*norm(vector(s.spin.x,s.spin.y,0))
                        s.spin.z = spinA_zmag
        
                    if ((J/abs(J))**(s.indices[0] + s.indices[1]) == -1):           # DOWN spin
                        s.spin = spinB_sigma*norm(vector(s.spin.x,s.spin.y,0))
                        s.spin.z = -spinB_zmag
        
                    s.spin = norm(s.spin)                                           # ensure spin remains magnitude 1

            #        s.spinvec.axis = spinvec_length*s.spin

                    s.spinvec.axis = spinvec_sigma*norm(vector(s.spin.x,s.spin.y,0)) + (spinvec_zmag/spinA_zmag)*vector(0,0,s.spin.z)                # spinvec attributes now represent a scaling
                    s.spinvec.pos = s.pos - s.spinvec.axis/2.
        
                    if colormode == 1:
                        s.spinvec.color = (.5+.5*(s.spin.z/abs(s.spin.z))*dot(norm(vector(s.spin.x,s.spin.y,0)),vector(0,-.707,0)) , 0 , .5+.5*(s.spin.z/abs(s.spin.z))*dot(norm(vector(s.spin.x,s.spin.y,0)),vector(0,.707,0)))
                    else:
                        s.spinvec.color = (.5+.5*(s.spin.z/abs(s.spin.z))*dot(norm(vector(s.spin.x,s.spin.y,0)),vector(-.707,-.707,0)) , 0 , .5+.5*(s.spin.z/abs(s.spin.z))*dot(norm(vector(s.spin.x,s.spin.y,0)),vector(.707,.707,0)))


                    if torque_toggle == 1:
                        s.torquevec.axis = .5*s.torque/(mag(s.torque)+.001)
                        s.torquevec.pos = s.pos + s.spinvec.axis/2.


                if Nsteps == 1000:
                    tt = clock()-tt
                    print '%0.1f' % tt, 'seconds for', Nsteps, 'steps with', Ntotal, 'spins'
                Nsteps = Nsteps+1
    
            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#

            pp.run = "stopped"

tkr = Tk()
pp = ProgClass(tkr)


scalerr = Scale(tkr,orient=VERTICAL,from_=100,to=0,label="R",command=lambda str: pp.set_red(str))
scalerr.pack(side=LEFT)
scalegg = Scale(tkr,orient=VERTICAL,from_=100,to=0,label="G",command=lambda str: pp.set_green(str))
scalegg.pack(side=LEFT)
scalebb = Scale(tkr,orient=VERTICAL,from_=100,to=0,label="B",command=lambda str: pp.set_blue(str))

scalebb.pack(side=LEFT)

tkr.mainloop() # Tk event loop

pp.run = "stop" # stop annimation
while pp.run == "stop":
    time.sleep(0.1) # give animation thread time to respond as "stopped"

scene.hide() # close visual
