updated code for visualization