use std::collections::HashMap;
use serde_json::json;

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

/// Fixed Mean Reversion Strategy that properly tracks positions
/// 
/// This version doesn't maintain its own position tracking but relies
/// on the position information passed through on_fill callbacks
pub struct MeanReversionFixedStrategy {
    name: String,
    id: String,
    lookback_period: usize,
    entry_threshold: f64,  // Number of standard deviations
    exit_threshold: f64,   // Exit when price moves back this fraction toward mean
    position_size: f64,
    
    // State
    price_history: HashMap<String, Vec<f64>>,
    current_positions: HashMap<String, f64>, // Track actual positions from fills
    entry_prices: HashMap<String, f64>, // Track entry prices for exit decisions
}

impl MeanReversionFixedStrategy {
    pub fn new(
        name: String, 
        id: String,
        lookback_period: usize,
        entry_threshold: f64,
        position_size: f64,
    ) -> Self {
        Self {
            name,
            id,
            lookback_period,
            entry_threshold,
            exit_threshold: 0.3, // Exit when price moves 30% back to mean
            position_size,
            price_history: HashMap::new(),
            current_positions: HashMap::new(),
            entry_prices: HashMap::new(),
        }
    }
    
    fn calculate_mean(prices: &[f64]) -> f64 {
        prices.iter().sum::<f64>() / prices.len() as f64
    }
    
    fn calculate_std_dev(prices: &[f64], mean: f64) -> f64 {
        let variance = prices.iter()
            .map(|p| (p - mean).powi(2))
            .sum::<f64>() / prices.len() as f64;
        variance.sqrt()
    }
}

impl Strategy for MeanReversionFixedStrategy {
    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 price history
            let history = self.price_history.entry(symbol.clone()).or_insert_with(Vec::new);
            history.push(price);
            
            // Keep only necessary history
            if history.len() > self.lookback_period {
                history.remove(0);
            }
            
            // Need enough data
            if history.len() >= self.lookback_period {
                // Calculate statistics
                let mean = Self::calculate_mean(history);
                let std_dev = Self::calculate_std_dev(history, mean);
                
                // Calculate bands
                let upper_band = mean + self.entry_threshold * std_dev;
                let lower_band = mean - self.entry_threshold * std_dev;
                
                // Get actual position from our tracking
                let current_position = self.current_positions.get(symbol).copied().unwrap_or(0.0);
                
                // Entry signals - allow pyramiding up to 3x base position
                let max_long_position = self.position_size * 3.0;
                let max_short_position = -self.position_size * 3.0;
                
                if price < lower_band && current_position < max_long_position {
                    // Price is oversold, buy (or add to long position)
                    let remaining_capacity = max_long_position - current_position;
                    let trade_size = self.position_size.min(remaining_capacity);
                    
                    if trade_size > 0.0 {
                        eprintln!("Mean reversion: {} oversold at ${:.2}, buying {} shares (current: {}, lower band: ${:.2}, mean: ${:.2})",
                            symbol, price, trade_size, current_position, lower_band, mean);
                        
                        signals.push(Signal {
                            symbol: symbol.clone(),
                            signal_type: SignalType::Buy,
                            strength: 1.0,
                            quantity: Some(trade_size),
                            reason: Some(format!(
                                "Mean reversion buy: price ${:.2} < lower band ${:.2} (mean: ${:.2}, std: ${:.2})",
                                price, lower_band, mean, std_dev
                            )),
                            metadata: Some(json!({
                                "mean": mean,
                                "std_dev": std_dev,
                                "upper_band": upper_band,
                                "lower_band": lower_band,
                                "price": price,
                            })),
                        });
                    }
                } else if price > upper_band && current_position > max_short_position {
                    // Price is overbought, sell short (or add to short position)
                    let remaining_capacity = current_position - max_short_position;
                    let trade_size = self.position_size.min(remaining_capacity);
                    
                    if trade_size > 0.0 {
                        eprintln!("Mean reversion: {} overbought at ${:.2}, selling {} shares short (current: {}, upper band: ${:.2}, mean: ${:.2})",
                            symbol, price, trade_size, current_position, upper_band, mean);
                        
                        signals.push(Signal {
                            symbol: symbol.clone(),
                            signal_type: SignalType::Sell,
                            strength: 1.0,
                            quantity: Some(trade_size),
                            reason: Some(format!(
                                "Mean reversion sell: price ${:.2} > upper band ${:.2} (mean: ${:.2}, std: ${:.2})",
                                price, upper_band, mean, std_dev
                            )),
                            metadata: Some(json!({
                                "mean": mean,
                                "std_dev": std_dev,
                                "upper_band": upper_band,
                                "lower_band": lower_band,
                                "price": price,
                            })),
                        });
                    }
                }
                
                // Exit signals based on current position
                if current_position > 0.0 {
                    // We're long - check exit conditions
                    let entry_price = self.entry_prices.get(symbol).copied().unwrap_or(price);
                    let target_price = entry_price + (mean - entry_price) * self.exit_threshold;
                    let stop_loss = lower_band - std_dev; // Stop loss below lower band
                    
                    // Exit if price reaches target or stop loss
                    if price >= target_price || price <= stop_loss {
                        let exit_reason = if price >= target_price { "target" } else { "stop_loss" };
                        eprintln!("Mean reversion: {} exit long at ${:.2} ({}), closing {} shares",
                            symbol, price, exit_reason, current_position);
                        
                        signals.push(Signal {
                            symbol: symbol.clone(),
                            signal_type: SignalType::Sell,
                            strength: 1.0,
                            quantity: Some(current_position),
                            reason: Some(format!(
                                "Mean reversion exit long: price ${:.2} reached target ${:.2} (entry: ${:.2})",
                                price, target_price, entry_price
                            )),
                            metadata: Some(json!({
                                "mean": mean,
                                "price": price,
                                "entry_price": entry_price,
                                "target_price": target_price,
                                "stop_loss": stop_loss,
                                "exit_type": exit_reason,
                            })),
                        });
                    }
                } else if current_position < 0.0 {
                    // We're short - check exit conditions
                    let entry_price = self.entry_prices.get(symbol).copied().unwrap_or(price);
                    let target_price = entry_price - (entry_price - mean) * self.exit_threshold;
                    let stop_loss = upper_band + std_dev; // Stop loss above upper band
                    
                    // Exit if price reaches target or stop loss
                    if price <= target_price || price >= stop_loss {
                        let exit_reason = if price <= target_price { "target" } else { "stop_loss" };
                        eprintln!("Mean reversion: {} exit short at ${:.2} ({}), covering {} shares",
                            symbol, price, exit_reason, current_position.abs());
                        
                        signals.push(Signal {
                            symbol: symbol.clone(),
                            signal_type: SignalType::Buy,
                            strength: 1.0,
                            quantity: Some(current_position.abs()),
                            reason: Some(format!(
                                "Mean reversion exit short: price ${:.2} {} (entry: ${:.2}, target: ${:.2}, stop: ${:.2})",
                                price, exit_reason, entry_price, target_price, stop_loss
                            )),
                            metadata: Some(json!({
                                "mean": mean,
                                "price": price,
                                "entry_price": entry_price,
                                "target_price": target_price,
                                "stop_loss": stop_loss,
                                "exit_type": exit_reason,
                            })),
                        });
                    }
                }
            }
        }
        
        signals
    }
    
    fn on_fill(&mut self, symbol: &str, quantity: f64, price: f64, side: &str) {
        // Update our position tracking based on actual fills
        let current = self.current_positions.get(symbol).copied().unwrap_or(0.0);
        let new_position = if side == "buy" {
            current + quantity
        } else {
            current - quantity
        };
        
        eprintln!("Mean reversion fill: {} {} @ {} - {}, position: {} -> {}", 
            quantity, symbol, price, side, current, new_position);
        
        // Update position
        if new_position.abs() < 0.001 {
            // Position closed
            self.current_positions.remove(symbol);
            self.entry_prices.remove(symbol);
            eprintln!("Position closed for {}", symbol);
        } else {
            self.current_positions.insert(symbol.to_string(), new_position);
            
            // Track average entry price
            if current.abs() < 0.001 {
                // New position - set initial entry price
                self.entry_prices.insert(symbol.to_string(), price);
            } else if (current > 0.0 && new_position > current) || (current < 0.0 && new_position < current) {
                // Adding to existing position - update average entry price
                let old_price = self.entry_prices.get(symbol).copied().unwrap_or(price);
                let avg_price = (old_price * current.abs() + price * quantity) / new_position.abs();
                self.entry_prices.insert(symbol.to_string(), avg_price);
                eprintln!("Updated avg entry price for {}: ${:.2}", symbol, avg_price);
            }
        }
    }
    
    fn get_name(&self) -> &str {
        &self.name
    }
    
    fn get_parameters(&self) -> serde_json::Value {
        json!({
            "lookback_period": self.lookback_period,
            "entry_threshold": self.entry_threshold,
            "exit_threshold": self.exit_threshold,
            "position_size": self.position_size,
        })
    }
}