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


from sfctools import Agent, BalanceSheet, BalanceEntry
import numpy as np
import matplotlib.pyplot as plt

from sfctools.examples.example_wrapper import Example


def run():

    class MyAgent(Agent):
        """
        An example agent
        """

        def __init__(self):
            super().__init__()

            self.change_equity(20.0, "Asset1")
            self.change_equity(10.0, "Asset2")
            self.change_debt(18.0)

        def change_equity(self, q, which):  # increase or decrease equity by q / -q
            with self.balance_sheet.modify:
                self.balance_sheet.change_item(which, BalanceEntry.ASSETS, q)
                self.balance_sheet.change_item(
                    "Equity", BalanceEntry.EQUITY, q)

        def change_debt(self, q):  # increase or decrease debt by q / -q
            with self.balance_sheet.modify:
                self.balance_sheet.change_item(
                    "Debt", BalanceEntry.LIABILITIES, q)
                self.balance_sheet.change_item(
                    "Asset2", BalanceEntry.ASSETS, q)

        def swap_assets(self, q):  # swap between assets
            try:
                with self.balance_sheet.modify:
                    self.balance_sheet.change_item(
                        "Asset2", BalanceEntry.ASSETS, -q)
                    self.balance_sheet.change_item(
                        "Asset1", BalanceEntry.ASSETS, q)
            except:
                pass

        def update(self):

            q1 = 0.5*(np.random.rand())
            q2 = 0.5*(np.random.rand())
            q3 = .4*np.random.rand()

            self.change_equity(q1, "Asset1")
            self.change_debt(q1)
            self.swap_assets(q3)

    my_agent = MyAgent()  # create an agent

    my_balances = []  # allocate empty list

    for i in range(10):  # iterate simulation
        my_agent.update()
        # append balance sheet data to list
        my_balances.append(my_agent.balance_sheet.raw_data)

    # plot balance sheet evolution
    if __name__ == "__main__":
        BalanceSheet.plot_list(
            my_balances, dt=5, xlabel="Time", title="My Agent", show_liabilities=True)


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


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