So I managed to get the functionality by hacking the annotations. Basically, just place a text annotation where you want the group names in the legend. Only thing I had to do was to group the different traces and then increase the tracegroupgap argument.