Updated code + bugs