Re-organized the code in better way