Updated whole code with latest working version