
__author__ = "Thomas Baldauf"
__email__ = "thomas.baldauf@dlr.de"
__license__ = "MIT"
__birthdate__ = '25.03.2025'
__status__ = 'test'  # options are: dev, test, prod

from sfctools.examples.example_wrapper import Example
import pandas as pd


def run():
    # part 1 of the example

    from sfctools import FlowMatrix, Accounts, Agent

    CA = Accounts.CA
    KA = Accounts.KA

    class A(Agent):
        def __init__(self):
            super().__init__()

    class B(Agent):
        def __init__(self):
            super().__init__()

    agent1 = A()
    agent2 = B()

    FlowMatrix().log_flow((CA, CA), 4.0, agent1, agent2, subject="subject1")
    FlowMatrix().log_flow((CA, CA), 2.2, agent1, agent2, subject="subject2")
    FlowMatrix().log_flow((CA, CA), 3.0, agent2, agent1, subject="subject3")
    FlowMatrix().log_flow((CA, KA), 4.0, agent1, agent1, subject="subject4")

    FlowMatrix().plot_sankey(
        colors=["orange", "steelblue", "lightgray", "gray"], use_legend=False)

    # part 2 of the example
    from sfctools.misc.mpl_plotting import plot_sankey

    my_sankey_source = pd.DataFrame({
        'from': ['A', 'A', 'A', 'B'],
        'to': ['C', 'C', 'D', 'D'],
        'value': [-1.0, -.2, 2.0, 3.0],
        'color_id': [0, 1, 1, 2]
    })

    my_sankey_sink = pd.DataFrame({
        'from': ['C', 'C', 'B', 'A'],  # , 'C'],
        'to': ['E', 'F', 'D', 'G'],
        'value': [0.2, .8, -2.0, -1.1],  # 1.0],
        'color_id': [0, 3, 4, 5],  # 5]
    })

    my_sankey_source = pd.DataFrame({
        'from': ['A', "A"],
        'to': ['C', "D"],
        'value': [-1.0, 1.0],
        'color_id': [0, 1]
    })

    my_sankey_sink = pd.DataFrame({
        'from': ['C', 'C', "A", "C", "A"],  # , 'C'],
        'to': ['E', 'F', "F", "E", "E"],
        'value': [0.2, .8, 1.5, -1.0, -1.4],  # 1.0],
        'color_id': [0, 3, 4, 4, 2],  # 5]
    })

    fig = plot_sankey([my_sankey_source, my_sankey_sink],
                      use_legend=False, round_values=2, space_layout=False, flip_labels=False, show_values=True, label_rot=0)
    # plt.close()


class SankeyExample(Example):
    def __init__(self):
        super().__init__(run)


if __name__ == "__main__":
    my_instance = SankeyExample()
    my_instance.run()
