Fixes, code reogranization
