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

from sfctools.examples.example_wrapper import Example
from sfctools import Agent, FlowMatrix, Accounts, World
from sfctools import AST as ASSETS
from sfctools import EQT as EQUITY


def run():

    FlowMatrix().reset()
    World().reset()

    # create some agent classes
    class MyAgent(Agent):
        def __init__(self):
            super().__init__()
            with self.balance_sheet.modify:

                self.balance_sheet.change_item(
                    "Cash", ASSETS, +100.0, suppress_stock=True)
                self.balance_sheet.change_item(
                    "Equity", EQUITY, +100.0, suppress_stock=True)

    class TypeA(MyAgent):
        def __init__(self):
            super().__init__()

    class TypeB(MyAgent):
        def __init__(self):
            super().__init__()

    class TypeC(MyAgent):
        def __init__(self):
            super().__init__()

    class TypeD(MyAgent):
        def __init__(self):
            super().__init__()

    class TypeE(MyAgent):
        def __init__(self):
            super().__init__()

    my_a = TypeA()  # create agents
    my_b = TypeB()
    my_c = TypeC()
    my_d = TypeD()
    my_e = TypeE()
    my_e2 = TypeE()

    def transfer(t, agent1, agent2, subject, quantity):  # define a generic transaction
        with agent1.balance_sheet.modify:
            with agent2.balance_sheet.modify:

                agent1.balance_sheet.change_item("Cash", ASSETS, -quantity)
                agent1.balance_sheet.change_item("Equity", EQUITY, -quantity)

                agent2.balance_sheet.change_item("Cash", ASSETS, +quantity)
                agent2.balance_sheet.change_item("Equity", EQUITY, +quantity)

                FlowMatrix().log_flow(t, quantity, agent1, agent2, subject=subject)

    # do some transactions
    transfer((Accounts.CA, Accounts.CA), my_a, my_b, "my subject", quantity=42)
    transfer((Accounts.CA, Accounts.CA), my_a,
             my_c, "other subject", quantity=3.1415)
    transfer((Accounts.CA, Accounts.KA), my_c, my_d,
             "different subject", quantity=10.0)
    transfer((Accounts.KA, Accounts.KA), my_a,
             my_e, "no subject", quantity=3.1415)
    transfer((Accounts.KA, Accounts.KA), my_a,
             my_e2, "no subject", quantity=3.1415)

    # output flow matrix as table and Sankey chart
    if __name__ == "__main__":
        print(FlowMatrix().to_dataframe())
        FlowMatrix().plot_colored(group=False)
        FlowMatrix().plot_colored(group=True)
        FlowMatrix().plot_sankey()


class FlowMatrixExample(Example):
    def __init__(self):
        super().__init__(lambda: run())


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