cleaned up code . . . almost finished