Refactored heatmap