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,at_sales, max_price_adj_rate) -> None:
        super().__init__(simulation, business, resource, exchanges)
        self.lr = lr
        self.lr_at_sales=at_sales
        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
        self.lp_threshold=0.05
        self.hp_threshold=0.80

    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],2)
            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):
        half=self.max_price_adj_rate/2
        ran=random.randrange(-half,half)
        self.open_orders[cx_id].append({
            'id': order.order_id,
            'lifetime': self.max_price_adj_rate+ran,
            '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
                    succsess=self.calc_order_success(cx,o)
                    modifier=0
                    if succsess>=self.hp_threshold:
                        modifier=1
                    elif succsess<=self.lp_threshold:
                        modifier=-1
                    self.update_believe(cx_id,o.price, succsess)  # 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()
                    succsess=self.calc_order_success(cx,o)
                    
                    cx.cancel_order(i["id"])
                    self.collect_balance_from_cxs()
                    self.collect_resource_from_cxs(self.resource)
                    self.update_believe(cx_id,o.price, succsess)
                    self.open_orders[cx_id].remove(i)
    def calc_order_success(self,cx, o):
        """
        Calculate how we should adjust the price belive
        """
        sold = o.qty-o.leaves_qty
        nsold=min(o.leaves_qty,cx.total_demand[self.resource])
        dem = cx.total_demand[self.resource]+sold
        if dem == 0:
            dem = 1
        score=sold-nsold
        return score

    def update_believe(self, cx_id, used_price, 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_at_sales)*self.lr

        self.price_believe[cx_id] = round(self.price_believe[cx_id], 2)
        if self.price_believe[cx_id] < 1:
            self.price_believe[cx_id] = 1

    def reset(self, episode):
        # Clean shop for today
        self.update_trades()
        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.open_orders = {i: [] for i in range(len(self.exchanges))}
        self.orders = {i: {} for i in range(len(self.exchanges))}
        return super().reset(episode)