Removed unused code
Added necessary headers