Added styling for axis
Refactored js a bit