Update code and prepare everything for Assigment 2