Fixed bugs after testing