from Tkinter import *
from visual import *
from random import random
import time,sys,thread


##  Simulation class definition
class Simulation:
    def __init__(self,parent):
        self.parent = parent

        self.setup_scene()

        self.dt = .0
        self.J = -1
         
        self.Ntotx = 16
        self.Ntoty = 16
        self.Ntot = self.Ntotx*self.Ntoty
        self.dx = .5
        self.dy = .5

        self.spins = []
        self.spins = self.create_lattice(self.spins, self.Ntotx,self.Ntoty,self.dx,self.dy)

        self.k = 2*pi*vector(1./(self.Ntotx*self.dx),1./(self.Ntoty*self.dy),0)
        self.spin_length = 1.
        self.spinA_sigma = .001

        self.spin_mag_ratio = (1-sin(abs(self.k.x*self.dx)))/cos(abs(self.k.x*self.dx))
#        self.spin_mag_ratio = cos(_k.x*self.dx)/((1+(sin(_k.x*self.dx))**2)**(1/2.))
#        self.spin_mag_ratio = 1/(((1+(sin(_k.x*self.dx))**2)**(1/2))/cos(_k.x*self.dx))

	self.spin_mag_ratio = self.spin_mag_ratio
	
	# 2 2 .758

        self.spins = self.initialize_spins(self.spins, self.k, self.spin_length, self.spinA_sigma, self.spin_mag_ratio)


        self.spinvec_length = 2
        self.spinvec_angleratio = .5/self.spinA_sigma
        self.spinvec_shaftwidth = .2
        self.spins = self.initialize_spinvecs(self.spins, self.spinvec_length, self.spinvec_angleratio, self.spinvec_shaftwidth) 

        self.toggle_torques = IntVar()
        self.toggle_colors = IntVar()
        self.toggle_norm_S = IntVar()

        thread.start_new_thread(self.animate_spins,(self.spins,))


    def create_lattice(self, _lattice=[], _Ntotx=8, _Ntoty=8, _deltax=1.0, _deltay=1.0):
        self.xmin = -_Ntotx*_deltax/2.
        self.ymin = -_Ntoty*_deltay/2.
        self.nz = 0

        for nx in range(_Ntotx):
              x = self.xmin + nx*_deltax
              for ny in range(_Ntoty):
                  y = self.ymin + ny*_deltay

                  _lattice.append(frame())
                  _lattice[-1].pos = vector(x,y,0)
                  _lattice[-1].spin = vector()
                  _lattice[-1].spinvec = arrow(pos=(x,y,0), axis=(0,0,1), color=(.9,.9,.9), shaftwidth=0.2)
              
                  _lattice[-1].torque = vector()
                  _lattice[-1].torquevec = arrow(pos=(x,y,0), axis=(0,0,1), color=(.3,.3,1), shaftwidth=0.05)
              
                  _lattice[-1].nearx = range(2)
                  _lattice[-1].neary = range(2)
                  _lattice[-1].indices = (nx,ny,self.nz)
              
        for s in _lattice:
              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 _lattice


    def initialize_spins(self, _lattice, _k=(1/8.,1/8.,1/8,), _spin_length=1.0, _spinA_sigma=.00001, _spin_mag_ratio=1.0):
        print 'wavenumber in x direction', _k.x

        print 'spin up to down magnitude ratio', _spin_mag_ratio

        ## For ratios over the critical value we see an ellipsoidal motion of the individual spins, as the wave juts forward in burts at regular intervals.  The problem is no longer a normal mode, and is experiencing nonlinear behavior.
        ## For ratios under the critical value we also see ellipsoids, but they appear in a different direction, along the chain for an up spin.
        ## APS Absolutely gives the correct results.

        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 _lattice:
            if ((self.J/abs(self.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 ((self.J/abs(self.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])
        return _lattice


    def initialize_spinvecs(self, _lattice, _spinvec_length=2.0, _spinvec_angleratio=50.0, _spinvec_shaftwidth=0.2):
        for s in _lattice:
            s.spinvec.axis = s.spin*_spinvec_length
            s.spinvec.pos = s.pos - s.spinvec.axis/2.
            s.spinvec.shaftwidth = _spinvec_shaftwidth
        return _lattice

   
    def animate_spins(self, _lattice):
##        self.tt = clock()
        self.Nsteps = 0

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

                for nnx in range(2):
                    nspinx = s.nearx[nnx]
                    s.torque = s.torque+cross(s.spin,_lattice[nspinx].spin)
                for nny in range(2):
                    nspiny = s.neary[nny]
                    s.torque = s.torque+cross(s.spin,_lattice[nspiny].spin)

##                s.torque.z = 0
                
            for s in _lattice:
                s.spin = s.spin + s.torque*self.dt                           # add torque

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

                angle_phi = atan(mag(vector(s.spin.x,s.spin.y,0))/s.spin.z)
                
                s.spinvec.axis = self.spinvec_length*rotate(s.spin, angle=(self.spinvec_angleratio-1)*atan(angle_phi), axis=cross(vector(0,0,1),s.spin))
                
                s.spinvec.pos = s.pos - s.spinvec.axis/2.
                
                if self.toggle_colors.get() == 1:
                    s.spinvec.color = (.5+.5*(s.spin.z/abs(s.spin.z))*dot(norm(vector(s.spin.x,s.spin.y,0)),norm(-self.k)) , 0 , .5+.5*(s.spin.z/abs(s.spin.z))*dot(norm(vector(s.spin.x,s.spin.y,0)),norm(self.k)))

                if self.toggle_torques.get() == 1:
                    s.torquevec.axis = .5*norm(s.torque)
                    s.torquevec.pos = s.pos + s.spinvec.axis/2.
                    
##
##            if self.Nsteps == 90:
##                self.tt = self.clock()-self.tt
##                print '%0.1f' % tt, 'seconds for', self.Nsteps, 'steps with', self.Ntotal, 'spins'
##            self.Nsteps = self.Nsteps+1


    def setup_scene(self):
        scene.title = "Simulation"
        scene.width = 700
        scene.height = 600
        scene.x = 50
        scene.y = 150
        scene.autoscale = 0
        scene.up = (0,0,1)
        scene.forward = (0,-.5,-4)
        scene.lights = [vector(-.5,.5,.5), vector(.5,-.5,.5)]


    def t_torques(self):
        for s in self.spins:
            s.torquevec.visible = self.toggle_torques.get()
            
    def t_colors(self):
        for s in self.spins:
            s.spinvec.color = (1,1,1)

    def t_norm_S(self):
        pass

    def set_angleratio(self,angleratio_string):
        self.spinvec_angleratio = string.atoi(angleratio_string)

    def set_gammaratio(self,gammaratio_string):
        self.spin_mag_ratio = string.atoi(gammaratio_string)/1000.

    def set_dt(self,dt_string):
        self.dt = string.atoi(dt_string)/1000.0

##    def set_Ntotx(self,Ntotx_string):
##        self.Ntotx = string.atoi(Ntotx_string)
###        print 'Hit reset to see changes'
##
##    def set_Ntoty(self,Ntoty_string):
##        self.Ntoty = string.atoi(Ntoty_string)
###        print 'Hit reset to see changes'

    def reset(self):
#        newspins = []
#        newspins = self.create_lattice(newspins, self.Ntotx,self.Ntoty,self.dx,self.dy)
        print self.spin_mag_ratio
        self.spins = self.initialize_spins(self.spins, self.k, self.spin_length, self.spinA_sigma, self.spin_mag_ratio)
        self.spins = self.initialize_spinvecs(self.spins, self.spinvec_length, self.spinvec_angleratio, self.spinvec_shaftwidth)




##  Instance creation
tkr = Tk()
simu = Simulation(tkr)

##  TKR appearance and widget creation
tkr.wm_geometry(newGeometry="400x300+750+150")   
tkr.wm_title("Controls")

##Ntotx_widget = Scale(tkr, orient=VERTICAL, from_=100, to=0, label="Nx", command=lambda str: simu.set_Ntotx(str))
##Ntotx_widget.set(simu.Ntotx)
##Ntotx_widget.pack(side=LEFT)
##
##Ntoty_widget = Scale(tkr, orient=VERTICAL, from_=100, to=0, label="Ny", command=lambda str: simu.set_Ntoty(str))
##Ntoty_widget.set(simu.Ntoty)
##Ntoty_widget.pack(side=LEFT)

dt_widget = Scale(tkr, orient=VERTICAL, from_=100, to=0, label="dt", command=lambda str: simu.set_dt(str))
dt_widget.set(simu.dt*1000.0)
dt_widget.pack(side=RIGHT)

angleratio_widget = Scale(tkr, orient=VERTICAL, from_=1/simu.spinA_sigma, to=1, label="TR", command=lambda str: simu.set_angleratio(str))
angleratio_widget.set(simu.spinvec_angleratio)
angleratio_widget.pack(side=RIGHT)

gammaratio_widget = Scale(tkr, orient=VERTICAL, from_=2000, to=0, label="SR", command=lambda str: simu.set_gammaratio(str))
gammaratio_widget.set(simu.spin_mag_ratio*1000)
gammaratio_widget.pack(side=RIGHT)

torque_widget = Checkbutton(tkr, text="Torques", variable=simu.toggle_torques, command=simu.t_torques)
simu.t_torques()
torque_widget.pack()

norm_S_widget = Checkbutton(tkr, text="Keep Spins Normalized", variable=simu.toggle_norm_S, command=simu.t_norm_S)
norm_S_widget.toggle()
norm_S_widget.pack()

color_widget = Checkbutton(tkr, text="Colors", variable=simu.toggle_colors, command=simu.t_colors)
color_widget.toggle()
color_widget.pack()

reset_widget = Button(tkr, text="Reset", command=simu.reset)
reset_widget.pack()



##  Enter the main TKR loop, starting the interface and program
tkr.mainloop()
