WIP: added some logic for sorting