_bal="balance"
from lightmatchingengine.lightmatchingengine import LightMatchingEngine,Side,Trade,Order,OrderBook
import log
from influxdb_client.client.write_api import SYNCHRONOUS
from influxdb_client import InfluxDBClient, Point
import time
import pandas as pd
import uuid

software_start=int(time.time())
class Exchange():
    """
    Basic Commodity exchange.
    """
    
    def __init__(self) -> None:
        
        self.reset()
        pass

    def reset(self):
        self.account={}
        self.escrow={}
        self.lme=LightMatchingEngine()
        self.order_account_map={}
        self.orders={}
        self.executed_trades=[]
        self.order_trades_map={}
        self.market_rate={}
        self.best_ask={}
        self.best_bid={}
        self.total_demand={}
        self.demand={}
        self.total_supply={}
        self.supply={}
        self.traded_commoditys={}
    
    def add_to_account(self,account_id,resource,amount):
        """
        Adds resources to account escrow
        """
        #check if account exists
        if account_id not in self.account:
            self.account[account_id]={_bal: 0}
        #check if ressource exists
        if resource not in self.account[account_id]:
            self.account[account_id][resource]=0
        
        self.account[account_id][resource]+=amount
    
    def remove_from_account(self,account_id,resource,amount) -> int:
        """
        Remove resources from account. Returns amount retrieved
        """
        #check if account exists
        if account_id not in self.account:
            self.account[account_id]={_bal: 0}
        #check if ressource exists
        if resource not in self.account[account_id]:
            self.account[account_id][resource]=0
        ret=amount
        remain=self.account[account_id][resource]-amount
        if remain<0:
            ret=self.account[account_id][resource]
            self.account[account_id][resource]=0
        else:
            self.account[account_id][resource]-=amount
        return ret

    def _move_to_escrow(self,account_id,resource,amount):
        """
        Adds resources from account to escrow. Return true if done.
        """
        if not self.check_account_resource(account_id,resource,amount):
            return False
        
        #check if account exists
        if account_id not in self.escrow:
            self.escrow[account_id]={_bal: 0}
        #check if ressource exists
        if resource not in self.escrow[account_id]:
            self.escrow[account_id][resource]=0
        
        self.escrow[account_id][resource]+=amount
        self.account[account_id][resource]-=amount
        return True
    
    def _move_to_account(self,account_id,resource,amount):
        """
        Move resources from escrow to account.
        """
        ret=amount
        self.escrow[account_id][resource]-=amount
        
       
        if resource not in self.account[account_id]:
            self.account[account_id][resource]=0
        self.account[account_id][resource]+=amount


    def check_account_resource(self,account_id,resource,amount):
        """
        Checks if account has sufficient resources in inventory
        """
        if account_id not in self.account:
            self.account[account_id]={_bal: 0}
        #check if ressource exists
        if resource not in self.account[account_id]:
            self.account[account_id][resource]=0
        
        if self.account[account_id][resource]>=amount:
            return True
        return False

    def submit_order(self,account_id,resource,amount,price,side :int) -> Order:
        """
        Submits an order to the exchange if resources are sufficent.
        Price is price per single item.
        Returns an order if the order is active. Returns None if submit has failed.
        """
        # calculate price for complete order fullfilment
        amount=int(amount)
        full_price=price*amount


        # Move resources into escrow
        if side==Side.BUY:
            escrow=self._move_to_escrow(account_id,_bal,full_price)
        else:
            escrow=self._move_to_escrow(account_id,resource,amount)
        
        if not escrow:
            # no sufficient resources
            return None
        # create order and execude any trades
        
        order,trades=self.lme.add_order(resource,price,amount,side)

        self.orders[order.order_id]=order
        self.order_account_map[order.order_id]=account_id
        self._execute_trades(trades)
        self.calculate_resource_metrics(resource)
        return order

    def cancel_order(self,order_id):
        """
        Cancel open order.
        Returns True if order has been canceled.
        """
        if not order_id in self.orders:
            return False
        order=self.orders[order_id]
        if self._is_order_complete(order):
            return False
        order=self.lme.cancel_order(order_id,order.instmt)
        
        if order_id in self.order_trades_map: # has order any kind of trades
            trades=self.order_trades_map[order_id]
            # return assets from escrow to account
            if order.side==Side.BUY:
                total_payed=0
                total_submitted=order.price*order.qty
                for trade in trades:
                    total_payed+=trade.trade_price*trade.trade_qty
                total_diff=total_submitted-total_payed
                self._move_to_account(self.order_account_map[order_id],_bal,total_diff)
            else:
                qty_traded=0
                qty_submitted=order.qty
                for trade in trades:
                    qty_traded+=trade.trade_qty
                qty_diff=qty_submitted-qty_traded
                self._move_to_account(self.order_account_map[order_id],order.instmt,qty_diff)
        return True



    def _execute_trades(self, trades):
        for i in trades:
            self._execute_trade(i)

    def _is_order_complete(self,order: Order):
        return order.leaves_qty==0

    def _execute_trade(self, trade: Trade):
        account_id=self.order_account_map[trade.order_id]

        # calculate price for trade fullfilment
        full_price=trade.trade_price*trade.trade_qty
        exprected_price=self.orders[trade.order_id].price*trade.trade_qty
        price_diff=exprected_price-full_price

        # save trade to order trade map
        if trade.order_id not in self.order_trades_map:
            self.order_trades_map[trade.order_id]=[]

        self.order_trades_map[trade.order_id].append(trade)


        # init escrow if needet
        if trade.instmt not in self.escrow[account_id]:
            self.escrow[account_id][trade.instmt]=0

        # first edit escrow
        if trade.trade_side==Side.BUY:
            self.escrow[account_id][trade.instmt]+=trade.trade_qty
            self.escrow[account_id][_bal]-=full_price
            # move new resources to account
            self._move_to_account(account_id,trade.instmt,trade.trade_qty)
            # if deal was better then expected move money back
            self._move_to_account(account_id,_bal,price_diff)
        else:
            self.escrow[account_id][trade.instmt]-=trade.trade_qty
            self.escrow[account_id][_bal]+=full_price
            # move new bal into account
            self._move_to_account(account_id,_bal,full_price)
        
        self.executed_trades.append(trade)
        #update market rate
        self.market_rate[trade.instmt]=trade.trade_price
    

    def calculate_resource_metrics(self, resource):
        order_book = self.lme.order_books.setdefault(resource, OrderBook())
        best_bid = max(order_book.bids.keys()) if len(order_book.bids) > 0 else None
        best_ask = max(order_book.asks.keys()) if len(order_book.asks) > 0 else None
        self.total_demand[resource]=0
        self.demand[resource]={}
        for k,v in order_book.bids.items():
            self.demand[resource][k]=0
            for o in v:
                self.demand[resource][k]+=o.leaves_qty
                self.total_demand[resource]+=o.leaves_qty
        self.supply[resource]={}
        self.total_supply[resource]=0
        for k,v in order_book.asks.items():
            self.supply[resource][k]=0
            for o in v:
                self.supply[resource][k]+=o.leaves_qty
                self.total_supply[resource]+=o.leaves_qty
                
   
        self.best_ask[resource]=best_ask
        self.best_bid[resource]=best_bid
        self.traded_commoditys[resource]=True

    def get_account_resource_amount(self,account_id,resource):
        """
        Returns the amount of selected resource in account
        """
        if account_id in self.account:
            if resource in self.account[account_id]:
                return self.account[account_id][resource]
        return 0
    
    def get_total_demand(self,resource):
        if resource in self.total_demand:
            return self.total_demand[resource]
        return 0
    def get_total_supply(self,resource):
        if resource in self.total_supply:
            return self.total_supply[resource]
        return 0
    
    def log_step(self,name,episode,episode_length,step,autolog=True):
        timepoint=episode*episode_length+step
        localStepDB={}

        for instrm,order in self.traded_commoditys.items():
            data={}
            data["cxid"]=name
            data["episode"]=episode
            #data["step"]=step
            data["tstep"]=timepoint
            data["instrm"]=instrm
            data["best_ask"]=0
            data["best_bid"]=0
            data["total_demand"]=0
            data["total_supply"]=0
            data["market_rate"]=0
            localStepDB[instrm]=data
    
        ## add best ask/best bid
        for item , value in self.best_ask.items():
            if value!=None:
                localStepDB[item]["best_ask"]=int(value)
        for item , value in self.best_bid.items():
            if value!=None:
                localStepDB[item]["best_bid"]=int(value)
        #demand supply
        for item , value in self.total_demand.items():
            if value!=None:
                localStepDB[item]["total_demand"]=int(value)
        for item , value in self.total_supply.items():
            if value!=None:
                localStepDB[item]["total_supply"]=int(value)
        #last market rate
        for item , value in self.market_rate.items():
            if value!=None:
                localStepDB[item]["market_rate"]=int(value)

        if autolog:
            for id,item in localStepDB.items():
                log.EXBooksData.append(item)
        return localStepDB
    
    def log_episode(self,name,episode):
        transformed={}
        
        for t in self.executed_trades:
            if t.trade_side==1:
                # will get logged twice so only log one side    
                if t.instmt not in transformed:
                    transformed[t.instmt]={}
                  
                if t.trade_price not in transformed[t.instmt]:
                    transformed[t.instmt][t.trade_price]=0
                transformed[t.instmt][t.trade_price]+=t.trade_qty
               
        for inst,stats in transformed.items():
            for price,qty in stats.items():
                data={}
                #data["id"]=uuid.uuid4()
                data["cxid"]=name
                data["episode"]=int(episode)
                #data["order_id"]=t.order_id
                data["instmt"]=t.instmt
                data["price"]=int(price)
                data["qty"]=int(qty)
               
                #data["side"]=int(t.trade_side)
                log.EXTradeData.append(data)
           
