Added basic prediction functionality