Fix layout and move updateLayout to the corresponding screen state classes