Cleaned code, added documentation