from .base_aquire_agent import Base_Aquire_Agent
import random
class Price_Believe_Aquire_Agent(Base_Aquire_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: 0 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,tick,episode):

        
        order_error=self.open_qty+self.target_error()
        if order_error<0:
            # aquire based on current price belive
            cx_id=self.select_best_cx()
            order=self.order_resource(self.price_believe[cx_id],order_error*-1,cx_id)
            if not order==None:
                self.register_order(cx_id,order)
            else:
                # order failed due to missing balance.. we need to adjust our price believe
                self.collect_balance_from_cxs()
                self.collect_resource_from_cxs(self.resource)
                self.update_believe(cx_id,-1)
        self.tick_open_orders()    

    def select_best_cx(self):
        best_id=-1
        best=0
        for cx_id in range(len(self.exchanges)):
            cx=self.exchanges[cx_id]
            potential=cx.get_total_supply(self.resource)*self.price_believe[cx_id]
            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,
        })

    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"]]
                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)
                    continue
                if i["lifetime"]>0:
                    self.open_qty+=o.leaves_qty
                    i["lifetime"]-=1 # subtract lifetime
                else:
                    # timeout
                    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]<0:
            self.price_believe[cx_id]=0

    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()

        return super().reset(episode)