Updated Code after Testing