import { BaseStrategy, Signal } from '../BaseStrategy';
import { MarketData } from '../../types';
import { getLogger } from '@stock-bot/logger';

const logger = getLogger('MLEnhancedStrategy');
import * as tf from '@tensorflow/tfjs-node';

interface MLModelConfig {
  modelPath?: string;
  features: string[];
  lookbackPeriod: number;
  updateFrequency: number; // How often to retrain in minutes
  minTrainingSize: number;
}

export class MLEnhancedStrategy extends BaseStrategy {
  private model: tf.LayersModel | null = null;
  private featureBuffer: Map<string, number[][]> = new Map();
  private predictions: Map<string, number> = new Map();
  private lastUpdate: number = 0;
  private trainingData: { features: number[][]; labels: number[] } = { features: [], labels: [] };
  
  // Feature extractors
  private indicators: Map<string, any> = new Map();
  
  // ML Configuration
  private mlConfig: MLModelConfig = {
    features: [
      'returns_20', 'returns_50', 'volatility_20', 'rsi_14',
      'volume_ratio', 'price_position', 'macd_signal'
    ],
    lookbackPeriod: 50,
    updateFrequency: 1440, // Daily
    minTrainingSize: 1000
  };

  protected async onStart(): Promise<void> {
    logger.info('ML Enhanced Strategy starting...');
    
    // Try to load existing model
    if (this.mlConfig.modelPath) {
      try {
        this.model = await tf.loadLayersModel(`file://${this.mlConfig.modelPath}`);
        logger.info('Loaded existing ML model');
      } catch (error) {
        logger.warn('Could not load model, will train new one');
      }
    }
    
    // Initialize feature buffers for each symbol
    this.config.symbols.forEach(symbol => {
      this.featureBuffer.set(symbol, []);
      this.indicators.set(symbol, {
        prices: [],
        volumes: [],
        returns: [],
        sma20: 0,
        sma50: 0,
        rsi: 50,
        macd: 0,
        signal: 0
      });
    });
  }

  protected updateIndicators(data: MarketData): void {
    if (data.type !== 'bar') return;
    
    const symbol = data.data.symbol;
    const indicators = this.indicators.get(symbol);
    if (!indicators) return;
    
    // Update price and volume history
    indicators.prices.push(data.data.close);
    indicators.volumes.push(data.data.volume);
    
    if (indicators.prices.length > 200) {
      indicators.prices.shift();
      indicators.volumes.shift();
    }
    
    // Calculate returns
    if (indicators.prices.length >= 2) {
      const ret = (data.data.close - indicators.prices[indicators.prices.length - 2]) / 
                  indicators.prices[indicators.prices.length - 2];
      indicators.returns.push(ret);
      
      if (indicators.returns.length > 50) {
        indicators.returns.shift();
      }
    }
    
    // Update technical indicators
    if (indicators.prices.length >= 20) {
      indicators.sma20 = this.calculateSMA(indicators.prices, 20);
      indicators.volatility20 = this.calculateVolatility(indicators.returns, 20);
    }
    
    if (indicators.prices.length >= 50) {
      indicators.sma50 = this.calculateSMA(indicators.prices, 50);
    }
    
    if (indicators.prices.length >= 14) {
      indicators.rsi = this.calculateRSI(indicators.prices, 14);
    }
    
    // Extract features
    const features = this.extractFeatures(symbol, data);
    if (features) {
      const buffer = this.featureBuffer.get(symbol)!;
      buffer.push(features);
      
      if (buffer.length > this.mlConfig.lookbackPeriod) {
        buffer.shift();
      }
      
      // Make prediction if we have enough data
      if (buffer.length === this.mlConfig.lookbackPeriod && this.model) {
        this.makePrediction(symbol, buffer);
      }
    }
    
    // Check if we should update the model
    const now = Date.now();
    if (now - this.lastUpdate > this.mlConfig.updateFrequency * 60 * 1000) {
      this.updateModel();
      this.lastUpdate = now;
    }
  }

  protected async generateSignal(data: MarketData): Promise<Signal | null> {
    if (data.type !== 'bar') return null;
    
    const symbol = data.data.symbol;
    const prediction = this.predictions.get(symbol);
    
    if (!prediction || Math.abs(prediction) < 0.01) {
      return null; // No strong signal
    }
    
    const position = this.getPosition(symbol);
    const indicators = this.indicators.get(symbol);
    
    // Risk management checks
    const volatility = indicators?.volatility20 || 0.02;
    const maxPositionRisk = 0.02; // 2% max risk per position
    const positionSize = this.calculatePositionSize(volatility, maxPositionRisk);
    
    // Generate signals based on ML predictions
    if (prediction > 0.02 && position <= 0) {
      // Strong bullish prediction
      return {
        type: 'buy',
        symbol,
        strength: Math.min(prediction * 50, 1), // Scale prediction to 0-1
        reason: `ML prediction: ${(prediction * 100).toFixed(2)}% expected return`,
        metadata: {
          prediction,
          confidence: this.calculateConfidence(symbol),
          features: this.getLatestFeatures(symbol)
        }
      };
    } else if (prediction < -0.02 && position >= 0) {
      // Strong bearish prediction
      return {
        type: 'sell',
        symbol,
        strength: Math.min(Math.abs(prediction) * 50, 1),
        reason: `ML prediction: ${(prediction * 100).toFixed(2)}% expected return`,
        metadata: {
          prediction,
          confidence: this.calculateConfidence(symbol),
          features: this.getLatestFeatures(symbol)
        }
      };
    } else if (position !== 0 && Math.sign(position) !== Math.sign(prediction)) {
      // Exit if prediction reverses
      return {
        type: 'close',
        symbol,
        strength: 1,
        reason: 'ML prediction reversed',
        metadata: { prediction }
      };
    }
    
    return null;
  }

  private extractFeatures(symbol: string, data: MarketData): number[] | null {
    const indicators = this.indicators.get(symbol);
    if (!indicators || indicators.prices.length < 50) return null;
    
    const features: number[] = [];
    
    // Price returns
    const currentPrice = indicators.prices[indicators.prices.length - 1];
    features.push((currentPrice / indicators.prices[indicators.prices.length - 20] - 1)); // 20-day return
    features.push((currentPrice / indicators.prices[indicators.prices.length - 50] - 1)); // 50-day return
    
    // Volatility
    features.push(indicators.volatility20 || 0);
    
    // RSI
    features.push((indicators.rsi - 50) / 50); // Normalize to -1 to 1
    
    // Volume ratio
    const avgVolume = indicators.volumes.slice(-20).reduce((a, b) => a + b, 0) / 20;
    features.push(data.data.volume / avgVolume - 1);
    
    // Price position in daily range
    const pricePosition = (data.data.close - data.data.low) / (data.data.high - data.data.low);
    features.push(pricePosition * 2 - 1); // Normalize to -1 to 1
    
    // MACD signal
    if (indicators.macd && indicators.signal) {
      features.push((indicators.macd - indicators.signal) / currentPrice);
    } else {
      features.push(0);
    }
    
    // Store for training
    if (indicators.returns.length >= 21) {
      const futureReturn = indicators.returns[indicators.returns.length - 1];
      this.trainingData.features.push([...features]);
      this.trainingData.labels.push(futureReturn);
      
      // Limit training data size
      if (this.trainingData.features.length > 10000) {
        this.trainingData.features.shift();
        this.trainingData.labels.shift();
      }
    }
    
    return features;
  }

  private async makePrediction(symbol: string, featureBuffer: number[][]): Promise<void> {
    if (!this.model) return;
    
    try {
      // Prepare input tensor
      const input = tf.tensor3d([featureBuffer]);
      
      // Make prediction
      const prediction = await this.model.predict(input) as tf.Tensor;
      const value = (await prediction.data())[0];
      
      this.predictions.set(symbol, value);
      
      // Cleanup tensors
      input.dispose();
      prediction.dispose();
    } catch (error) {
      logger.error('ML prediction error:', error);
    }
  }

  private async updateModel(): Promise<void> {
    if (this.trainingData.features.length < this.mlConfig.minTrainingSize) {
      logger.info('Not enough training data yet');
      return;
    }
    
    logger.info('Updating ML model...');
    
    try {
      // Create or update model
      if (!this.model) {
        this.model = this.createModel();
      }
      
      // Prepare training data
      const features = tf.tensor2d(this.trainingData.features);
      const labels = tf.tensor1d(this.trainingData.labels);
      
      // Train model
      await this.model.fit(features, labels, {
        epochs: 50,
        batchSize: 32,
        validationSplit: 0.2,
        shuffle: true,
        callbacks: {
          onEpochEnd: (epoch, logs) => {
            if (epoch % 10 === 0) {
              logger.debug(`Epoch ${epoch}: loss = ${logs?.loss.toFixed(4)}`);
            }
          }
        }
      });
      
      logger.info('Model updated successfully');
      
      // Cleanup tensors
      features.dispose();
      labels.dispose();
      
      // Save model if path provided
      if (this.mlConfig.modelPath) {
        await this.model.save(`file://${this.mlConfig.modelPath}`);
      }
    } catch (error) {
      logger.error('Model update error:', error);
    }
  }

  private createModel(): tf.LayersModel {
    const model = tf.sequential({
      layers: [
        // LSTM layer for sequence processing
        tf.layers.lstm({
          units: 64,
          returnSequences: true,
          inputShape: [this.mlConfig.lookbackPeriod, this.mlConfig.features.length]
        }),
        tf.layers.dropout({ rate: 0.2 }),
        
        // Second LSTM layer
        tf.layers.lstm({
          units: 32,
          returnSequences: false
        }),
        tf.layers.dropout({ rate: 0.2 }),
        
        // Dense layers
        tf.layers.dense({
          units: 16,
          activation: 'relu'
        }),
        tf.layers.dropout({ rate: 0.1 }),
        
        // Output layer
        tf.layers.dense({
          units: 1,
          activation: 'tanh' // Output between -1 and 1
        })
      ]
    });
    
    model.compile({
      optimizer: tf.train.adam(0.001),
      loss: 'meanSquaredError',
      metrics: ['mae']
    });
    
    return model;
  }

  private calculateConfidence(symbol: string): number {
    // Simple confidence based on prediction history accuracy
    // In practice, would track actual vs predicted returns
    const prediction = this.predictions.get(symbol) || 0;
    return Math.min(Math.abs(prediction) * 10, 1);
  }

  private getLatestFeatures(symbol: string): Record<string, number> {
    const buffer = this.featureBuffer.get(symbol);
    if (!buffer || buffer.length === 0) return {};
    
    const latest = buffer[buffer.length - 1];
    return {
      returns_20: latest[0],
      returns_50: latest[1],
      volatility_20: latest[2],
      rsi_normalized: latest[3],
      volume_ratio: latest[4],
      price_position: latest[5],
      macd_signal: latest[6]
    };
  }

  private calculateVolatility(returns: number[], period: number): number {
    if (returns.length < period) return 0;
    
    const recentReturns = returns.slice(-period);
    const mean = recentReturns.reduce((a, b) => a + b, 0) / period;
    const variance = recentReturns.reduce((sum, r) => sum + Math.pow(r - mean, 2), 0) / period;
    
    return Math.sqrt(variance * 252); // Annualized
  }

  private calculatePositionSize(volatility: number, maxRisk: number): number {
    // Kelly-inspired position sizing with volatility adjustment
    const targetVolatility = 0.15; // 15% annual target
    const volAdjustment = targetVolatility / volatility;
    
    return Math.min(volAdjustment, 2.0); // Max 2x leverage
  }

  protected onStop(): void {
    logger.info('ML Enhanced Strategy stopped');
    
    // Save final model if configured
    if (this.model && this.mlConfig.modelPath) {
      this.model.save(`file://${this.mlConfig.modelPath}`)
        .then(() => logger.info('Model saved'))
        .catch(err => logger.error('Failed to save model:', err));
    }
    
    // Cleanup
    this.featureBuffer.clear();
    this.predictions.clear();
    this.indicators.clear();
    if (this.model) {
      this.model.dispose();
    }
  }

  protected onConfigUpdate(updates: any): void {
    logger.info('ML Enhanced Strategy config updated:', updates);
    
    if (updates.mlConfig) {
      this.mlConfig = { ...this.mlConfig, ...updates.mlConfig };
    }
  }
}