Code clean ups