use std::collections::HashMap;
use chrono::{DateTime, Utc};
use serde_json::json;

use crate::{
    MarketUpdate, MarketDataType, 
    backtest::{Strategy, Signal, SignalType},
};

/// Pairs Trading Strategy
/// 
/// This strategy trades the spread between two correlated securities. When the spread
/// deviates from its historical mean, we trade expecting it to revert.
/// 
/// Entry Signals:
/// - Long pair A, Short pair B when spread < (mean - threshold * std)
/// - Short pair A, Long pair B when spread > (mean + threshold * std)
/// 
/// Exit Signals:
/// - Exit when spread returns to mean
pub struct PairsTradingStrategy {
    name: String,
    id: String,
    pair_a: String,
    pair_b: String,
    lookback_period: usize,
    entry_threshold: f64,  // Number of standard deviations
    position_size: f64,
    hedge_ratio: f64,      // How many shares of B per share of A
    
    // State
    price_history_a: Vec<f64>,
    price_history_b: Vec<f64>,
    spread_history: Vec<f64>,
    positions: HashMap<String, f64>,
    last_prices: HashMap<String, f64>,
}

impl PairsTradingStrategy {
    pub fn new(
        name: String, 
        id: String,
        pair_a: String,
        pair_b: String,
        lookback_period: usize,
        entry_threshold: f64,
        position_size: f64,
    ) -> Self {
        Self {
            name,
            id,
            pair_a,
            pair_b,
            lookback_period,
            entry_threshold,
            position_size,
            hedge_ratio: 1.0, // Default 1:1, could be calculated dynamically
            price_history_a: Vec::new(),
            price_history_b: Vec::new(),
            spread_history: Vec::new(),
            positions: HashMap::new(),
            last_prices: HashMap::new(),
        }
    }
    
    fn calculate_spread(&self, price_a: f64, price_b: f64) -> f64 {
        price_a - self.hedge_ratio * price_b
    }
    
    fn calculate_mean(values: &[f64]) -> f64 {
        values.iter().sum::<f64>() / values.len() as f64
    }
    
    fn calculate_std_dev(values: &[f64], mean: f64) -> f64 {
        let variance = values.iter()
            .map(|v| (v - mean).powi(2))
            .sum::<f64>() / values.len() as f64;
        variance.sqrt()
    }
}

impl Strategy for PairsTradingStrategy {
    fn on_market_data(&mut self, data: &MarketUpdate) -> Vec<Signal> {
        let mut signals = Vec::new();
        
        // Only process bar data
        if let MarketDataType::Bar(bar) = &data.data {
            let symbol = &data.symbol;
            let price = bar.close;
            
            // Update last prices
            self.last_prices.insert(symbol.clone(), price);
            
            // Update price histories
            if symbol == &self.pair_a {
                self.price_history_a.push(price);
                if self.price_history_a.len() > self.lookback_period {
                    self.price_history_a.remove(0);
                }
            } else if symbol == &self.pair_b {
                self.price_history_b.push(price);
                if self.price_history_b.len() > self.lookback_period {
                    self.price_history_b.remove(0);
                }
            }
            
            // Only generate signals when we have prices for both pairs
            if let (Some(&price_a), Some(&price_b)) = 
                (self.last_prices.get(&self.pair_a), self.last_prices.get(&self.pair_b)) {
                
                // Calculate current spread
                let spread = self.calculate_spread(price_a, price_b);
                
                // Update spread history
                self.spread_history.push(spread);
                if self.spread_history.len() > self.lookback_period {
                    self.spread_history.remove(0);
                }
                
                // Need enough data
                if self.spread_history.len() >= self.lookback_period {
                    // Calculate statistics
                    let mean = Self::calculate_mean(&self.spread_history);
                    let std_dev = Self::calculate_std_dev(&self.spread_history, mean);
                    
                    // Calculate bands
                    let upper_band = mean + self.entry_threshold * std_dev;
                    let lower_band = mean - self.entry_threshold * std_dev;
                    
                    let position_a = self.positions.get(&self.pair_a).copied().unwrap_or(0.0);
                    let position_b = self.positions.get(&self.pair_b).copied().unwrap_or(0.0);
                    
                    // Check for entry signals
                    if position_a == 0.0 && position_b == 0.0 {
                        if spread < lower_band {
                            // Spread too low: Buy A, Sell B
                            signals.push(Signal {
                                symbol: self.pair_a.clone(),
                                signal_type: SignalType::Buy,
                                strength: 1.0,
                                quantity: Some(self.position_size),
                                reason: Some(format!(
                                    "Pairs trade: spread ${:.2} < lower band ${:.2}",
                                    spread, lower_band
                                )),
                                metadata: Some(json!({
                                    "spread": spread,
                                    "mean": mean,
                                    "std_dev": std_dev,
                                    "pair": "A",
                                })),
                            });
                            
                            signals.push(Signal {
                                symbol: self.pair_b.clone(),
                                signal_type: SignalType::Sell,
                                strength: 1.0,
                                quantity: Some(self.position_size * self.hedge_ratio),
                                reason: Some(format!(
                                    "Pairs trade hedge: spread ${:.2} < lower band ${:.2}",
                                    spread, lower_band
                                )),
                                metadata: Some(json!({
                                    "spread": spread,
                                    "mean": mean,
                                    "std_dev": std_dev,
                                    "pair": "B",
                                })),
                            });
                            
                            self.positions.insert(self.pair_a.clone(), self.position_size);
                            self.positions.insert(self.pair_b.clone(), -self.position_size * self.hedge_ratio);
                        } else if spread > upper_band {
                            // Spread too high: Sell A, Buy B
                            signals.push(Signal {
                                symbol: self.pair_a.clone(),
                                signal_type: SignalType::Sell,
                                strength: 1.0,
                                quantity: Some(self.position_size),
                                reason: Some(format!(
                                    "Pairs trade: spread ${:.2} > upper band ${:.2}",
                                    spread, upper_band
                                )),
                                metadata: Some(json!({
                                    "spread": spread,
                                    "mean": mean,
                                    "std_dev": std_dev,
                                    "pair": "A",
                                })),
                            });
                            
                            signals.push(Signal {
                                symbol: self.pair_b.clone(),
                                signal_type: SignalType::Buy,
                                strength: 1.0,
                                quantity: Some(self.position_size * self.hedge_ratio),
                                reason: Some(format!(
                                    "Pairs trade hedge: spread ${:.2} > upper band ${:.2}",
                                    spread, upper_band
                                )),
                                metadata: Some(json!({
                                    "spread": spread,
                                    "mean": mean,
                                    "std_dev": std_dev,
                                    "pair": "B",
                                })),
                            });
                            
                            self.positions.insert(self.pair_a.clone(), -self.position_size);
                            self.positions.insert(self.pair_b.clone(), self.position_size * self.hedge_ratio);
                        }
                    }
                    // Check for exit signals
                    else if position_a != 0.0 && position_b != 0.0 {
                        // Exit when spread returns to mean
                        let spread_distance = (spread - mean).abs();
                        let exit_threshold = std_dev * 0.1; // Exit near mean
                        
                        if spread_distance < exit_threshold {
                            // Close positions
                            if position_a > 0.0 {
                                // We're long A, short B
                                signals.push(Signal {
                                    symbol: self.pair_a.clone(),
                                    signal_type: SignalType::Sell,
                                    strength: 1.0,
                                    quantity: Some(position_a),
                                    reason: Some(format!(
                                        "Pairs trade exit: spread ${:.2} returned to mean ${:.2}",
                                        spread, mean
                                    )),
                                    metadata: None,
                                });
                                
                                signals.push(Signal {
                                    symbol: self.pair_b.clone(),
                                    signal_type: SignalType::Buy,
                                    strength: 1.0,
                                    quantity: Some(position_b.abs()),
                                    reason: Some("Pairs trade exit: closing hedge".to_string()),
                                    metadata: None,
                                });
                            } else {
                                // We're short A, long B
                                signals.push(Signal {
                                    symbol: self.pair_a.clone(),
                                    signal_type: SignalType::Buy,
                                    strength: 1.0,
                                    quantity: Some(position_a.abs()),
                                    reason: Some(format!(
                                        "Pairs trade exit: spread ${:.2} returned to mean ${:.2}",
                                        spread, mean
                                    )),
                                    metadata: None,
                                });
                                
                                signals.push(Signal {
                                    symbol: self.pair_b.clone(),
                                    signal_type: SignalType::Sell,
                                    strength: 1.0,
                                    quantity: Some(position_b),
                                    reason: Some("Pairs trade exit: closing hedge".to_string()),
                                    metadata: None,
                                });
                            }
                            
                            self.positions.clear();
                        }
                    }
                }
            }
        }
        
        signals
    }
    
    fn on_fill(&mut self, symbol: &str, quantity: f64, price: f64, side: &str) {
        eprintln!("Pairs trading fill: {} {} @ {} - {}", quantity, symbol, price, side);
    }
    
    fn get_name(&self) -> &str {
        &self.name
    }
    
    fn get_parameters(&self) -> serde_json::Value {
        json!({
            "pair_a": self.pair_a,
            "pair_b": self.pair_b,
            "lookback_period": self.lookback_period,
            "entry_threshold": self.entry_threshold,
            "position_size": self.position_size,
            "hedge_ratio": self.hedge_ratio,
        })
    }
}