Lots of code cleanup

- Cleaning up code
- Logic fixups

Signed-off-by: Satbir Jhuti <satbirjhuti@gmail.com>