Refactorization - fixed many bugs