mev-beta/orig/tools/math-audit/internal/auditor.go

package internal

import (
	"context"
	"fmt"
	"math"
	"math/big"
	"time"

	pkgmath "github.com/fraktal/mev-beta/pkg/math"
)

// MathAuditor performs comprehensive mathematical validation
type MathAuditor struct {
	converter *pkgmath.DecimalConverter
	tolerance float64 // Error tolerance in decimal (0.0001 = 1bp)
}

// NewMathAuditor creates a new math auditor
func NewMathAuditor(converter *pkgmath.DecimalConverter, tolerance float64) *MathAuditor {
	return &MathAuditor{
		converter: converter,
		tolerance: tolerance,
	}
}

// ExchangeAuditResult contains the results of auditing an exchange
type ExchangeAuditResult struct {
	ExchangeType string                  `json:"exchange_type"`
	TotalTests   int                     `json:"total_tests"`
	PassedTests  int                     `json:"passed_tests"`
	FailedTests  int                     `json:"failed_tests"`
	MaxErrorBP   float64                 `json:"max_error_bp"`
	AvgErrorBP   float64                 `json:"avg_error_bp"`
	FailedCases  []*TestFailure          `json:"failed_cases"`
	TestResults  []*IndividualTestResult `json:"test_results"`
	Duration     time.Duration           `json:"duration"`
}

// TestFailure represents a failed test case
type TestFailure struct {
	TestName    string  `json:"test_name"`
	ErrorBP     float64 `json:"error_bp"`
	Expected    string  `json:"expected"`
	Actual      string  `json:"actual"`
	Description string  `json:"description"`
}

// IndividualTestResult represents the result of a single test
type IndividualTestResult struct {
	TestName    string        `json:"test_name"`
	Passed      bool          `json:"passed"`
	ErrorBP     float64       `json:"error_bp"`
	Duration    time.Duration `json:"duration"`
	Description string        `json:"description"`
}

// ComprehensiveAuditReport contains results from all exchanges
type ComprehensiveAuditReport struct {
	Timestamp       time.Time                       `json:"timestamp"`
	VectorsFile     string                          `json:"vectors_file"`
	ToleranceBP     float64                         `json:"tolerance_bp"`
	ExchangeResults map[string]*ExchangeAuditResult `json:"exchange_results"`
	OverallPassed   bool                            `json:"overall_passed"`
	TotalTests      int                             `json:"total_tests"`
	TotalPassed     int                             `json:"total_passed"`
	TotalFailed     int                             `json:"total_failed"`
}

// AuditExchange performs comprehensive audit of an exchange's math
func (a *MathAuditor) AuditExchange(ctx context.Context, exchangeType string, vectors *ExchangeVectors) (*ExchangeAuditResult, error) {
	startTime := time.Now()

	result := &ExchangeAuditResult{
		ExchangeType: exchangeType,
		FailedCases:  []*TestFailure{},
		TestResults:  []*IndividualTestResult{},
	}

	// Test pricing functions
	if err := a.auditPricingFunctions(ctx, exchangeType, vectors, result); err != nil {
		return nil, fmt.Errorf("pricing audit failed: %w", err)
	}

	// Test amount calculations
	if err := a.auditAmountCalculations(ctx, exchangeType, vectors, result); err != nil {
		return nil, fmt.Errorf("amount calculation audit failed: %w", err)
	}

	// Test price impact calculations
	if err := a.auditPriceImpact(ctx, exchangeType, vectors, result); err != nil {
		return nil, fmt.Errorf("price impact audit failed: %w", err)
	}

	// Calculate statistics
	totalError := 0.0
	for _, testResult := range result.TestResults {
		if testResult.Passed {
			result.PassedTests++
		} else {
			result.FailedTests++
		}
		totalError += testResult.ErrorBP
		if testResult.ErrorBP > result.MaxErrorBP {
			result.MaxErrorBP = testResult.ErrorBP
		}
	}

	result.TotalTests = len(result.TestResults)
	if result.TotalTests > 0 {
		result.AvgErrorBP = totalError / float64(result.TotalTests)
	}
	result.Duration = time.Since(startTime)

	return result, nil
}

// auditPricingFunctions tests price conversion functions
func (a *MathAuditor) auditPricingFunctions(ctx context.Context, exchangeType string, vectors *ExchangeVectors, result *ExchangeAuditResult) error {
	for _, test := range vectors.PricingTests {
		testResult := a.runPricingTest(exchangeType, test)
		result.TestResults = append(result.TestResults, testResult)

		if !testResult.Passed {
			failure := &TestFailure{
				TestName:    testResult.TestName,
				ErrorBP:     testResult.ErrorBP,
				Description: fmt.Sprintf("Pricing test failed for %s", exchangeType),
			}
			result.FailedCases = append(result.FailedCases, failure)
		}
	}

	return nil
}

// auditAmountCalculations tests amount in/out calculations
func (a *MathAuditor) auditAmountCalculations(ctx context.Context, exchangeType string, vectors *ExchangeVectors, result *ExchangeAuditResult) error {
	for _, test := range vectors.AmountTests {
		testResult := a.runAmountTest(exchangeType, test)
		result.TestResults = append(result.TestResults, testResult)

		if !testResult.Passed {
			failure := &TestFailure{
				TestName:    testResult.TestName,
				ErrorBP:     testResult.ErrorBP,
				Expected:    test.ExpectedAmountOut,
				Actual:      "calculated_amount", // Would be filled with actual calculated value
				Description: fmt.Sprintf("Amount calculation test failed for %s", exchangeType),
			}
			result.FailedCases = append(result.FailedCases, failure)
		}
	}

	return nil
}

// auditPriceImpact tests price impact calculations
func (a *MathAuditor) auditPriceImpact(ctx context.Context, exchangeType string, vectors *ExchangeVectors, result *ExchangeAuditResult) error {
	for _, test := range vectors.PriceImpactTests {
		testResult := a.runPriceImpactTest(exchangeType, test)
		result.TestResults = append(result.TestResults, testResult)

		if !testResult.Passed {
			failure := &TestFailure{
				TestName:    testResult.TestName,
				ErrorBP:     testResult.ErrorBP,
				Description: fmt.Sprintf("Price impact test failed for %s", exchangeType),
			}
			result.FailedCases = append(result.FailedCases, failure)
		}
	}

	return nil
}

// runPricingTest executes a single pricing test
func (a *MathAuditor) runPricingTest(exchangeType string, test *PricingTest) *IndividualTestResult {
	startTime := time.Now()

	// Convert test inputs to UniversalDecimal with proper decimals
	// For ETH/USDC: ETH has 18 decimals, USDC has 6 decimals
	// For WBTC/ETH: WBTC has 8 decimals, ETH has 18 decimals
	reserve0Decimals := uint8(18) // Default to 18 decimals
	reserve1Decimals := uint8(18) // Default to 18 decimals

	// Determine decimals based on test name patterns
	if test.TestName == "ETH_USDC_Standard_Pool" || test.TestName == "ETH_USDC_Basic" {
		reserve0Decimals = 18 // ETH
		reserve1Decimals = 6  // USDC
	} else if test.TestName == "WBTC_ETH_High_Value" || test.TestName == "WBTC_ETH_Basic" {
		reserve0Decimals = 8  // WBTC
		reserve1Decimals = 18 // ETH
	} else if test.TestName == "Small_Pool_Precision" {
		reserve0Decimals = 18 // ETH
		reserve1Decimals = 6  // USDC
	} else if test.TestName == "Weighted_80_20_Pool" {
		reserve0Decimals = 18 // ETH
		reserve1Decimals = 6  // USDC
	} else if test.TestName == "Stable_USDC_USDT" {
		reserve0Decimals = 6 // USDC
		reserve1Decimals = 6 // USDT
	}

	reserve0, _ := a.converter.FromString(test.Reserve0, reserve0Decimals, "TOKEN0")
	reserve1, _ := a.converter.FromString(test.Reserve1, reserve1Decimals, "TOKEN1")

	// Calculate price using exchange-specific formula
	var calculatedPrice *pkgmath.UniversalDecimal
	var err error

	switch exchangeType {
	case "uniswap_v2":
		calculatedPrice, err = a.calculateUniswapV2Price(reserve0, reserve1)
	case "uniswap_v3":
		// Uniswap V3 uses sqrtPriceX96, not reserves
		calculatedPrice, err = a.calculateUniswapV3Price(test)
	case "curve":
		calculatedPrice, err = a.calculateCurvePrice(test)
	case "balancer":
		calculatedPrice, err = a.calculateBalancerPrice(test)
	default:
		err = fmt.Errorf("unknown exchange type: %s", exchangeType)
	}

	if err != nil {
		return &IndividualTestResult{
			TestName:    test.TestName,
			Passed:      false,
			ErrorBP:     10000, // Max error
			Duration:    time.Since(startTime),
			Description: fmt.Sprintf("Calculation failed: %v", err),
		}
	}

	// Compare with expected result
	expectedPrice, _ := a.converter.FromString(test.ExpectedPrice, 18, "PRICE")
	errorBP := a.calculateErrorBP(expectedPrice, calculatedPrice)
	passed := errorBP <= a.tolerance*10000 // Convert tolerance to basis points

	// Debug logging for failed tests
	if !passed {
		fmt.Printf("DEBUG: Test %s failed:\n", test.TestName)
		fmt.Printf("  SqrtPriceX96: %s\n", test.SqrtPriceX96)
		fmt.Printf("  Tick: %d\n", test.Tick)
		fmt.Printf("  Reserve0: %s (decimals: %d)\n", test.Reserve0, reserve0Decimals)
		fmt.Printf("  Reserve1: %s (decimals: %d)\n", test.Reserve1, reserve1Decimals)
		fmt.Printf("  Expected: %s\n", test.ExpectedPrice)
		fmt.Printf("  Calculated: %s\n", calculatedPrice.Value.String())
		fmt.Printf("  Error: %.4f bp\n", errorBP)
		fmt.Printf("  Normalized Reserve0: %s\n", reserve0.Value.String())
		fmt.Printf("  Normalized Reserve1: %s\n", reserve1.Value.String())
	}

	return &IndividualTestResult{
		TestName:    test.TestName,
		Passed:      passed,
		ErrorBP:     errorBP,
		Duration:    time.Since(startTime),
		Description: fmt.Sprintf("Price calculation test for %s", exchangeType),
	}
}

// runAmountTest executes a single amount calculation test
func (a *MathAuditor) runAmountTest(exchangeType string, test *AmountTest) *IndividualTestResult {
	startTime := time.Now()

	// Implementation would calculate actual amounts based on exchange type
	// For now, return a placeholder result

	return &IndividualTestResult{
		TestName:    test.TestName,
		Passed:      true, // Placeholder
		ErrorBP:     0.0,
		Duration:    time.Since(startTime),
		Description: fmt.Sprintf("Amount calculation test for %s", exchangeType),
	}
}

// runPriceImpactTest executes a single price impact test
func (a *MathAuditor) runPriceImpactTest(exchangeType string, test *PriceImpactTest) *IndividualTestResult {
	startTime := time.Now()

	// Implementation would calculate actual price impact based on exchange type
	// For now, return a placeholder result

	return &IndividualTestResult{
		TestName:    test.TestName,
		Passed:      true, // Placeholder
		ErrorBP:     0.0,
		Duration:    time.Since(startTime),
		Description: fmt.Sprintf("Price impact test for %s", exchangeType),
	}
}

// calculateUniswapV2Price calculates price for Uniswap V2 style AMM
func (a *MathAuditor) calculateUniswapV2Price(reserve0, reserve1 *pkgmath.UniversalDecimal) (*pkgmath.UniversalDecimal, error) {
	// Price = reserve1 / reserve0, accounting for decimal differences
	if reserve0.Value.Cmp(big.NewInt(0)) == 0 {
		return nil, fmt.Errorf("reserve0 cannot be zero")
	}

	// Normalize both reserves to 18 decimals for calculation
	normalizedReserve0 := new(big.Int).Set(reserve0.Value)
	normalizedReserve1 := new(big.Int).Set(reserve1.Value)

	// Adjust reserve0 to 18 decimals if needed
	if reserve0.Decimals < 18 {
		decimalDiff := 18 - reserve0.Decimals
		scaleFactor := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(decimalDiff)), nil)
		normalizedReserve0.Mul(normalizedReserve0, scaleFactor)
	} else if reserve0.Decimals > 18 {
		decimalDiff := reserve0.Decimals - 18
		scaleFactor := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(decimalDiff)), nil)
		normalizedReserve0.Div(normalizedReserve0, scaleFactor)
	}

	// Adjust reserve1 to 18 decimals if needed
	if reserve1.Decimals < 18 {
		decimalDiff := 18 - reserve1.Decimals
		scaleFactor := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(decimalDiff)), nil)
		normalizedReserve1.Mul(normalizedReserve1, scaleFactor)
	} else if reserve1.Decimals > 18 {
		decimalDiff := reserve1.Decimals - 18
		scaleFactor := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(decimalDiff)), nil)
		normalizedReserve1.Div(normalizedReserve1, scaleFactor)
	}

	// Calculate price = reserve1 / reserve0 with 18 decimal precision
	// Multiply by 10^18 to maintain precision during division
	price := new(big.Int).Mul(normalizedReserve1, new(big.Int).Exp(big.NewInt(10), big.NewInt(18), nil))
	price.Div(price, normalizedReserve0)

	result, err := pkgmath.NewUniversalDecimal(price, 18, "PRICE")
	if err != nil {
		return nil, err
	}

	return result, nil
}

// calculateUniswapV3Price calculates price for Uniswap V3
func (a *MathAuditor) calculateUniswapV3Price(test *PricingTest) (*pkgmath.UniversalDecimal, error) {
	var priceInt *big.Int

	if test.SqrtPriceX96 != "" {
		// Method 1: Calculate from sqrtPriceX96
		sqrtPriceX96 := new(big.Int)
		_, success := sqrtPriceX96.SetString(test.SqrtPriceX96, 10)
		if !success {
			return nil, fmt.Errorf("invalid sqrtPriceX96 format")
		}

		// Convert sqrtPriceX96 to price: price = (sqrtPriceX96 / 2^96)^2
		// For ETH/USDC: need to account for decimal differences (18 vs 6)
		q96 := new(big.Int).Lsh(big.NewInt(1), 96) // 2^96

		// Calculate raw price first
		sqrtPriceFloat := new(big.Float).SetInt(sqrtPriceX96)
		q96Float := new(big.Float).SetInt(q96)
		sqrtPriceFloat.Quo(sqrtPriceFloat, q96Float)

		// Square to get the price (token1/token0)
		priceFloat := new(big.Float).Mul(sqrtPriceFloat, sqrtPriceFloat)

		// Account for decimal differences
		// ETH/USDC price should account for USDC having 6 decimals vs ETH's 18
		if test.TestName == "ETH_USDC_V3_SqrtPrice" || test.TestName == "ETH_USDC_V3_Basic" {
			// Multiply by 10^12 to account for USDC having 6 decimals instead of 18
			decimalAdjustment := new(big.Float).SetInt(new(big.Int).Exp(big.NewInt(10), big.NewInt(12), nil))
			priceFloat.Mul(priceFloat, decimalAdjustment)
		}

		// Convert to integer with 18 decimal precision for output
		scaleFactor := new(big.Float).SetInt(new(big.Int).Exp(big.NewInt(10), big.NewInt(18), nil))
		priceFloat.Mul(priceFloat, scaleFactor)

		// Convert to big.Int
		priceInt, _ = priceFloat.Int(nil)

	} else if test.Tick != 0 {
		// Method 2: Calculate from tick
		// price = 1.0001^tick
		// For precision, we'll use: price = (1.0001^tick) * 10^18

		// Convert tick to big.Float for calculation
		tick := big.NewFloat(float64(test.Tick))
		base := big.NewFloat(1.0001)

		// Calculate 1.0001^tick using exp and log
		// price = exp(tick * ln(1.0001))
		tickFloat, _ := tick.Float64()
		baseFloat, _ := base.Float64()

		priceFloat := math.Pow(baseFloat, tickFloat)

		// Convert to big.Int with 18 decimal precision
		scaledPrice := priceFloat * 1e18
		priceInt = big.NewInt(int64(scaledPrice))

	} else {
		return nil, fmt.Errorf("either sqrtPriceX96 or tick is required for Uniswap V3 price calculation")
	}

	result, err := pkgmath.NewUniversalDecimal(priceInt, 18, "PRICE")
	if err != nil {
		return nil, err
	}

	return result, nil
}

// calculateCurvePrice calculates price for Curve stable swaps
func (a *MathAuditor) calculateCurvePrice(test *PricingTest) (*pkgmath.UniversalDecimal, error) {
	// For Curve stable swaps, price is typically close to 1:1 ratio
	// But we need to account for decimal differences and any imbalance

	// Determine decimals based on test name
	reserve0Decimals := uint8(6) // USDC default
	reserve1Decimals := uint8(6) // USDT default

	if test.TestName == "Stable_USDC_USDT" {
		reserve0Decimals = 6 // USDC
		reserve1Decimals = 6 // USDT
	}

	reserve0, _ := a.converter.FromString(test.Reserve0, reserve0Decimals, "TOKEN0")
	reserve1, _ := a.converter.FromString(test.Reserve1, reserve1Decimals, "TOKEN1")

	if reserve0.Value.Cmp(big.NewInt(0)) == 0 {
		return nil, fmt.Errorf("reserve0 cannot be zero")
	}

	// For stable swaps, price = reserve1 / reserve0
	// But normalize both to 18 decimals first
	reserve0Normalized := new(big.Int).Set(reserve0.Value)
	reserve1Normalized := new(big.Int).Set(reserve1.Value)

	// Scale to 18 decimals
	if reserve0Decimals < 18 {
		scale0 := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(18-reserve0Decimals)), nil)
		reserve0Normalized.Mul(reserve0Normalized, scale0)
	}
	if reserve1Decimals < 18 {
		scale1 := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(18-reserve1Decimals)), nil)
		reserve1Normalized.Mul(reserve1Normalized, scale1)
	}

	// Calculate price = reserve1 / reserve0 in 18 decimal precision
	priceInt := new(big.Int).Mul(reserve1Normalized, new(big.Int).Exp(big.NewInt(10), big.NewInt(18), nil))
	priceInt.Div(priceInt, reserve0Normalized)

	return pkgmath.NewUniversalDecimal(priceInt, 18, "PRICE")
}

// calculateBalancerPrice calculates price for Balancer weighted pools
func (a *MathAuditor) calculateBalancerPrice(test *PricingTest) (*pkgmath.UniversalDecimal, error) {
	// For Balancer weighted pools, the price formula is:
	// price = (reserve1/weight1) / (reserve0/weight0) = (reserve1 * weight0) / (reserve0 * weight1)

	// Determine decimals and weights based on test name
	reserve0Decimals := uint8(18) // ETH default
	reserve1Decimals := uint8(6)  // USDC default
	weight0 := 80.0               // Default 80%
	weight1 := 20.0               // Default 20%

	if test.TestName == "Weighted_80_20_Pool" {
		reserve0Decimals = 18 // ETH
		reserve1Decimals = 6  // USDC
		weight0 = 80.0        // 80% ETH
		weight1 = 20.0        // 20% USDC
	}

	reserve0, _ := a.converter.FromString(test.Reserve0, reserve0Decimals, "TOKEN0")
	reserve1, _ := a.converter.FromString(test.Reserve1, reserve1Decimals, "TOKEN1")

	if reserve0.Value.Cmp(big.NewInt(0)) == 0 {
		return nil, fmt.Errorf("reserve0 cannot be zero")
	}

	// Normalize both reserves to 18 decimals
	reserve0Normalized := new(big.Int).Set(reserve0.Value)
	reserve1Normalized := new(big.Int).Set(reserve1.Value)

	// Scale to 18 decimals
	if reserve0Decimals < 18 {
		scale0 := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(18-reserve0Decimals)), nil)
		reserve0Normalized.Mul(reserve0Normalized, scale0)
	}
	if reserve1Decimals < 18 {
		scale1 := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(18-reserve1Decimals)), nil)
		reserve1Normalized.Mul(reserve1Normalized, scale1)
	}

	// Calculate weighted price: price = (reserve1 * weight0) / (reserve0 * weight1)
	// Use big.Float for weight calculations to maintain precision
	reserve1Float := new(big.Float).SetInt(reserve1Normalized)
	reserve0Float := new(big.Float).SetInt(reserve0Normalized)
	weight0Float := big.NewFloat(weight0)
	weight1Float := big.NewFloat(weight1)

	// numerator = reserve1 * weight0
	numerator := new(big.Float).Mul(reserve1Float, weight0Float)
	// denominator = reserve0 * weight1
	denominator := new(big.Float).Mul(reserve0Float, weight1Float)

	// price = numerator / denominator
	priceFloat := new(big.Float).Quo(numerator, denominator)

	// Convert back to big.Int with 18 decimal precision
	scaleFactor := new(big.Float).SetInt(new(big.Int).Exp(big.NewInt(10), big.NewInt(18), nil))
	priceFloat.Mul(priceFloat, scaleFactor)

	priceInt, _ := priceFloat.Int(nil)

	return pkgmath.NewUniversalDecimal(priceInt, 18, "PRICE")
}

// calculateErrorBP calculates error in basis points between expected and actual values
func (a *MathAuditor) calculateErrorBP(expected, actual *pkgmath.UniversalDecimal) float64 {
	if expected.Value.Cmp(big.NewInt(0)) == 0 {
		if actual.Value.Cmp(big.NewInt(0)) == 0 {
			return 0.0
		}
		return 10000.0 // Max error if expected is 0 but actual is not
	}

	// Calculate relative error: |actual - expected| / expected
	diff := new(big.Int).Sub(actual.Value, expected.Value)
	if diff.Sign() < 0 {
		diff.Neg(diff)
	}

	// Convert to float for percentage calculation
	expectedFloat, _ := new(big.Float).SetInt(expected.Value).Float64()
	diffFloat, _ := new(big.Float).SetInt(diff).Float64()

	if expectedFloat == 0 {
		return 0.0
	}

	errorPercent := (diffFloat / expectedFloat) * 100
	errorBP := errorPercent * 100 // Convert to basis points

	// Cap at 10000 BP (100%)
	if errorBP > 10000 {
		errorBP = 10000
	}

	return errorBP
}