Fix bugs and implement new layout