More cleanup + some logic fixes.