Modified the function
Fixed some bugs. Final program.