Fix bugs and add more styling