Added some functionality @DataGenTask, Changed some ugly design mistakes