[CODE IMPROVEMENT] Rework plotting

[psinger]

The current plotting functionality could need some rework, some suggestions:

• For training plotting, I would just plot the whole sample, and (color)mark the labels.
• For validation insights we need better visual depictions of the different parts and clearer separation between samples.