Made some changes to mainScreen to make it generic