Reverted code back to working version