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

from .world import World
from .clock import Clock

from ..datastructs.balance import BalanceSheet, BalanceEntry
from ..datastructs.income_statement import IncomeStatement
from ..datastructs.cash_flow_statement import CashFlowStatement
from ..datastructs.signalslot import Signal, Slot

from ..sfc_structs import SfcArray, Sfc2DArray, Sfc3DArray
from .flow_matrix import FlowMatrix

import numpy as np
import warnings
from collections import defaultdict
from functools import wraps

from ..datastructs.autopickle import AutoPickle
import pickle


class AgentVisuals:
    """
    Wrapper class for visualization properties of the agent 
    """

    def __init__(self, parent):

        # information about shape of the agent
        self._target_pos = None  # targeted new position
        self._pos = 0 + 0j      # current position
        self._z = 0  # 3d position (optional)

        self.visual_pos = (0, 0)  # visual position on drawing canvas

        self.shape = "rect"
        self.fill_color = (100, 150, 200)  # (200, 200, 200)
        self.edge_color = (100, 150, 200)  # (200, 200, 200)
        self.edge_width = 0
        self.size = (10, 10)
        self.speed = 1.0
        self.hover = False

        self.block_rotation = False  # block rotation transform
        self.block_scaling = False  # block scaling transform
        self.block_shifting = False  # block shifting

        self.parent = parent
        World().visuals.object_buffer.append(self)

    @property
    def info(self):
        return self.parent.info

    def __repr__(self):
        return "<visual %s hover=%s>" % (self._target_pos, self.hover)

    def x(self):
        """
        x position of the agent 
        """
        return self.pos.real

    def y(self):
        """
        y positfion of the agent
        """
        return self.pos.imag

    def vx(self):
        # if self.block_rotation:
        #    return self.x()
        return self.visual_pos[0]

    def vy(self):
        # if self.block_rotation:
        #    return self.y()
        return self.visual_pos[1]

    def z(self):
        return self._z

    @property
    def pos(self):
        return self._pos

    def move(self, pos, z=0, speed=1.0):
        self.speed = speed
        if speed >= 1.0:
            self._pos = pos
        self._target_pos = pos
        self._z = z

    @pos.setter
    def pos(self, value):
        self._target_pos = value

    def update(self):
        if self._target_pos is not None:
            self._pos = (self._pos) + self.speed*(self._target_pos-self._pos)
        # pass


class Agent():
    """
    This is the base class for any agent. It takes care of the bookkeeping system
    and other low-level operations of the agent.

    """

    # (optional) place balance sheet items as list of str to determine balance structure of this agent class
    bal_items = None
    # maps str to index in the balance item (relevant for str indexing, if desired by the user
    bal_items_map = {}

    def __getstate__(self):
        state = {
            "SFC_ID": self._SFC_ID,
            "name": self.name,
            "balance_sheet": self.balance_sheet.to_dataframe(),
            "income_statement":  self.income_statement.to_dataframe(),
            "cash_flow_statement": self.cash_flow_statement.to_dataframe(),
            "bankrupt": self.bankrupt
        }
        # pickle all other
        for key, value in list(self.__dict__.items()):
            if key in state:  # Skip already included attributes
                continue
            if key in ["_balance_sheet"]:
                continue
            try:
                pickle.dumps(value)
                state[key] = value
            except Exception as e:
                warnings.warn(
                    f"[{self.__class__.__name__}]: Can't pickle '{key}', skipping.")

        return state

    def __setstate__(self, state):
        self._SFC_ID = state["SFC_ID"]
        self.name = state["name"]
        self.bankrupt = state["bankrupt"]
        for key, value in list(state.items()):
            if key in ["SFC_ID", "name", "bankrupt"]:
                continue
            self.__dict__[key] = value

    def __init__(self, bal_items=None, alias=None, verbose=False, signals: list = None, slots: list = None, visualize=False):
        """
        Instantiation of agents is done here.
        Note: Each agent will be given a name and an alias name. The name will be set automatically and
        the alias name will be a (possibly non-unique) name, for example 'name' could be 'Firm_0001' and
        alias_name could be 'MyEnergyProducer'.

        :param alias: (optional) alias or None(default)
        :param verbose: boolean switch to enable console output. Default is False
        :param signals: list of str or None(default), signal names connected to this node
        :param slots: list of str or None(default), slot names connected to this node.
        :param visualize (boolean): if set True, a visualization instance is created for this agent. See World.visuals
        """

        self.name = "New Agent"  # set dummy name for registering agent

        world = World()  # get singleton world
        # register agent at world registry
        self._SFC_ID = world.register_agent(self)

        # set name
        self.name = str(self.__class__.__name__) + \
            "__%05i" % world.agent_registry.get_count(self.__class__.__name__)

        if alias is None:  # set the alias if desired
            self.alias = self.name
        else:
            self.alias = alias

        if bal_items is None:
            if self.__class__.bal_items is not None:
                bal_items = self.__class__.bal_items
            else:
                bal_items = []

        # agent gets a blank balance sheet
        self._balance_sheet = BalanceSheet(self, bal_items)
        self.income_statement = IncomeStatement(self)  # blank income statement
        self.cash_flow_statement = CashFlowStatement(
            self)  # blank cash flow statement

        # triggers for tree struct to enable event-based approach
        self.trigger_dict = {}  # this is meant for automation

        # bankruptcy flag
        self.bankrupt = False

        # print notification in verbose mode
        if verbose:
            if self.name != self.alias:
                print("New Agent %s" % self.name, "alias", self.alias)
            else:
                print("New Agent %s" % self.name)

        # triggers for signals and slots
        if signals is not None or slots is not None:
            self.signals = {}  # agent only has attribute "signals" if given
            self.slots = {}  # agent only has attribute "slots" if given

            # fill signals and slots dictionaries...
            assert isinstance(
                signals, list) or signals is None, "signal must be a list of str or None"
            if signals is not None:
                for signal_name in signals:
                    self.signals[signal_name] = Signal.retrieve(signal_name)

            assert isinstance(
                slots, list) or slots is None, "signal must be a list of str or None"
            if slots is not None:
                for slot_name in slots:
                    self.slots[slot_name] = Slot.retrieve(slot_name)

        if visualize:
            self.visuals = AgentVisuals(parent=self)
        else:
            self.visuals = None

        self._info = None

    @property
    def info(self):
        return self._info

    @info.setter
    def info(self, value):
        self._info = value

    def trigger(self, event_key, *args, **kwargs):
        """
        trigger a certain default method of this agent
        :param event_key str: a key for the event
        """
        method = self.trigger_dict[event_key]
        method()

    @property
    def sfc_id(self):
        return self._SFC_ID

    @property
    def balance_sheet(self):
        """
        This will return the (private) BalanceSheet object of the agent.
        """
        return self._balance_sheet

    @property
    def bal(self):
        # short version of balance sheet
        return self.balance_sheet

    @property
    def ics(self):
        # short version of income statement
        return self.income_statement

    @property
    def cfs(self):
        # short version of cash flow statement
        return self.cash_flow_statement

    def filter_balance_key(self, key):
        # assert isinstance(key, tuple), "Wrong indexation."
        if key[1] == 0:
            key = (key[0], BalanceEntry.ASSETS)
        elif key[1] == 1:
            key = (key[0], BalanceEntry.LIABILITIES)
        elif key[1] == 2:
            key = (key[0], BalanceEntry.EQUITY)

        return key

    def fetch_src_arr_balance(self, key):
        cls_name = self.__class__.__name__
        q = 0
        key1 = key[1]

        data1d = SfcArray.sorted_instances[cls_name][key1]
        for arr in data1d:
            q += arr[self._SFC_ID]

        data2d = Sfc2DArray.sorted_instances[cls_name][key1]
        for arr in data2d:
            q += float(arr.fetch(cls_name)[self._SFC_ID])

        data3d = Sfc3DArray.sorted_instances[cls_name][key1]
        for arr in data3d:
            # get the entry from the current time step
            q += float(arr.fetch(cls_name)[self._SFC_ID, Clock().get_time()])
        return q

    def __getitem__(self, key):
        """
        Quick access function for the agent's balance sheet
        """
        key = self.filter_balance_key(key)
        if not World.use_sfc_array:
            return self.balance_sheet.get_balance(key[0], key[1])
        # NOTE local balance sheet is ignored if there is an external SfcArray registered
        return self.fetch_src_arr_balance(key)

    def __setitem__(self, key, value):
        """
        This magic method supports item assignment for agents, which is linked to balance sheet operations
        """

        key = self.filter_balance_key(key)
        if not World.use_sfc_array:
            self.balance_sheet.change_item(
                key[0], key[1], value, relative=False)

        cls_name = self.__class__.__name__
        key1 = key[1]

        data1d = SfcArray.sorted_instances[cls_name][key1]
        for arr in data1d:
            arr[self._SFC_ID] = value

        data2d = Sfc2DArray.sorted_instances[cls_name][key1]
        if len(data2d):
            raise RuntimeError(
                "Value assignment not supported because Sfc2DArray is being used.")

        data3d = Sfc3DArray.sorted_instances[cls_name][key1]
        if len(data3d):
            raise RuntimeError(
                "Value assignment not supported because Sfc3DArray is being used.")

    def __repr__(self):
        s = ""
        if hasattr(self, "alias") and self.name != self.alias:
            s = " alias " + self.alias

        return "<Agent: " + self.name + s + ">"

    def __str__(self):
        if hasattr(self, "alias") and self.alias is not None and self.name != self.alias:
            return self.name + " (" + self.alias + ")"  # + " [ %i ]"%id(self)
        else:
            return self.name  # + " [ %i ]"%id(self)

    def __lt__(self, other):  # required for pandas
        return self.name < other.name

    @property
    def leverage(self) -> float:
        """
        leverage of own balance sheet
        """
        return self.balance_sheet.leverage

    @property
    def net_worth(self) -> float:
        """
        net worth at own balance sheet
        """
        return self.balance_sheet.net_worth

    @property
    def total_assets(self):
        """
        sum of assets column of balance sheet
        """
        return self.balance_sheet.total_assets

    @property
    def total_liabilities(self):
        """
        sum of liabilities column of balance sheet
        """
        return self.balance_sheet.total_liabilities

    @property
    def cash_balance(self) -> float:
        """
        'Shortcut' property for balance of 'Cash' stored in 'Assets' of the agent
        """
        return self.balance_sheet.get_balance("Cash", 0)

    def file_bankruptcy(self, event=None):
        """ file agent for bankrupcy

        :param event: preferrably str, can be None for manual trigger but should be something like 'negative cash' or 'negative equity' (default).

        NOTE This should be called by a sub-module rather than manually by the user
        """
        self.bankrupt = True
        raise RuntimeError("%s went bankrupt (reason: %s)!\n\n%s" % (
            self, event, self.balance_sheet.to_string()))

    def get_balance(self, key, kind) -> float:
        """
        Gets the nominal balance of a certain entry in the balance sheet.
        :param key: str, which item is requested (e.g. 'Cash', 'Apples',...)
        :param kind: BalanceEntry
        """
        # get balance sheet asset balance
        return self.balance_sheet.get_balance(key, kind=kind)


class SingletonAgent(Agent):
    """
    SingletonAgent is an agent which is singleton, meaning that only one instance creation of this Agent will be allowed.
    """
    _instance = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super().__new__(cls)

        return cls._instance

    def __init__(self, alias=None, verbose=False):
        super().__init__(alias=alias, verbose=verbose)


def block_on_bankrupt(method):
    """
    This is a decorator you can use when building an agent.
    It will block this action whenever the agent is bankrupt.

    Usage:

    .. code-block:: python

        from sfctools import block_on_bankrupt

        class MyAgent(Agent):
            ...

        @block_on_bankrupt
        def my_fun(self,...):
            ...

    """

    @wraps(method)
    def _impl(self, *args, **kwargs):
        if not self.bankrupt:
            method_output = method(self, *args, **kwargs)
            return method_output
        else:
            warnings.warn("%s: tried to call forbidden method in bankruptcy: %s" % (
                self, method))  # <- TODO test this
            return None

    return _impl
