Cleaned up code added a few lines.