Added some AI logic