Refactor cost and gradient functions