import fresnel
import gsd.hoomd
traj = gsd.hoomd.open('doped_configuration.gsd', 'rb') 
from PySide2 import QtCore
%gui qt
import fresnel.interact

Nchiral = 600
Nachiral = 5400

snap=traj[0]
box = snap.configuration.box

main=[]
dopant=[]
for i in range (0, Nachiral):
 main.append(i)
for i in range (Nachiral, Npart):
 dopant.append(i)
  
N = snap.particles.N
particle_body = snap.particles.body
particle_types = snap.particles.typeid

colors = np.empty((N, 3))

for i in range (0, Npart):
 if i in main:
  colors[particle_body == i] = fresnel.color.linear([0.586, 0.152, 0.758]) #purple
 elif i in dopant:
  colors[particle_body == i] = fresnel.color.linear([0.82, 0.828, 0.281]) #gold
 


colors[particle_types == 0] = fresnel.color.linear([1.0, 1.0, 1.0]) #centre of mass
colors[particle_types == 1] = fresnel.color.linear([1.0, 1.0, 1.0]) #centre of mass

scene = fresnel.Scene()

geometry = fresnel.geometry.Sphere(scene, N=len(snap.particles.position), radius=0.5)
geometry.material = fresnel.material.Material(roughness=0.5, primitive_color_mix=1.0)
geometry.color[:] = fresnel.color.linear(colors)
geometry.position[:] = snap.particles.position[:]
geometry.radius[snap.particles.typeid[:] == 0] = 0.5
geometry.outline_width = 0.04
fresnel.geometry.Box(scene, box, box_radius=0.3)


scene.camera = fresnel.camera.Orthographic.fit(scene)
scene.lights = fresnel.light.lightbox()
fresnel.pathtrace(scene, light_samples=5)

view = fresnel.interact.SceneView(scene)
view.show()