Untitled

package co.prey.trading.pricing

import org.roboquant.brokers.sim.NoCostPricingEngine
import org.roboquant.brokers.sim.Pricing
import org.roboquant.brokers.sim.PricingEngine
import org.roboquant.common.Logging
import org.roboquant.common.Size
import org.roboquant.feeds.OrderBook
import org.roboquant.feeds.PriceAction
import java.math.BigDecimal
import java.time.Instant

class OrderBookPricingEngine(
    private val orderBookSorted: Boolean = true
) : PricingEngine {

    private class OrderBookPricing(val orderBook: OrderBook, val sorted: Boolean) : Pricing {

        private val logger = Logging.getLogger(OrderBookPricingEngine::class)

        override fun marketPrice(size: Size): Double {
            val bigDecimalSize = size.toBigDecimal().abs()
            var index = 0
            var remainingSize = bigDecimalSize
            val entries = getSortedEntries(size)
            val pricesUsed = mutableListOf<Pair<BigDecimal, BigDecimal>>()

            while (remainingSize.signum() > 0 && index < entries.size) {
                val entry = entries[index++] // be careful of out of bounds here
                val entrySize = entry.size.toBigDecimal()
                val entryLimit = entry.limit.toBigDecimal()
                if (remainingSize < entrySize) {
                    pricesUsed.add(entryLimit to remainingSize)
                } else {
                    pricesUsed.add(entryLimit to entrySize)
                }
                remainingSize -= entrySize
            }

            if (remainingSize.signum() > 0) {
                logger.error { "Insufficient liquidity for size $size, returning last price" }
                return entries.last().limit
            }

            return pricesUsed
                .map { (price, levelSize) -> levelSize * price / bigDecimalSize }
                .fold(BigDecimal.ZERO) { prev, next -> prev + next }
                .toDouble()
        }

        private fun getSortedEntries(size: Size): List<OrderBook.OrderBookEntry> {
            val entries = if (size > 0) orderBook.asks else orderBook.bids

            return if (!sorted) {
                if (size > 0) entries.sortedBy { it.limit } else entries.sortedByDescending { it.limit }
            } else {
                entries
            }
        }
    }

    override fun getPricing(action: PriceAction, time: Instant): Pricing {
        val orderBook = action as? OrderBook ?: return NoCostPricingEngine().getPricing(action, time)
        return OrderBookPricing(orderBook, orderBookSorted)
    }
}