Updated design a bit