added set_color function, optimazed code, tried to document it better.