Total code rewrite, bug fixes