Fixed include and added comments