fixing bugs and improving setup