reformat + improve code