stock-bot/libs/vector-engine/src/index.ts

import { DataFrame } from '@stock-bot/data-frame';
import { getLogger } from '@stock-bot/logger';
import { atr, bollingerBands, ema, macd, rsi, sma } from '@stock-bot/utils';

// Vector operations interface
export interface VectorOperation {
  name: string;
  inputs: string[];
  output: string;
  operation: (inputs: number[][]) => number[];
}

// Vectorized strategy context
export interface VectorizedContext {
  data: DataFrame;
  lookback: number;
  indicators: Record<string, number[]>;
  signals: Record<string, number[]>;
}

// Performance metrics for vectorized backtesting
export interface VectorizedMetrics {
  totalReturns: number;
  sharpeRatio: number;
  maxDrawdown: number;
  winRate: number;
  profitFactor: number;
  totalTrades: number;
  avgTrade: number;
  returns: number[];
  drawdown: number[];
  equity: number[];
}

// Vectorized backtest result
export interface VectorizedBacktestResult {
  metrics: VectorizedMetrics;
  trades: VectorizedTrade[];
  equity: number[];
  timestamps: number[];
  signals: Record<string, number[]>;
}

export interface VectorizedTrade {
  entryIndex: number;
  exitIndex: number;
  entryPrice: number;
  exitPrice: number;
  quantity: number;
  side: 'LONG' | 'SHORT';
  pnl: number;
  return: number;
  duration: number;
}

// Vectorized strategy engine
export class VectorEngine {
  private logger = getLogger('vector-engine');
  private operations: Map<string, VectorOperation> = new Map();

  constructor() {
    this.registerDefaultOperations();
  }

  private registerDefaultOperations(): void {
    // Register common mathematical operations
    this.registerOperation({
      name: 'add',
      inputs: ['a', 'b'],
      output: 'result',
      operation: ([a, b]) => a.map((val, i) => val + b[i]),
    });

    this.registerOperation({
      name: 'subtract',
      inputs: ['a', 'b'],
      output: 'result',
      operation: ([a, b]) => a.map((val, i) => val - b[i]),
    });

    this.registerOperation({
      name: 'multiply',
      inputs: ['a', 'b'],
      output: 'result',
      operation: ([a, b]) => a.map((val, i) => val * b[i]),
    });

    this.registerOperation({
      name: 'divide',
      inputs: ['a', 'b'],
      output: 'result',
      operation: ([a, b]) => a.map((val, i) => (b[i] !== 0 ? val / b[i] : NaN)),
    });

    // Register comparison operations
    this.registerOperation({
      name: 'greater_than',
      inputs: ['a', 'b'],
      output: 'result',
      operation: ([a, b]) => a.map((val, i) => (val > b[i] ? 1 : 0)),
    });

    this.registerOperation({
      name: 'less_than',
      inputs: ['a', 'b'],
      output: 'result',
      operation: ([a, b]) => a.map((val, i) => (val < b[i] ? 1 : 0)),
    });

    this.registerOperation({
      name: 'crossover',
      inputs: ['a', 'b'],
      output: 'result',
      operation: ([a, b]) => {
        const result = new Array(a.length).fill(0);
        for (let i = 1; i < a.length; i++) {
          if (a[i] > b[i] && a[i - 1] <= b[i - 1]) {
            result[i] = 1;
          }
        }
        return result;
      },
    });

    this.registerOperation({
      name: 'crossunder',
      inputs: ['a', 'b'],
      output: 'result',
      operation: ([a, b]) => {
        const result = new Array(a.length).fill(0);
        for (let i = 1; i < a.length; i++) {
          if (a[i] < b[i] && a[i - 1] >= b[i - 1]) {
            result[i] = 1;
          }
        }
        return result;
      },
    });
  }

  registerOperation(operation: VectorOperation): void {
    this.operations.set(operation.name, operation);
    this.logger.debug(`Registered operation: ${operation.name}`);
  }

  // Execute vectorized strategy
  async executeVectorizedStrategy(
    data: DataFrame,
    strategyCode: string
  ): Promise<VectorizedBacktestResult> {
    try {
      const context = this.prepareContext(data);
      const signals = this.executeStrategy(context, strategyCode);
      const trades = this.generateTrades(data, signals);
      const metrics = this.calculateMetrics(data, trades);

      return {
        metrics,
        trades,
        equity: metrics.equity,
        timestamps: data.getColumn('timestamp'),
        signals,
      };
    } catch (error) {
      this.logger.error('Vectorized strategy execution failed', error);
      throw error;
    }
  }

  private prepareContext(data: DataFrame): VectorizedContext {
    const close = data.getColumn('close');
    const high = data.getColumn('high');
    const low = data.getColumn('low');
    const volume = data.getColumn('volume');

    // Calculate common indicators
    const indicators: Record<string, number[]> = {
      sma_20: sma(close, 20),
      sma_50: sma(close, 50),
      ema_12: ema(close, 12),
      ema_26: ema(close, 26),
      rsi: rsi(close),
    };

    const m = macd(close);
    indicators.macd = m.macd;
    indicators.macd_signal = m.signal;
    indicators.macd_histogram = m.histogram;

    const bb = bollingerBands(close);
    indicators.bb_upper = bb.upper;
    indicators.bb_middle = bb.middle;
    indicators.bb_lower = bb.lower;

    return {
      data,
      lookback: 100,
      indicators,
      signals: {},
    };
  }

  private executeStrategy(
    context: VectorizedContext,
    strategyCode: string
  ): Record<string, number[]> {
    // This is a simplified strategy execution
    // In production, you'd want a more sophisticated strategy compiler/interpreter
    const signals: Record<string, number[]> = {
      buy: new Array(context.data.length).fill(0),
      sell: new Array(context.data.length).fill(0),
    };

    // Example: Simple moving average crossover strategy
    if (strategyCode.includes('sma_crossover')) {
      const sma20 = context.indicators.sma_20;
      const sma50 = context.indicators.sma_50;

      for (let i = 1; i < sma20.length; i++) {
        // Buy signal: SMA20 crosses above SMA50
        if (!isNaN(sma20[i]) && !isNaN(sma50[i]) && !isNaN(sma20[i - 1]) && !isNaN(sma50[i - 1])) {
          if (sma20[i] > sma50[i] && sma20[i - 1] <= sma50[i - 1]) {
            signals.buy[i] = 1;
          }
          // Sell signal: SMA20 crosses below SMA50
          else if (sma20[i] < sma50[i] && sma20[i - 1] >= sma50[i - 1]) {
            signals.sell[i] = 1;
          }
        }
      }
    }

    return signals;
  }

  private generateTrades(data: DataFrame, signals: Record<string, number[]>): VectorizedTrade[] {
    const trades: VectorizedTrade[] = [];
    const close = data.getColumn('close');
    const timestamps = data.getColumn('timestamp');

    let position: { index: number; price: number; side: 'LONG' | 'SHORT' } | null = null;

    for (let i = 0; i < close.length; i++) {
      if (signals.buy[i] === 1 && !position) {
        // Open long position
        position = {
          index: i,
          price: close[i],
          side: 'LONG',
        };
      } else if (signals.sell[i] === 1) {
        if (position && position.side === 'LONG') {
          // Close long position
          const trade: VectorizedTrade = {
            entryIndex: position.index,
            exitIndex: i,
            entryPrice: position.price,
            exitPrice: close[i],
            quantity: 1, // Simplified: always trade 1 unit
            side: 'LONG',
            pnl: close[i] - position.price,
            return: (close[i] - position.price) / position.price,
            duration: timestamps[i] - timestamps[position.index],
          };
          trades.push(trade);
          position = null;
        } else if (!position) {
          // Open short position
          position = {
            index: i,
            price: close[i],
            side: 'SHORT',
          };
        }
      } else if (signals.buy[i] === 1 && position && position.side === 'SHORT') {
        // Close short position
        const trade: VectorizedTrade = {
          entryIndex: position.index,
          exitIndex: i,
          entryPrice: position.price,
          exitPrice: close[i],
          quantity: 1,
          side: 'SHORT',
          pnl: position.price - close[i],
          return: (position.price - close[i]) / position.price,
          duration: timestamps[i] - timestamps[position.index],
        };
        trades.push(trade);
        position = null;
      }
    }

    return trades;
  }

  private calculateMetrics(data: DataFrame, trades: VectorizedTrade[]): VectorizedMetrics {
    if (trades.length === 0) {
      return {
        totalReturns: 0,
        sharpeRatio: 0,
        maxDrawdown: 0,
        winRate: 0,
        profitFactor: 0,
        totalTrades: 0,
        avgTrade: 0,
        returns: [],
        drawdown: [],
        equity: [],
      };
    }

    const returns = trades.map(t => t.return);
    const pnls = trades.map(t => t.pnl);

    // Calculate equity curve
    const equity: number[] = [10000]; // Starting capital
    let currentEquity = 10000;

    for (const trade of trades) {
      currentEquity += trade.pnl;
      equity.push(currentEquity);
    }

    // Calculate drawdown
    const drawdown: number[] = [];
    let peak = equity[0];

    for (const eq of equity) {
      if (eq > peak) {peak = eq;}
      drawdown.push((peak - eq) / peak);
    }

    const totalReturns = (equity[equity.length - 1] - equity[0]) / equity[0];
    const avgReturn = returns.reduce((sum, r) => sum + r, 0) / returns.length;
    const returnStd = Math.sqrt(
      returns.reduce((sum, r) => sum + Math.pow(r - avgReturn, 2), 0) / returns.length
    );

    const winningTrades = trades.filter(t => t.pnl > 0);
    const losingTrades = trades.filter(t => t.pnl < 0);

    const grossProfit = winningTrades.reduce((sum, t) => sum + t.pnl, 0);
    const grossLoss = Math.abs(losingTrades.reduce((sum, t) => sum + t.pnl, 0));

    return {
      totalReturns,
      sharpeRatio: returnStd !== 0 ? (avgReturn / returnStd) * Math.sqrt(252) : 0,
      maxDrawdown: Math.max(...drawdown),
      winRate: winningTrades.length / trades.length,
      profitFactor: grossLoss !== 0 ? grossProfit / grossLoss : Infinity,
      totalTrades: trades.length,
      avgTrade: pnls.reduce((sum, pnl) => sum + pnl, 0) / trades.length,
      returns,
      drawdown,
      equity,
    };
  }

  // Utility methods for vectorized operations
  applyOperation(operationName: string, inputs: Record<string, number[]>): number[] {
    const operation = this.operations.get(operationName);
    if (!operation) {
      throw new Error(`Operation '${operationName}' not found`);
    }

    const inputArrays = operation.inputs.map(inputName => {
      if (!inputs[inputName]) {
        throw new Error(`Input '${inputName}' not provided for operation '${operationName}'`);
      }
      return inputs[inputName];
    });

    return operation.operation(inputArrays);
  }

  // Batch processing for multiple strategies
  async batchBacktest(
    data: DataFrame,
    strategies: Array<{ id: string; code: string }>
  ): Promise<Record<string, VectorizedBacktestResult>> {
    const results: Record<string, VectorizedBacktestResult> = {};

    for (const strategy of strategies) {
      try {
        this.logger.info(`Running vectorized backtest for strategy: ${strategy.id}`);
        results[strategy.id] = await this.executeVectorizedStrategy(data, strategy.code);
      } catch (error) {
        this.logger.error(`Backtest failed for strategy: ${strategy.id}`, error);
        // Continue with other strategies
      }
    }

    return results;
  }
}