mev-beta/pkg/arbitrum/capital_optimizer.go

package arbitrum

import (
	"fmt"
	"math/big"
	"sort"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/common"

	"github.com/fraktal/mev-beta/internal/logger"
)

// CapitalOptimizer manages optimal capital allocation for limited budget MEV operations
type CapitalOptimizer struct {
	logger *logger.Logger
	mu     sync.RWMutex

	// Capital configuration
	totalCapital     *big.Int // Total available capital in wei (13 ETH)
	availableCapital *big.Int // Currently available capital
	reservedCapital  *big.Int // Capital reserved for gas and safety
	maxPositionSize  *big.Int // Maximum per-trade position size
	minPositionSize  *big.Int // Minimum viable position size

	// Risk management
	maxConcurrentTrades int     // Maximum number of concurrent trades
	maxCapitalPerTrade  float64 // Maximum % of capital per trade
	emergencyReserve    float64 // Emergency reserve %

	// Position tracking
	activeTrades  map[string]*ActiveTrade
	tradeHistory  []*CompletedTrade
	profitTracker *ProfitTracker

	// Performance metrics
	totalProfit      *big.Int
	totalGasCost     *big.Int
	successfulTrades uint64
	failedTrades     uint64
	startTime        time.Time
}

// ActiveTrade represents a currently executing trade
type ActiveTrade struct {
	ID                string
	TokenIn           common.Address
	TokenOut          common.Address
	AmountIn          *big.Int
	ExpectedProfit    *big.Int
	MaxGasCost        *big.Int
	StartTime         time.Time
	EstimatedDuration time.Duration
	RiskScore         float64
}

// CompletedTrade represents a finished trade for analysis
type CompletedTrade struct {
	ID            string
	TokenIn       common.Address
	TokenOut      common.Address
	AmountIn      *big.Int
	ActualProfit  *big.Int
	GasCost       *big.Int
	ExecutionTime time.Duration
	Success       bool
	Timestamp     time.Time
	ProfitMargin  float64
	ROI           float64 // Return on Investment
}

// ProfitTracker tracks profitability metrics
type ProfitTracker struct {
	DailyProfit       *big.Int
	WeeklyProfit      *big.Int
	MonthlyProfit     *big.Int
	LastResetDaily    time.Time
	LastResetWeekly   time.Time
	LastResetMonthly  time.Time
	TargetDailyProfit *big.Int // Target daily profit
}

// TradeOpportunity represents a potential arbitrage opportunity for capital allocation
type TradeOpportunity struct {
	ID              string
	TokenIn         common.Address
	TokenOut        common.Address
	AmountIn        *big.Int
	ExpectedProfit  *big.Int
	GasCost         *big.Int
	ProfitMargin    float64
	ROI             float64
	RiskScore       float64
	Confidence      float64
	ExecutionWindow time.Duration
	Priority        int
}

// NewCapitalOptimizer creates a new capital optimizer for the given budget
func NewCapitalOptimizer(logger *logger.Logger, totalCapitalETH float64) *CapitalOptimizer {
	// Convert ETH to wei
	totalCapitalWei := new(big.Int).Mul(
		big.NewInt(int64(totalCapitalETH*1e18)),
		big.NewInt(1))

	// Reserve 10% for gas and emergencies
	emergencyReserve := 0.10
	reservedWei := new(big.Int).Div(
		new(big.Int).Mul(totalCapitalWei, big.NewInt(10)),
		big.NewInt(100))

	availableWei := new(big.Int).Sub(totalCapitalWei, reservedWei)

	// Set position size limits (optimized for $13 ETH budget)
	maxPositionWei := new(big.Int).Div(totalCapitalWei, big.NewInt(4))   // Max 25% per trade
	minPositionWei := new(big.Int).Div(totalCapitalWei, big.NewInt(100)) // Min 1% per trade

	// Daily profit target: 1-3% of total capital
	targetDailyProfitWei := new(big.Int).Div(
		new(big.Int).Mul(totalCapitalWei, big.NewInt(2)), // 2% target
		big.NewInt(100))

	return &CapitalOptimizer{
		logger:              logger,
		totalCapital:        totalCapitalWei,
		availableCapital:    availableWei,
		reservedCapital:     reservedWei,
		maxPositionSize:     maxPositionWei,
		minPositionSize:     minPositionWei,
		maxConcurrentTrades: 3,    // Conservative for limited capital
		maxCapitalPerTrade:  0.25, // 25% max per trade
		emergencyReserve:    emergencyReserve,
		activeTrades:        make(map[string]*ActiveTrade),
		tradeHistory:        make([]*CompletedTrade, 0),
		totalProfit:         big.NewInt(0),
		totalGasCost:        big.NewInt(0),
		startTime:           time.Now(),
		profitTracker: &ProfitTracker{
			DailyProfit:       big.NewInt(0),
			WeeklyProfit:      big.NewInt(0),
			MonthlyProfit:     big.NewInt(0),
			LastResetDaily:    time.Now(),
			LastResetWeekly:   time.Now(),
			LastResetMonthly:  time.Now(),
			TargetDailyProfit: targetDailyProfitWei,
		},
	}
}

// CanExecuteTrade checks if a trade can be executed with current capital constraints
func (co *CapitalOptimizer) CanExecuteTrade(opportunity *TradeOpportunity) (bool, string) {
	co.mu.RLock()
	defer co.mu.RUnlock()

	// Check if we have enough concurrent trade slots
	if len(co.activeTrades) >= co.maxConcurrentTrades {
		return false, "maximum concurrent trades reached"
	}

	// Check if we have sufficient capital
	totalRequired := new(big.Int).Add(opportunity.AmountIn, opportunity.GasCost)
	if totalRequired.Cmp(co.availableCapital) > 0 {
		return false, "insufficient available capital"
	}

	// Check position size limits
	if opportunity.AmountIn.Cmp(co.maxPositionSize) > 0 {
		return false, "position size exceeds maximum limit"
	}

	if opportunity.AmountIn.Cmp(co.minPositionSize) < 0 {
		return false, "position size below minimum threshold"
	}

	// Check profitability after gas costs
	netProfit := new(big.Int).Sub(opportunity.ExpectedProfit, opportunity.GasCost)
	if netProfit.Sign() <= 0 {
		return false, "trade not profitable after gas costs"
	}

	// Check ROI threshold (minimum 0.5% ROI for limited capital)
	minROI := 0.005
	if opportunity.ROI < minROI {
		return false, fmt.Sprintf("ROI %.3f%% below minimum threshold %.3f%%", opportunity.ROI*100, minROI*100)
	}

	// Check risk score (maximum 0.7 for conservative trading)
	maxRisk := 0.7
	if opportunity.RiskScore > maxRisk {
		return false, fmt.Sprintf("risk score %.3f exceeds maximum %.3f", opportunity.RiskScore, maxRisk)
	}

	return true, ""
}

// AllocateCapital allocates capital for a trade and returns the optimal position size
func (co *CapitalOptimizer) AllocateCapital(opportunity *TradeOpportunity) (*big.Int, error) {
	co.mu.Lock()
	defer co.mu.Unlock()

	// Check if trade can be executed
	canExecute, reason := co.CanExecuteTrade(opportunity)
	if !canExecute {
		return nil, fmt.Errorf("cannot execute trade: %s", reason)
	}

	// Calculate optimal position size based on Kelly criterion and risk management
	optimalSize := co.calculateOptimalPositionSize(opportunity)

	// Ensure position size is within bounds
	if optimalSize.Cmp(co.maxPositionSize) > 0 {
		optimalSize = new(big.Int).Set(co.maxPositionSize)
	}
	if optimalSize.Cmp(co.minPositionSize) < 0 {
		optimalSize = new(big.Int).Set(co.minPositionSize)
	}

	// Reserve capital for this trade
	totalRequired := new(big.Int).Add(optimalSize, opportunity.GasCost)
	co.availableCapital = new(big.Int).Sub(co.availableCapital, totalRequired)

	// Create active trade record
	activeTrade := &ActiveTrade{
		ID:                opportunity.ID,
		TokenIn:           opportunity.TokenIn,
		TokenOut:          opportunity.TokenOut,
		AmountIn:          optimalSize,
		ExpectedProfit:    opportunity.ExpectedProfit,
		MaxGasCost:        opportunity.GasCost,
		StartTime:         time.Now(),
		EstimatedDuration: opportunity.ExecutionWindow,
		RiskScore:         opportunity.RiskScore,
	}

	co.activeTrades[opportunity.ID] = activeTrade

	co.logger.Info(fmt.Sprintf("💰 CAPITAL ALLOCATED: $%.2f for trade %s (%.1f%% of capital, ROI: %.2f%%)",
		co.weiToUSD(optimalSize), opportunity.ID[:8],
		co.getCapitalPercentage(optimalSize)*100, opportunity.ROI*100))

	return optimalSize, nil
}

// CompleteTradeand updates capital allocation
func (co *CapitalOptimizer) CompleteTrade(tradeID string, actualProfit *big.Int, gasCost *big.Int, success bool) {
	co.mu.Lock()
	defer co.mu.Unlock()

	activeTrade, exists := co.activeTrades[tradeID]
	if !exists {
		co.logger.Warn(fmt.Sprintf("Trade %s not found in active trades", tradeID))
		return
	}

	// Return capital to available pool
	capitalReturned := activeTrade.AmountIn
	if success && actualProfit.Sign() > 0 {
		// Add profit to returned capital
		capitalReturned = new(big.Int).Add(capitalReturned, actualProfit)
	}

	co.availableCapital = new(big.Int).Add(co.availableCapital, capitalReturned)

	// Update profit tracking
	netProfit := new(big.Int).Sub(actualProfit, gasCost)
	if success && netProfit.Sign() > 0 {
		co.totalProfit = new(big.Int).Add(co.totalProfit, netProfit)
		co.profitTracker.DailyProfit = new(big.Int).Add(co.profitTracker.DailyProfit, netProfit)
		co.profitTracker.WeeklyProfit = new(big.Int).Add(co.profitTracker.WeeklyProfit, netProfit)
		co.profitTracker.MonthlyProfit = new(big.Int).Add(co.profitTracker.MonthlyProfit, netProfit)
		co.successfulTrades++
	} else {
		co.failedTrades++
	}

	co.totalGasCost = new(big.Int).Add(co.totalGasCost, gasCost)

	// Create completed trade record
	executionTime := time.Since(activeTrade.StartTime)
	roi := 0.0
	if activeTrade.AmountIn.Sign() > 0 {
		roi = float64(netProfit.Int64()) / float64(activeTrade.AmountIn.Int64())
	}

	completedTrade := &CompletedTrade{
		ID:            tradeID,
		TokenIn:       activeTrade.TokenIn,
		TokenOut:      activeTrade.TokenOut,
		AmountIn:      activeTrade.AmountIn,
		ActualProfit:  actualProfit,
		GasCost:       gasCost,
		ExecutionTime: executionTime,
		Success:       success,
		Timestamp:     time.Now(),
		ProfitMargin:  float64(netProfit.Int64()) / float64(actualProfit.Int64()),
		ROI:           roi,
	}

	co.tradeHistory = append(co.tradeHistory, completedTrade)

	// Remove from active trades
	delete(co.activeTrades, tradeID)

	// Log completion
	if success {
		co.logger.Info(fmt.Sprintf("✅ TRADE COMPLETED: %s, Profit: $%.2f, ROI: %.2f%%, Time: %v",
			tradeID[:8], co.weiToUSD(netProfit), roi*100, executionTime))
	} else {
		co.logger.Error(fmt.Sprintf("❌ TRADE FAILED: %s, Loss: $%.2f, Time: %v",
			tradeID[:8], co.weiToUSD(gasCost), executionTime))
	}

	// Check profit targets and adjust strategy if needed
	co.checkProfitTargets()
}

// calculateOptimalPositionSize calculates optimal position size using modified Kelly criterion
func (co *CapitalOptimizer) calculateOptimalPositionSize(opportunity *TradeOpportunity) *big.Int {
	// Modified Kelly Criterion: f = (bp - q) / b
	// where b = odds received on the wager, p = probability of winning, q = probability of losing

	// Convert confidence to win probability
	winProbability := opportunity.Confidence
	lossProbability := 1.0 - winProbability

	// Calculate odds from ROI
	odds := opportunity.ROI
	if odds <= 0 {
		return co.minPositionSize
	}

	// Kelly fraction
	kellyFraction := (odds*winProbability - lossProbability) / odds

	// Apply conservative scaling (25% of Kelly for risk management)
	conservativeKelly := kellyFraction * 0.25

	// Ensure we don't bet more than max position size
	if conservativeKelly > co.maxCapitalPerTrade {
		conservativeKelly = co.maxCapitalPerTrade
	}

	// Apply risk adjustment
	riskAdjustment := 1.0 - opportunity.RiskScore
	conservativeKelly *= riskAdjustment

	// Calculate position size
	positionSize := new(big.Int).Mul(
		co.availableCapital,
		big.NewInt(int64(conservativeKelly*1000)),
	)
	positionSize = new(big.Int).Div(positionSize, big.NewInt(1000))

	// Ensure minimum viability
	if positionSize.Cmp(co.minPositionSize) < 0 {
		positionSize = new(big.Int).Set(co.minPositionSize)
	}

	return positionSize
}

// GetOptimalOpportunities returns prioritized opportunities based on capital allocation strategy
func (co *CapitalOptimizer) GetOptimalOpportunities(opportunities []*TradeOpportunity) []*TradeOpportunity {
	co.mu.RLock()
	defer co.mu.RUnlock()

	// Filter opportunities that can be executed
	var viable []*TradeOpportunity
	for _, opp := range opportunities {
		if canExecute, _ := co.CanExecuteTrade(opp); canExecute {
			viable = append(viable, opp)
		}
	}

	if len(viable) == 0 {
		return []*TradeOpportunity{}
	}

	// Sort by profitability score (ROI * Confidence / Risk)
	sort.Slice(viable, func(i, j int) bool {
		scoreI := (viable[i].ROI * viable[i].Confidence) / (1.0 + viable[i].RiskScore)
		scoreJ := (viable[j].ROI * viable[j].Confidence) / (1.0 + viable[j].RiskScore)
		return scoreI > scoreJ
	})

	// Return top opportunities that fit within concurrent trade limits
	maxReturn := co.maxConcurrentTrades - len(co.activeTrades)
	if len(viable) > maxReturn {
		viable = viable[:maxReturn]
	}

	return viable
}

// Helper methods

func (co *CapitalOptimizer) weiToUSD(wei *big.Int) float64 {
	// Assume ETH = $2000 for rough USD calculations
	ethPrice := 2000.0
	ethAmount := new(big.Float).Quo(new(big.Float).SetInt(wei), big.NewFloat(1e18))
	ethFloat, _ := ethAmount.Float64()
	return ethFloat * ethPrice
}

func (co *CapitalOptimizer) getCapitalPercentage(amount *big.Int) float64 {
	ratio := new(big.Float).Quo(new(big.Float).SetInt(amount), new(big.Float).SetInt(co.totalCapital))
	percentage, _ := ratio.Float64()
	return percentage
}

func (co *CapitalOptimizer) checkProfitTargets() {
	now := time.Now()

	// Reset daily profits if needed
	if now.Sub(co.profitTracker.LastResetDaily) >= 24*time.Hour {
		co.profitTracker.DailyProfit = big.NewInt(0)
		co.profitTracker.LastResetDaily = now
	}

	// Check if we've hit daily target
	if co.profitTracker.DailyProfit.Cmp(co.profitTracker.TargetDailyProfit) >= 0 {
		co.logger.Info(fmt.Sprintf("🎯 DAILY PROFIT TARGET ACHIEVED: $%.2f (target: $%.2f)",
			co.weiToUSD(co.profitTracker.DailyProfit),
			co.weiToUSD(co.profitTracker.TargetDailyProfit)))
	}
}

// GetStatus returns current capital allocation status
func (co *CapitalOptimizer) GetStatus() map[string]interface{} {
	co.mu.RLock()
	defer co.mu.RUnlock()

	totalRuntime := time.Since(co.startTime)
	successRate := 0.0
	if co.successfulTrades+co.failedTrades > 0 {
		successRate = float64(co.successfulTrades) / float64(co.successfulTrades+co.failedTrades)
	}

	netProfit := new(big.Int).Sub(co.totalProfit, co.totalGasCost)

	return map[string]interface{}{
		"total_capital_usd":     co.weiToUSD(co.totalCapital),
		"available_capital_usd": co.weiToUSD(co.availableCapital),
		"reserved_capital_usd":  co.weiToUSD(co.reservedCapital),
		"active_trades":         len(co.activeTrades),
		"max_concurrent_trades": co.maxConcurrentTrades,
		"successful_trades":     co.successfulTrades,
		"failed_trades":         co.failedTrades,
		"success_rate":          successRate,
		"total_profit_usd":      co.weiToUSD(co.totalProfit),
		"total_gas_cost_usd":    co.weiToUSD(co.totalGasCost),
		"net_profit_usd":        co.weiToUSD(netProfit),
		"daily_profit_usd":      co.weiToUSD(co.profitTracker.DailyProfit),
		"daily_target_usd":      co.weiToUSD(co.profitTracker.TargetDailyProfit),
		"runtime_hours":         totalRuntime.Hours(),
		"capital_utilization":   co.getCapitalPercentage(new(big.Int).Sub(co.totalCapital, co.availableCapital)),
	}
}