Updated and reorganized code