import evince as ev
import bubblebox as bb
import matplotlib.pyplot as plt
import numpy as np

from bubblebox.mdbox import lj_force, hook_force, no_forces, coulomb_force

ev.enable_notebook()
# simple harmonic oscillator (classical)

b = bb.mdbox(n_bubbles = 2, size = (3,3))
b.set_masses([1.0, 2.0])

b.pos[:,0] = np.array([-.6,0])
b.pos[:,1] = np.array([ .6,0])


print("Particle 0 position:", b.pos[:,0])
print("Particle 1 position:", b.pos[:,1])


# set forces between particles to harmonic oscillator
# V = constant * (separation - equilibrium_distance)**2
constant = 4.0
equilibrium_distance = 2
b.set_forces(hook_force, np.array([constant,equilibrium_distance]))



b.view()
while True:
    b.advance()
    b.update_view()
# simple coulomb attraction  (or repulsion)

b = bb.mdbox(n_bubbles = 2, size = (10000,10000))

b.set_masses([1.0,2.0])

b.pos[:,0] = np.array([0,1])
b.pos[:,1] = np.array([1,0])

b.r2_cut = 10000 # infinite attraction
print("Particle 0 position:", b.pos[:,0])
print("Particle 1 position:", b.pos[:,1])

# Coulomb force
# F = -constant*Q1*Q1/separation


#charges
Q1 = 1.0
Q2 = -1.0

Q1Q2 = Q1*Q2



# charges
constant = 10.0

b.set_forces(coulomb_force, np.array([constant, Q1Q2]))


print("Interactions:")
print(b.interactions)
print("Potential between particle 0 and 1:")

b.pos[:,0] = np.array([1,0])
b.pos[:,1] = np.array([0,1])
# set charges:





b.view()
while True:
    b.advance()
    b.update_view()
b = bb.mdbox(2, size = (-10,-10))
b.pos *= .1
b.set_masses(1000, np.array([0])) #adjust the mass of one particle

b.view() # set manually the number of iterations per frame 
while True:
    b.advance()
    b.update_view()
b = bb.showcase.fcc_system(lattice_parameter=1.5)

# select at random 100 unique bubbles 
selection = np.random.choice(b.n_bubbles, 100, replace = False)

# set the selected bubbles mass to 4
b.set_masses(4.0, selection)

# set the selected bubbles to inactive
b.active[selection] = False
b.view()
while True:
    b.advance()
    b.update_view()
b = bb.showcase.fcc_system(lattice_parameter=1.4)

# select a spherical shell
r = 2.0 #radius
selection = np.abs(np.sqrt(b.pos[0]**2 + b.pos[1]**2 + b.pos[2]**2)-r)<0.8

# a solid sphere
selection = np.sqrt(b.pos[0]**2 + b.pos[1]**2)<r




# set the selected bubbles mass to 4
b.masses[selection] = 4

# set the selected bubbles to inactive
b.active[selection] = False
b.view()
n = 12
b = bb.mdbox(n_bubbles = n**2, size = (7,7), vel = 0.1)

species_a = np.arange(int(.5*n**2))
species_b = np.arange(int(.5*n**2), n**2)

b.set_masses(2.0, species_a)



d2, d1 = np.sqrt(2), 1

# set forces between a and b
b.set_forces(lj_force, np.array([10.0, d1]), species_a, species_b)

# set forces within a and a
b.set_forces(lj_force, np.array([10.0, d2]),  species_a, species_b)

# set forces within b and b
b.set_forces(lj_force, np.array([10.0, d2]),  species_a, species_b)



b.view()
while True:
    for j in range(10):
        b.advance()
    b.update_view()
    b.vel_ *= .95