Refactor stuff so makes more sense

@danielsuo