import numpy as np
import matplotlib.pyplot as plt

from matplotlib.sankey import Sankey


plt.rcParams["figure.figsize"] = (15,10)


system_1 = [
    {"label": "1st",  "value":  1.00, "orientation":  -1},
    {"label": "2nd",  "value":  0.15, "orientation": -1},
    {"label": "3rd",  "value":  0.60, "orientation": -1},
    {"label": "5th",  "value":  0.25, "orientation": -1},
    {"label": "6th",  "value":  0.25, "orientation": -1},
    {"label": "7th",  "value":  0.25, "orientation": 1},
    {"label": "8th",  "value":  0.25, "orientation": 1},
    {"label": "9th",  "value":  0.25, "orientation": 1},
    {"label": "4th",  "value":  -0.10, "orientation": -1},
]

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, xticks=[], yticks=[], title="Where are all my cows?")
sankey = Sankey(ax=ax, unit="cow")
system_2 = system_1[:4]
#     system_2.append({"label": "new",  "value":  0.25, "orientation": 1})
system_2.append({"label": "4th",  "value": -0.10, "orientation": -1})


flows  = [x["value"] for x in system_1]
labels = [x["label"] for x in system_1]
orientations=[x["orientation"] for x in system_1]
sankey.add(flows=flows,
           labels=labels,
           label='one',
           pathlengths=[1.5,1.5,0.8,.6,.5,.5,.5,.5,1.5],
           orientations=orientations,
          )

sankey.add(flows=[-x["value"] for x in system_2],
           labels=[x["label"] for x in system_2],
           label='two',
           trunklength=1.5,
           orientations=[x["orientation"] for x in system_2],
           pathlengths=[1.5,1.5,0.8,0.2,1.5],
           prior=0,
         connect=(8,4)

      )

diagrams = sankey.finish()
diagrams[-1].patch.set_hatch('/')
plt.legend(loc='best')


plt.show()