Improved main.js

Improved main.js