refactor visualization class