from __future__ import division
from visual.controls import *
from Tkinter import *
import tkFileDialog
import string

# Bruce Sherwood, begun Spring 2002, major revision Spring 2006.
# Based on "EM Field" by David Trowbridge and Bruce Sherwood,
# which won first prize in the 1991 Educational Software Contest
# of the journal Computers in Physics.

# Originally written in the cT programming language and distributed
# by Physics Academic Software (http://webassign.net/pas).
# See the cT archives at vpython.org.
# Feel free to use this code in any way you like.

# Omissions: The original program "EM Field" displayed field lines and equipotentials,
# and it contained some explanatory displays about how the integrals were computed.

print """
Drag charges or currents out of the storage boxes at top of screen.
Drag charges or currents into storage boxes to remove from active sources.
Drag charge or current to reposition.
Click or drag with left button away from charges or currents to show field.
Drag with left button to create
  a Gaussian surface around a long charged rod,
  or an Amperian contour around a line current,
  or in 3D (electric) a round-trip potential (line integral).
Positive charge is red, negative charge is blue.
Current out of screen is magenta, current into screen is green.
"""

def resetscene():
    scene.up = (0,1,0)
    scene.forward = (0,0,-1)

scene.x = scene.y = 0
wscene = 700
hscene = 700
scene.background = color.white
scene.foreground = color.white
scene.height = hscene
scene.width = wscene
scene.range = 10
scene.fov = 0.001 # look at charged rods or currents end-on

# Main modes
ELECTRIC = 1; MAGNETIC = 2
# Types of measurement (or SOURCE if dragging a charge)
SOURCE = 1; VECTOR = 2; UNITVECTOR = 3; LINE = 4; SURFACE = 5

# Initial settings:
dim = 2 # 2 for 2D, 3 for 3D
mode = ELECTRIC
measure = VECTOR

k = 1 # electric or magnetic constant
q = 1 # elementary charge
rq = 0.3 # radius of charge or wire
constrain = False # constrain sources and vectors to be placed on integer grid
Qpos = color.red
Qneg = color.blue
Ipos = color.magenta
Ineg = color.green
Ecolor = color.orange # electric field color
Mcolor = color.cyan # magnetic field color
EFcolor = (0.73, 0.62, 0) # flux color for E Gaussian surface
MFcolor = (0, 0.7, 0.7) # flux color for B Gaussian surface
Fcolor = EFcolor # initially
Eshaft = 0.3
Escale = 10
Gscale = 5 # for Gauss's law
Ascale = 1 # for Ampere's law
Gtolerance = 0.5 # shortest dl for Gaussian surface
dz = vector(0,0,0.1) # place curves in front of flux etc.
offset = vector(1.5,0,0) # for label on line/surface integrals
rcurve = 0.07
sources = [] # list of sources now on screen
fields = [] # list of field vectors now on screen
Gmeasures = [] # list of all the current Gaussian surfaces or Amperian loops

class ChargeBox(object):
    def __init__(self, center, sign):
        w = 1.5
        h = 1.5
        t = 0.2
        d = 0.2
        self.icon = None
        self.center = center
        self.region = (center.x-w/2,center.y-h/2,center.x+w/2,center.y+h/2)
        self.sign = sign
        self.frame = frame(pos=center)
        for x in (-w/2,w/2):
            box(frame=self.frame, pos=(x,0,0), size=(t,h+t,d), color=color.black)
        for y in (-h/2,h/2):
            box(frame=self.frame, pos=(0,y,0), size=(w,t,d), color=color.black)
        self.update()

    def update(self):
        if self.sign > 0:
            if mode == ELECTRIC:
                c = Qpos
            else:
                c = Ipos
        else:
            if mode == ELECTRIC:
                c = Qneg
            else:
                c = Ineg
        if self.icon:
            self.icon.visible = 0
        if dim == 3:
            self.icon = sphere(frame=self.frame, radius=rq, color=c)
        else:
            self.icon = cylinder(frame=self.frame, radius=rq, axis=(0,0,rq), color=c)

    def visible(self, vis):
        self.icon.visible = vis

    def inbox(self, point):
        return ( (self.region[0] <= point.x <= self.region[2]) and (self.region[1] <= point.y <= self.region[3]) )

class Gmeasure(object):
    def __init__(self, labelpos):
        self.elements = [] # (p1, p2, faces)
        # The reason for omitting pos in the creation of self.curve
        # is that with older versions of Visual that caused an error.
        self.curve = curve(color=FieldDrag.lcolor, radius=rcurve)
        self.curve.append(pos=labelpos)
        self.label = label(pos=labelpos+offset, box=0, opacity=1,
                           color=color.white, visible=0)
        self.measure = measure
        self.turning = 0
        self.flux = None

    def showsource(self):
        source = round(self.flux/(2*pi))
        if self.turning < 0:
            source = -source
        if mode == ELECTRIC:
            if self.measure == SURFACE:
                s = 'Q'
                if abs(abs(self.turning)-2*pi) > 0.001*pi:
                    self.label.text = 'Illegal surface'
                    self.label.visible = 1
                    return
            else:
                self.label.text = DVdisplay(self.flux)
                self.label.visible = 1
                return
        else:
            if self.measure == SURFACE:
                s = 'Qmonopole'
            else:
                s = 'I'
        if source > 0:
            info = '%s=+%d' % (s,source)
        elif source < 0:
            info = '%s=%d' % (s,source)
        else:
            info = '%s=0' % s
        self.label.text = info
        self.label.visible = 1

    def update(self):
        self.flux = 0
        for G in self.elements:
            p1 = G[0]
            p2 = G[1]
            points, dflux = Integralpoints(self.measure, p1, p2)
            self.flux += dflux
            G[2].pos = points
            if self.measure == LINE:
                Gcolor = color.red
                if dflux < 0:
                    Gcolor = color.blue
                G[2].color = Gcolor
        self.showsource()

    def Integral(self, p1, p2, final):
        newdir = norm(p2-p1)
        if self.flux is None:
            self.flux = 0
            initialdir = newdir
            self.curve.append(pos=p2+dz)
        else:
            olddir = self.curve.pos[-1]-self.curve.pos[-2]
            self.turning += turn(olddir,newdir)
            self.curve.append(pos=p2+dz)
            if final:
                # Calculate turning angle from last segment to first segment
                initdir = self.curve.pos[1]-self.curve.pos[0]
                self.turning += turn(newdir,initdir)
        points, dflux = Integralpoints(self.measure, p1,p2)
        self.flux += dflux
        if self.measure == SURFACE:
            Gcolor = FieldDrag.fcolor
        else:
            if mode == ELECTRIC: # display DV as we drag
                if dflux < 0:
                    Gcolor = color.red
                else:
                    Gcolor = color.blue
                self.label.pos = p2+offset
                self.label.text = DVdisplay(self.flux)
                self.label.visible = 1
            else:
                if dflux > 0:
                    Gcolor = color.red
                else:
                    Gcolor = color.blue
        self.elements.append((p1, p2,
            faces(pos=points, color=Gcolor, normal=[(0,0,1),(0,0,1),(0,0,1),(0,0,1),(0,0,1),(0,0,1)])))

    def visible(self, vis):
        self.curve.visible = vis
        self.label.visible = vis
        for el in self.elements:
            el[2].visible = vis


def CalcField(pos):
    field = vector(0,0,0)
    if mode == ELECTRIC:
        for Q in sources:
            r = pos-Q.pos
            if mag(r) > Q.radius:
                if dim == 3:
                    field += k*Q.source*norm(r)/mag(r)**2
                else:
                    field += k*Q.source*norm(r)/mag(r)
            else:
                return vector(0,0,0)
    else:
        for I in sources:
            r = pos-I.pos
            if mag(r) > I.radius:
                field += k*I.source*cross(vector(0,0,1),norm(r))/mag(r)
            else:
                return vector(0,0,0)
    return field

def DVdisplay(DV):
    if DV > 0:
        sign='-'
    else:
        sign='+'
    return 'DV = %s%0.1f' % (sign,-10*DV)

def FieldUpdate():
    for f in fields:
        if f[0] == VECTOR:
            f[1].axis = Escale*CalcField(f[1].pos)
        else:
            f[1].axis = norm(CalcField(f[1].pos))           
    for G in Gmeasures:
        G.update()

def Integralpoints(measuretype, p1, p2): # used for both SURFACE and LINE integrals
    # For LINE integrals, "flux" is actually dot(field,dl).
    field = CalcField(0.5*(p1+p2))
    if measuretype == SURFACE:
        G = Gscale*field
        if cross(G,p2-p1).z > 0:
            points = [p1, p1+G, p2+G, p1, p2+G, p2]
        else:
            points = [p1, p2+G, p1+G, p1, p2, p2+G]
        return (points, cross(field, p2-p1).z)
    else:
        Bdl = dot(field,p2-p1)
        if Bdl:
            A = Ascale*abs(Bdl/mag(p2-p1))
        else:
            A = 0
        rhat = norm(cross(p2-p1,vector(0,0,1)))
        points = [p1-A*rhat, p1+A*rhat, p2+A*rhat, p1-A*rhat, p2+A*rhat, p2-A*rhat]
        return (points, Bdl)

def GetMouseEvent(m):
    global drag, dragpos, initialpos, pick
    if m.events:
        m = scene.mouse.getevent()
        if m.click == 'left' and m.pick not in sources and m.pos.y < 7.4:
            pos = m.pos
            if constrain:
                pos = vector(round(m.pos.x),round(m.pos.y),round(m.pos.z))
            if measure == VECTOR or measure != UNITVECTOR:
                fields.append((VECTOR, arrow(pos=pos, axis=Escale*CalcField(pos),
                                    color=FieldDrag.color, shaftwidth=Eshaft)))
            elif measure == UNITVECTOR:
                fields.append((UNITVECTOR, arrow(pos=pos, axis=norm(CalcField(pos)),
                                    color=FieldDrag.color, shaftwidth=Eshaft)))
        elif m.drag == 'left':
            if m.pick in sources: # drag an existing source
                pick = m.pick
                drag = SOURCE
                dragpos = m.pickpos
                FieldDrag.axis = vector(0,0,0) # turn off FieldDrag while dragging source
            elif not m.pick:
                drag = measure
                pick = None
                dragpos = initialpos = m.pos
                if measure == LINE or measure == SURFACE:
                    Gmeasures.append(Gmeasure(dragpos))
            elif m.pick in [PosBox.icon, NegBox.icon]: # drag source out of storage box
                if m.pick is PosBox.icon:
                    PosBox.visible(0)
                    newq = q
                else:
                    NegBox.visible(0)
                    newq = -q
                createsource(m.pos, newq)
                pick = sources[-1]
                drag = SOURCE
                dragpos = m.pos
        elif m.drop: # finished dragging
            if drag == SOURCE:
                if PosBox.inbox(m.pos) or NegBox.inbox(m.pos): # dropped back into storage box
                    sources.remove(pick)
                    pick.visible = 0
                    del pick
                if constrain:
                    pick.pos = vector(round(pick.x),round(pick.y),round(pick.z))
                drag = None
                PosBox.visible(1)
                NegBox.visible(1)
                pick = None # dropped the source at a new location
                FieldUpdate() # make sure all fields correspond to final source position
                dragpos = None
            elif drag == VECTOR or drag == UNITVECTOR:
                # at end of showing E, leave an arrow at this location:
                pos = FieldDrag.pos
                if constrain:
                    pos = vector(round(FieldDrag.x),round(FieldDrag.y),round(FieldDrag.z))
                if drag == VECTOR:
                    fields.append((VECTOR, arrow(pos=pos, axis=Escale*CalcField(pos),
                                        color=FieldDrag.color, shaftwidth=Eshaft)))
                else:
                    fields.append((UNITVECTOR, arrow(pos=pos, axis=norm(CalcField(pos)),
                                        color=FieldDrag.color, shaftwidth=Eshaft)))                
                FieldDrag.axis = vector(0,0,0)
                drag = None
                dragpos = None
            elif drag == SURFACE or drag == LINE: # must complete the round trip if not DV measurement
                drag = None
                FieldDrag.axis = vector(0,0,0)
                if mode == ELECTRIC and measure == LINE: return # if measuring DV
                r = initialpos-dragpos
                nsteps = int((mag(r)/Gtolerance))+1
                dr = r/nsteps
                final = False
                for nn in range(nsteps):
                    if nn == nsteps-1:
                        final = True
                    Gmeasures[-1].Integral(dragpos,dragpos+dr,final)
                    dragpos = dragpos + dr # was += dr, which causes reference trouble
                Gmeasures[-1].showsource()
                dragpos = None

def Drag(m):
    global dragpos, drag
    if drag == SOURCE:
        pick.pos += m.pos-dragpos
        dragpos = pick.pos
        FieldUpdate()
    elif drag == VECTOR or drag == UNITVECTOR:
        FieldDrag.pos = m.pos
        if drag == VECTOR:
            FieldDrag.axis = Escale*CalcField(m.pos)
        else:
            FieldDrag.axis = norm(CalcField(m.pos))
        dragpos = m.pos
    elif drag == LINE or drag == SURFACE:
        if len(Gmeasures[-1].elements)>2 and (mag(m.pos-initialpos) <= Gtolerance):
            drag = None
            Gmeasures[-1].Integral(dragpos,initialpos,True)
            Gmeasures[-1].showsource()
            FieldDrag.axis = vector(0,0,0)
            dragpos = None
        elif mag(m.pos-dragpos) > Gtolerance:
            newpos = m.pos
            Gmeasures[-1].Integral(dragpos,newpos,False)
            dragpos = newpos
            FieldDrag.pos = newpos
            FieldDrag.axis = Escale*CalcField(newpos)

def turn(v1,v2):
    # angle to turn from vector v1 to v2
    # counterclockwise is positive
    c = dot(v1,v2)
    s = cross(v1,v2).z
    return atan2(s,c)

def createsource(newpos, newq):
    global sources
    if mode == ELECTRIC:
        if newq > 0:
            newcolor = Qpos
        else:
            newcolor = Qneg
    else:
        if newq > 0:
            newcolor = Ipos
        else:
            newcolor = Ineg
    if dim == 3:
        sources.append(sphere(pos=newpos, radius=rq, color=newcolor, source=newq))
    else:
        sources.append(cylinder(pos=newpos, axis=(0,0,rq), radius=rq, color=newcolor, source=newq))

class Grid(object):
    def __init__(self):
        self.points = []
        for x in range(-9,10):
            for y in range(-9,8):
                self.points.append(sphere(pos=(x,y), radius=0.05, color=color.black))

    def visible(self, vis):
        for s in self.points:
            s.visible = vis

def clearmeasures():
    global Gmeasures, fields
    for obj in fields:
        obj[1].visible = 0
    fields = []
    for obj in Gmeasures:
        obj.visible(0)
    Gmeasures = []

def clearall():
    global sources
    for obj in sources:
        obj.visible = 0
    sources = []
    clearmeasures()
    setmeasure(VECTOR)
    FieldDrag.axis = vector(0,0,0)
    FieldDrag.visible = 1

def setFielddragcolor():
    if mode == ELECTRIC:
        FieldDrag.color = Ecolor
        FieldDrag.fcolor = EFcolor
        FieldDrag.lcolor = color.cyan
    else:
        FieldDrag.color = Mcolor
        FieldDrag.fcolor = MFcolor
        FieldDrag.lcolor = color.yellow

def setdim():
    global dim
    scene.userspin = 0
    if bdim.value:
        if mode == MAGNETIC:
            bdim.value = 0
            return # can't toggle to 3D for magnetic fields
        dim = 3
        scene.userspin = 1
    else:
        dim = 2
        resetscene()
    PosBox.update()
    NegBox.update()
    setmode()
    clearall()

def setmode():
    global mode, dim
    if bmode.value:
        mode = MAGNETIC
        Title.text = 'Drag +z or -z current-carrying wire into scene.\nThen click or drag.'
        if dim == 3:
            dim = 2
            bdim.value = 0
            scene.userspin = 0
            resetscene()
    else:
        mode = ELECTRIC
        if dim == 3:
            Title.text = 'Drag positive or negative point charge into scene.\nThen click or drag. Rotate scene for 3D.'
            scene.userspin = 1
        else:
            Title.text = 'Drag positively or negatively charged long rod into scene.\nThen click or drag.'
            scene.userspin = 0
            resetscene()
    setupmeasuremenu()
    setFielddragcolor()
    PosBox.update()
    NegBox.update()
    clearall()

def setgrid():
    global grid
    if bgrid.value:
        if not grid:
            grid = Grid()
        else:
            grid.visible(1)
    else:
        grid.visible(0)

def setconstrain():
    global constrain
    constrain = bconstrain.value

def setscale(): # slider adjusts scaling of vectors
    global Escale, Gscale, Ascale
    newscale = bscale.value
    factor = newscale/Escale
    Escale *= factor
    Gscale *= factor
    Ascale *= factor
    FieldUpdate()

# em file format: title 'sources for fields version 1.0', then 'line charges' or 'point charges' or 'currents'
#   For each source, charge/current TAB x TAB y

def filedialog():
    global root
    if not root:
        # The following arcane sequence ensures that the file dialog box
        #   will appear in front of all other windows.
        root = Tk() # invoke Tk (graphics toolkit)
        root.wm_geometry(newGeometry='+0+0')
    else:
        root.deiconify()
    root.wait_visibility() # wait for root window to be displayed
    root.focus_force() # force it to be in focus (in front of all other windows)

def getsources():
    global root, sources, mode, dim
    filedialog()
    fd = tkFileDialog.askopenfile(filetypes=[('Sources', '.em')])
    root.withdraw() # make root window invisible
    if not fd:
        return
    data = fd.readlines()
    sourcetype = data[1][:-1]
    clearall()
    if sourcetype == 'currents':
        mode = MAGNETIC
        bmode.value = 1
        if dim == 3:
            dim = 2
            bdim.value = 0
    else:
        mode = ELECTRIC
        bmode.value = 0
        if sourcetype == 'point charges':
            dim = 3
            bdim.value = 1
        else:
            dim = 2
            bdim.value = 0
    setmode()
    PosBox.update()
    NegBox.update()
    for line in data[2:]:
        nums = string.split(line)
        newq = float(nums[0])
        x = float(nums[1])
        y = float(nums[2])
        z = 0.0
        if len(nums) == 4:
            z = float(nums[3])
        newpos = vector(x,y,z)
        createsource(newpos,newq)        

def savesources():
    global root
    filedialog()
    fname = tkFileDialog.asksaveasfilename(filetypes=[('Sources', '.em')])
    root.withdraw() # make root window invisible
    if not fname:
        return
    if fname[-3:] != '.em':
        fname = fname+'.em'
    fd = open(fname, 'w')
    fd.write('sources for fields version 1.0\n')
    if mode == ELECTRIC:
        if dim == 2:
            sourcetype = 'line charges'
        else:
            sourcetype = 'point charges'
    else:
        sourcetype = 'currents'
    fd.write(sourcetype+'\n')
    for s in sources:
        fd.write('%g\t%0.5f\t%0.5f\t%0.5f\n' % (s.source,s.x,s.y,s.z))

# mm file format: title 'measures for fields version 1.0', then
#   For each measure, code TAB x TAB y, where code tells what kind of measure
#   code = 1 (field), 2 (unit vector field), 3 (potential), 4 (equipotential)

def getmeasures():
    global root, sources, mode, dim
    filedialog()
    fd = tkFileDialog.askopenfile(filetypes=[('Measures', '.mm')])
    root.withdraw() # make root window invisible
    if not fd:
        return
    data = fd.readlines()
    for line in data[1:]:
        nums = string.split(line)
        measuretype = int(nums[0])
        x = float(nums[1])
        y = float(nums[2])
        z = 0.0
        if len(nums) == 4:
            z = float(nums[3])
        newpos = vector(x,y,z)
        if measuretype == VECTOR:
            fields.append((VECTOR, arrow(pos=newpos, axis=Escale*CalcField(newpos),
                                color=FieldDrag.color, shaftwidth=Eshaft)))
        elif measuretype == UNITVECTOR:
            fields.append((UNITVECTOR, arrow(pos=newpos, axis=norm(CalcField(newpos)),
                                color=FieldDrag.color, shaftwidth=Eshaft)))

def savemeasures():
    global root
    filedialog()
    fname = tkFileDialog.asksaveasfilename(filetypes=[('Measures', '.mm')])
    root.withdraw() # make root window invisible
    if not fname:
        return
    if fname[-3:] != '.mm':
        fname = fname+'.mm'
    fd = open(fname, 'w')
    fd.write('measures for fields version 1.0\n')
    for s in fields:
        fd.write('%d\t%0.5f\t%0.5f\t%0.5f\n' % (s[0],s[1].x,s[1].y,s[1].z))

def setupmeasuremenu():
    bmeasuremenu.items = []
    bmeasuremenu.items.append( ["Vector", lambda: setmeasure(VECTOR)] )
    bmeasuremenu.items.append( ["Unit vector", lambda: setmeasure(UNITVECTOR)] )
    if mode == ELECTRIC:
        bmeasuremenu.items.append( ["Potential", lambda: setmeasure(LINE)] )
    else:
        bmeasuremenu.items.append( ["Ampere's law", lambda: setmeasure(LINE)] )
    if dim == 2:
        bmeasuremenu.items.append( ["Gauss's law", lambda: setmeasure(SURFACE)] )
    setmeasure(measure)

def setmeasure(measuretype):
    global measure
    for item in bmeasuremenu.items:
        if item[0][0] == '*':
            item[0] = item[0][1:]
    if measuretype == VECTOR:
        bmeasuremenu.items[0][0] = '*'+bmeasuremenu.items[0][0]
    elif measuretype == UNITVECTOR:
        bmeasuremenu.items[1][0] = '*'+bmeasuremenu.items[1][0]
    elif measuretype == LINE:
        bmeasuremenu.items[2][0] = '*'+bmeasuremenu.items[2][0]
    elif measuretype == SURFACE and len(bmeasuremenu.items) == 4:
        bmeasuremenu.items[3][0] = '*'+bmeasuremenu.items[3][0]
    measure = measuretype

FieldDrag = arrow(axis=(0,0,0), shaftwidth=Eshaft) # for dragging field arrow
setFielddragcolor()
PosBox = ChargeBox(vector(-9,8.5), 1)
NegBox = ChargeBox(vector(9,8.5), -1)
Title = label(pos=(0,8.25), opacity=1)
grid = None
root = None

c = controls(x=wscene, y=0, width=1024-wscene, height=hscene)
c.display.background = color.white
c.display.foreground = color.black
bmode = toggle(pos=(-35,80), height=12, width=10, action=lambda: setmode(),
               text0='Electric', text1='Magnetic')
bdim = toggle(pos=(-10,80), height=12, width=10, action=lambda: setdim(),
               text0='2D', text1='3D')
bgrid = toggle(pos=(10,80), height=12, width=10, action=lambda: setgrid(),
               text1='Grid')
bconstrain = toggle(pos=(35,80), height=12, width=10, action=lambda: setconstrain(),
               text1='Constrain')
bclearmeasures = button(pos=(-15,40), height=30, width=30, action=lambda: clearmeasures(),
                text='Clear', color=color.red)
bclearall = button(pos=(27,40), height=30, width=30, action=lambda: clearall(),
                   text='Delete', color=color.red)

bscale = slider(pos=(-40,-60), axis=(0,120,0), min=1, max=50, width=5, action=lambda: setscale() )
label(display=c.display, pos=(-27,-65), text='Field scale factor', opacity=0, box=0)
bscale.value = Escale

bresetscene = button(pos=(10,-40), height=20, width=40, action=lambda: resetscene(),
                     text='Reset scene', color=color.red)

bfilemenu = menu(pos=(27,10), height=8, width=33, text='Read/write files')
bfilemenu.items.append( ('Save sources', lambda: savesources()) )
bfilemenu.items.append( ('Get sources', lambda: getsources()) )
bfilemenu.items.append( ('Save measures', lambda: savemeasures()) )
bfilemenu.items.append( ('Get measures', lambda: getmeasures()) )

bmeasuremenu = menu(pos=(-15,10), height=8, width=33, text='Measure type')
setupmeasuremenu()
label(display=c.display, pos=(0,-85), text='See detailed instructions in shell window', opacity=0, box=0)
setmode()
dragpos = None

while 1:
    rate(400)
    c.interact()
    m = scene.mouse
    GetMouseEvent(m)
    if dragpos and (m.pos != dragpos):
        Drag(m)