from floweaver import *
import pandas as pd

# Set the default size to fit the documentation better.
size = dict(width=570, height=300)

d = [
     {'source':'Brain', 'target': 'Brain', 'type':'1', 'value':1},
     {'source':'Heart', 'target': 'Heart', 'type':'2', 'value':1},
     {'source':'Banana', 'target': 'Banana', 'type':'3', 'value':1},
     {'source':'Lung', 'target': 'Lung', 'type':'4', 'value':1},
     {'source':'Skin', 'target': 'Skin', 'type':'5', 'value':1},
    ]

flows = pd.DataFrame(data=d)

nodes = {
    't1': ProcessGroup(['Brain', 'Heart', 'Lung', '']),
    #'t2': ProcessGroup(['Heart', 'Brain', 'Banana']),
    #'t3': ProcessGroup(['Banana', 'Heart', 'Brain']),
    't4': ProcessGroup(['Lung', 'Skin', 'Brain', 'Heart', 'Banana']),
}

nodes['t1'].partition = Partition.Simple('process', ['Brain', 'Heart', 'Lung'])
nodes['t2'] = Waypoint(Partition.Simple('process', ['Heart', 'Lung', 'Brain']))
# Can I group 2 process groups for T3?
nodes['t3'] = Waypoint(Partition.Simple('process', ['Banana', 'Skin' 'Heart', 'Lung' 'Brain']))
nodes['t4'].partition = Partition.Simple('process', ['Lung', 'Brain', 'Heart', 'Banana'])

nodes['t2_1'] = ProcessGroup(['Skin'], partition=Partition.Simple('process', ['Skin']))
nodes['t2_2'] = ProcessGroup(['Banana'], partition=Partition.Simple('process', ['Banana']))

ordering = [
    [['t1']],
    [['t2_1'], ['t2'], ['t2_2']],
    [['t3']],
    [['t4']],
]

bundles = [
    Bundle('t1', 't4', waypoints=['t2', 't3']),
    Bundle('t2_1', Elsewhere, waypoints=['t3']),
    Bundle('t2_2', 't4', waypoints=['t3'])
]

df = pd.DataFrame(d)

ranks = Partition.Simple("type", df.type.unique())

sdd = SankeyDefinition(nodes, bundles, ordering, flow_partition=ranks)
weave(sdd, df).to_widget(**size)
weave(sdd, df).to_widget(**size).auto_save_png('./flow-3-example.png')