from .base_distribution_agent import Base_Distribution_Agent
import random


class Price_Believe_Distribiute_Agent(Base_Distribution_Agent):
    """
    Aquire agent with internal price believe system. 
    """

    def __init__(self, simulation, business, resource, exchanges: list, lr, max_price_adj_rate) -> None:
        super().__init__(simulation, business, resource, exchanges)
        self.lr = lr
        self.max_price_adj_rate = max_price_adj_rate
        self.price_believe = {i: 1 for i in range(len(self.exchanges))}
        self.open_orders = {i: [] for i in range(len(self.exchanges))}
        self.open_qty = 0

    def tick(self, step, episode):

        order_error = self.target_error()
        if order_error > 0:
            # aquire based on current price belive
            cx_id = self.select_best_cx()

            price=round(self.price_believe[cx_id])
            if price<self.min_price:
                price=self.min_price
            order = self.distribute_resource(
                price, order_error, cx_id)
            self.register_order(cx_id, order)
        self.tick_open_orders()

    def select_best_cx(self):
        best_id = -1
        best = -1
        for cx_id in range(len(self.exchanges)):
            cx = self.exchanges[cx_id]
            available = (cx.get_total_demand(self.resource) > 0)
            if available:
                potential = 1*self.price_believe[cx_id]
            else:
                continue
            if potential > best:
                best = potential
                best_id = cx_id
        if best_id == -1:
            best_id = random.randint(0, len(self.exchanges)-1)
        return best_id

    def register_order(self, cx_id, order):
        self.open_orders[cx_id].append({
            'id': order.order_id,
            'lifetime': self.max_price_adj_rate,
            'leaves': order.leaves_qty
        })
        if order.leaves_qty!=order.qty:
            self.update_trades()

    def tick_open_orders(self) -> int:
        """
        Removes 1 from pending orders timeout timer and cancels timeed out orders.
        Returns ids of orders that timeed out. 
        """
        self.open_qty = 0
        for cx_id in range(len(self.exchanges)):
            cx = self.exchanges[cx_id]
            cx_orders = self.open_orders[cx_id]
            for i in cx_orders:
                # Check for each order if it is fullfiled or if it is timed
                o = self.orders[cx_id][i["id"]]
                leaves=o.leaves_qty
                if o.leaves_qty == 0:
                    # order is done
                    self.open_orders[cx_id].remove(i)  # remove order from open
                    self.update_believe(cx_id, 1)  # update price believe
                    self.collect_balance_from_cxs()
                    self.collect_resource_from_cxs(self.resource)
                    self.update_trades()
                    continue
                
                if not (i["leaves"]==leaves):
                    #update in order
                    i["leaves"]=leaves
                    self.update_trades()
                    #reset lifetime
                    i["lifetime"]=self.max_price_adj_rate

                if i["lifetime"] > 0:
                    self.open_qty += o.leaves_qty
                    i["lifetime"] -= 1  # subtract lifetime
                else:
                    # timeout
                    self.update_trades()
                    cx.cancel_order(i["id"])
                    self.collect_balance_from_cxs()
                    self.collect_resource_from_cxs(self.resource)
                    self.update_believe(cx_id, -1)
                    self.open_orders[cx_id].remove(i)

    def update_believe(self, cx_id, modifier):
        """
        Updates the believe based on the modifier.
        If positive will add lr to believe
        If negative will sub lr to believe
        """
        self.price_believe[cx_id] += modifier*self.lr
        if self.price_believe[cx_id] < 1:
            self.price_believe[cx_id] = 1

    def reset(self, episode):
        # Clean shop for today
        for cx_id in range(len(self.exchanges)):
            cx = self.exchanges[cx_id]
            cx_orders = self.open_orders[cx_id]
            for i in cx_orders:
                cx.cancel_order(i["id"])
            self.collect_balance_from_cxs()
            self.collect_resource_from_cxs(self.resource)
        # book keeping
        self.update_trades()
        self.open_orders = {i: [] for i in range(len(self.exchanges))}
        self.orders = {i: {} for i in range(len(self.exchanges))}
        return super().reset(episode)