added docstrs, cleaned code