Bug fixes

Excuse the shit code, I'm tired