from visual import *
from visual.graph import *
from random import *
N=25

jitter=0.0
unit=1.e-3


dt_small=0.005 #3.e-4
dt_large=0.2 # 3.e-2
dt=dt_small


quake_moment_graph=gdisplay(x=0, y=400, width=400, height=350,
                        title="Quake Moment ",xtitle="t",ytitle="quake moment",
                        background=color.white, foreground=color.blue)
quake_moment_plot=gcurve(color=color.blue)

quake_magnitude_graph=gdisplay(x=400, y=400, width=400, height=350,
                        title="Quake Magnitude Distribution",xtitle="Magnitude",ytitle="Count",
                        background=(0.1,0.1,0.1), foreground=color.blue)
quake_magnitude_plot=gdots(color=color.blue)
BINMAX=20
quake_magnitude_bins=zeros( (BINMAX) )

scene=display(x=0,y=0,width=800, height=400, background=color.yellow, fov=1e-8, range=(5*unit,5*unit,5*unit) )
#scene.center=vector( (float(N)-1)/2. ,0,0)



###############  VISUALIZATIONS
###############
###############
###
block_pos=zeros( (N) , "float")
block_vel=zeros( (N) , "float")

plate=[]
plate_edge=[]
pendulum=[]
pendulum_vertex=[]
block=[]
spring=[]
graph_frame=frame()
plate_plot=[]
block_plot=[]

def initialize_kinematics():
    for i in arange(N):
        block_pos[i]=(i+(random()-0.5)*jitter)*unit
        block_vel[i]=0.


##
v0=0.

def update_kinematics(x,v, v0, t, dt):
    xnew=zeros( (len(x)), "float" )
    vnew=zeros( (len(x)), "float" )

    ####
    ##
    while(true): #will emulate a do_while



        #### YOUR CODE GOES HERE
        ####
        ##
        ##
            


        ##
        ##
        ####
        ####

        small_step=(sum(block_vel*block_vel)>0)
        if (small_step):
            if dt==dt_small:
                loop_again=false
            else:
                loop_again=true
                dt=dt_small
        else:
            if dt==dt_large:
                loop_again=false
            else:
                loop_again=true
                dt=dt_large

        if loop_again==false:
            break       ##do... while(loop_again)

    ##
    ####


    return dt



def initialize_geometry():
    plate_edge.append(
        sphere(pos=vector(0,1.e3,0)*unit, radius=0.1*unit,color=color.orange)
    )#plate[0] is the upper plate
    plate_edge.append(
        sphere(pos=vector(0,0,0)*unit, radius=0.1*unit,color=(.7,.7,.7)  )
    )#plate[1] is the lower plate

    plank_length=1000
    for i in range(2):
        plate.append(
            box(pos=plate_edge[i].pos+vector( plank_length ,0,0 )*unit,
                axis=(1,0,0), size=(2*plank_length*unit,.05*unit,1*unit) , color=plate_edge[i].color)
        )
    
    
    for i in range(N):
        block.append(
            box( vel=float(block_vel[i]), pos=(float(block_pos[i]),0.125*unit,0), size=(0.25*unit,0.25*unit,0.25*unit), color=color.cyan )
        )
        pendulum_vertex.append(block[i].x)
        pendulum.append(
            curve( pos=[vector(block[i].x, block[i].width, block[i].z),
                        vector(block[i].x,plate_edge[0].y, block[i].z)], color=color.green)
        )

    for i in range(N-1):
        spring.append(
            helix(  pos=block[i].pos+vector(block[i].length/2.,0,0),
                    axis=( block[i+1].pos-block[i].pos-vector(block[i].length,0,0) ),
                    radius=0.1*unit,thickness=0.02*unit , color=color.red)
        )  
def initialize_plots():
    for i in range(2):
        plate_plot.append(
            curve(pos=(plate_edge[i].x, block[0].y ,0) , color=plate_edge[i].color, frame=graph_frame )
        )
    for i in range(N):
        block_plot.append(
            curve(pos=block[i].pos, color=block[i].color, frame=graph_frame )
            )

def update_geometry(pos,plate_pos):
    plate_edge[0].x = plate_pos
    plate[0].x = plate_pos+(plate[0].length)/2.0
    for i in arange(len(pos)):
        block[i].x=pos[i]
        pendulum[i].x[0]=block[i].x
        pendulum[i].x[1]=pendulum_vertex[i]+plate_pos
    for i in range(N-1):
        spring[i].pos=block[i].pos+vector(block[i].length/2.,0,0)
        spring[i].axis=( block[i+1].pos-block[i].pos-vector(block[i].length,0,0) )
        
def update_plots(ticker):
    graph_frame.pos=vector(0,0,ticker)
    for i in range(2):
        plate_plot[i].append(vector(plate_edge[i].x, block[0].y ,0)+vector(0,0,-ticker) )
    for i in range(N):
        block_plot[i].append(block[i].pos+vector(0,0,-ticker) )

###
###############
###############
###############  VISUALIZATIONS



initialize_kinematics()  
initialize_geometry()


t=0

block_pos=array([ b.x for b in block ])
block_vel=array([ b.vel for b in block ])

update_geometry(block_pos,v0*t)
initialize_plots()


total_energy=0

ticker=0
gcurve_ticker=0

localt=0
quake_moment_old=0.0
while 1:
    #rate(10)

    dt=update_kinematics(block_pos,block_vel, v0, t, dt) 

    #for quake_magnitude
    
    quake_moment=0
    for b in block:
        quake_moment += abs(b.vel)*dt
    if (quake_moment>0):
        if (quake_moment_old<1.e-9):#quaking starts
            total_moment=0.
            localt=0.
    
        total_moment +=quake_moment

        if ticker%100==0:
            quake_moment_plot.plot( pos=(t,quake_moment) )
            gcurve_ticker +=1
            if gcurve_ticker%500==0:
                quake_moment_plot=None
                quake_moment_plot=gcurve(color=color.blue) #start a new gcurve

    if (quake_moment==0.) and (quake_moment_old != 0.):
        #quaking ends
        bin=10+((30*log(total_moment))/float(BINMAX))
        #print bin
        bin=min(max(0,int(bin)),BINMAX-1)
        #print bin
        if bin>0 and bin<BINMAX-1:
            quake_magnitude_bins[bin] +=1
            for p in quake_magnitude_plot.dots:
                p.visible=0
            del quake_magnitude_plot.dots[:]
            for i in arange(0,BINMAX):
                quake_magnitude_plot.plot( pos=(i,quake_magnitude_bins[i]) )
   
    quake_moment_old=quake_moment

    t +=dt
    ticker +=1.



        

    block_pos=array([ b.x for b in block ])
    block_vel=array([ b.vel for b in block ])


        
    update_geometry(block_pos, v0*t )
    #update_plots(ticker*unit/1.e3)
    scene.title="%10.6f" % t            
    scene.center=block[N/2].pos