Cleaned up the ugly and fixed a bug