The current sankey library from JFreeChart does not allow the different flows from one node to be different colors. It also behaves strangely when you have more than 2 'stages' and its not easy to label the nodes / flows with text.
The data2viz kotlin library will hopefully improve on these things, hence I'm working on porting this to AnyLogic