Refactored the heat-map