Cleaned up code and organized layout a bit better