Added colorbar and legend to plots
Merged after some conflicts