Some logic fixes