Modified logic to make code work.