Reorganize code and improve structure