Bug fix
OOP seperation