Refactor custom.js