Fix bugs from testing. It works now!