moved code from main class into model class