Major re-work and bug fixes