Add code for data collection