Untitled

package co.prey.trading.broker

import org.roboquant.brokers.Account
import org.roboquant.brokers.Broker
import org.roboquant.brokers.sim.AccountModel
import org.roboquant.brokers.sim.CashAccount
import org.roboquant.brokers.sim.FeeModel
import org.roboquant.brokers.sim.NoFeeModel
import org.roboquant.brokers.sim.PricingEngine
import org.roboquant.brokers.sim.SimBroker
import org.roboquant.brokers.sim.SpreadPricingEngine
import org.roboquant.common.Currency
import org.roboquant.common.USD
import org.roboquant.common.Wallet
import org.roboquant.feeds.Event
import org.roboquant.orders.Order
import org.roboquant.orders.OrderState
import java.time.Duration
import java.time.Instant
import java.util.PriorityQueue
import java.util.Queue

class DelayedSimBroker(
    initialDeposit: Wallet = Wallet(1_000_000.00.USD),
    baseCurrency: Currency = initialDeposit.currencies.first(),
    feeModel: FeeModel = NoFeeModel(),
    accountModel: AccountModel = CashAccount(),
    pricingEngine: PricingEngine = SpreadPricingEngine(),
    private val artificialTradeDelay: Duration = Duration.ofMillis(10),
    private val simBroker: SimBroker = SimBroker(initialDeposit, baseCurrency, feeModel, accountModel, pricingEngine)
) : Broker by simBroker {

    private val orderQueue: Queue<Pair<Order, Instant>> = PriorityQueue { a, b ->
        a.second.compareTo(b.second)
    }
    private val openOrders: MutableMap<Int, OrderState> = mutableMapOf()

    override val account: Account
        get() {
            val simAccount = simBroker.account
            return Account(
                baseCurrency = simAccount.baseCurrency,
                lastUpdate = simAccount.lastUpdate,
                cash = simAccount.cash,
                trades = simAccount.trades,
                openOrders = simAccount.openOrders + openOrders.values,
                closedOrders = simAccount.closedOrders,
                positions = simAccount.positions,
                buyingPower = simAccount.buyingPower
            )
        }

    override fun place(orders: List<Order>, event: Event): Account {
        orderQueue.addAll(orders.map { it to (event.time + artificialTradeDelay) })
        orders.forEach { openOrders[it.id] = OrderState(it) }

        val ordersToProcess = mutableListOf<Order>()
        while (orderQueue.isNotEmpty() && orderQueue.peek().second <= event.time) {
            val (order) = orderQueue.poll()
            ordersToProcess.add(order)
            openOrders.remove(order.id)
        }

        simBroker.place(ordersToProcess, event)
        return account
    }

    fun closePositions(): Account {
        simBroker.closePositions()
        return account
    }
}