Code clean-up and solving some displaying bugs