Almost there, just need to clean stuff up