from yade import plot
import numpy as np



def exportStuffFunction(Omega,mySphere):
    plot.addData(t=Omega.time,Ek=utils.kineticEnergy(),posz=mySphere.state.pos[2])



def main():
    global mySphere1
    plot.reset()
    O.materials.append(FrictMat(young=10e9, poisson=0.11, density=2700, frictionAngle=90, label='Gomma'))
    mySphere1 = utils.sphere((0,0,0),radius=0.2,material='Gomma')
    O.bodies.append(mySphere1)
    O.bodies.append(utils.wall(-1,axis=2,material='Gomma'))

    O.engines=[
    ForceResetter(),
    InsertionSortCollider([Bo1_Sphere_Aabb(),Bo1_Wall_Aabb()]),
    InteractionLoop(
        [Ig2_Wall_Sphere_ScGeom()],
        [Ip2_FrictMat_FrictMat_FrictPhys()],
        [Law2_ScGeom_FrictPhys_CundallStrack()]),
        NewtonIntegrator(damping=0.0,gravity=(0,0,-9.81)),
        PyRunner(command='exportStuffFunction(O,mySphere1)',iterPeriod=10)

    ]


    O.dt = 0.25*utils.PWaveTimeStep()


    plot.plots={'t':('posz',None,'Ek')}
    O.trackEnergy=True

    O.run(100000)

    plot.plot()
    plot.saveDataTxt('Plot_data.txt',vars=('t'))


#----------------------


main()



myDict = plot.data
t_array = np.array(myDict["t"])
t_list = myDict["t"]


np.savetxt('np_array.txt', t_array,newline='\n')
np.savetxt('np_list.txt', t_list,newline='\n')


fid = open('write_data.txt','w')
for i in range (0,len(t_list)):
    fid.write(str(t_list[i])+'\n')
fid.close()

print('------------------------------------------------')
print(str(len(t_list))+" "+str(len(t_array))+" "+str(len(myDict["t"])))
print('------------------------------------------------')