Cleaned up jQuery code a bit