use crate::{MarketMicrostructure, PriceLevel, Quote, Trade, Bar, Side};
use chrono::{DateTime, Utc, Duration, Timelike};
use rand::prelude::*;
use rand_distr::{Normal, Pareto, Beta};

pub struct OrderBookReconstructor {
    tick_size: f64,
    lot_size: f64,
    num_levels: usize,
    spread_model: SpreadModel,
    depth_model: DepthModel,
}

#[derive(Clone)]
pub enum SpreadModel {
    Fixed { spread_ticks: u32 },
    Dynamic { base_bps: f64, volatility_factor: f64 },
    InformedTrader { base_bps: f64, information_decay: f64 },
}

#[derive(Clone)]
pub enum DepthModel {
    Linear { base_size: f64, decay_rate: f64 },
    Exponential { base_size: f64, decay_factor: f64 },
    PowerLaw { alpha: f64, x_min: f64 },
}

impl OrderBookReconstructor {
    pub fn new(tick_size: f64, lot_size: f64) -> Self {
        Self {
            tick_size,
            lot_size,
            num_levels: 10,
            spread_model: SpreadModel::Dynamic { 
                base_bps: 2.0, 
                volatility_factor: 1.5 
            },
            depth_model: DepthModel::Exponential { 
                base_size: 1000.0, 
                decay_factor: 0.7 
            },
        }
    }
    
    pub fn reconstruct_from_trades_and_quotes(
        &self,
        trades: &[(DateTime<Utc>, Trade)],
        quotes: &[(DateTime<Utc>, Quote)],
        timestamp: DateTime<Utc>,
    ) -> (Vec<PriceLevel>, Vec<PriceLevel>) {
        // Find the most recent quote before timestamp
        let recent_quote = quotes.iter()
            .filter(|(t, _)| *t <= timestamp)
            .last()
            .map(|(_, q)| q);
        
        // Find recent trades to estimate market conditions
        let recent_trades: Vec<_> = trades.iter()
            .filter(|(t, _)| {
                let age = timestamp - *t;
                age < Duration::minutes(5) && age >= Duration::zero()
            })
            .map(|(_, t)| t)
            .collect();
        
        if let Some(quote) = recent_quote {
            // Start with actual quote
            self.build_full_book(quote, &recent_trades, timestamp)
        } else if !recent_trades.is_empty() {
            // Reconstruct from trades only
            self.reconstruct_from_trades_only(&recent_trades, timestamp)
        } else {
            // No data - return empty book
            (vec![], vec![])
        }
    }
    
    fn build_full_book(
        &self,
        top_quote: &Quote,
        recent_trades: &[&Trade],
        _timestamp: DateTime<Utc>,
    ) -> (Vec<PriceLevel>, Vec<PriceLevel>) {
        let mut bids = Vec::with_capacity(self.num_levels);
        let mut asks = Vec::with_capacity(self.num_levels);
        
        // Add top of book
        bids.push(PriceLevel {
            price: top_quote.bid,
            size: top_quote.bid_size,
            order_count: Some(self.estimate_order_count(top_quote.bid_size)),
        });
        
        asks.push(PriceLevel {
            price: top_quote.ask,
            size: top_quote.ask_size,
            order_count: Some(self.estimate_order_count(top_quote.ask_size)),
        });
        
        // Calculate spread and volatility from recent trades
        let (_spread_bps, _volatility) = self.estimate_market_conditions(recent_trades, top_quote);
        
        // Build deeper levels
        for i in 1..self.num_levels {
            // Bid levels
            let bid_price = top_quote.bid - (i as f64 * self.tick_size);
            let bid_size = self.calculate_level_size(i, top_quote.bid_size, &self.depth_model);
            bids.push(PriceLevel {
                price: bid_price,
                size: bid_size,
                order_count: Some(self.estimate_order_count(bid_size)),
            });
            
            // Ask levels
            let ask_price = top_quote.ask + (i as f64 * self.tick_size);
            let ask_size = self.calculate_level_size(i, top_quote.ask_size, &self.depth_model);
            asks.push(PriceLevel {
                price: ask_price,
                size: ask_size,
                order_count: Some(self.estimate_order_count(ask_size)),
            });
        }
        
        (bids, asks)
    }
    
    fn reconstruct_from_trades_only(
        &self,
        recent_trades: &[&Trade],
        _timestamp: DateTime<Utc>,
    ) -> (Vec<PriceLevel>, Vec<PriceLevel>) {
        if recent_trades.is_empty() {
            return (vec![], vec![]);
        }
        
        // Estimate mid price from trades
        let prices: Vec<f64> = recent_trades.iter().map(|t| t.price).collect();
        let mid_price = prices.iter().sum::<f64>() / prices.len() as f64;
        
        // Estimate spread from trade price variance
        let variance = prices.iter()
            .map(|p| (p - mid_price).powi(2))
            .sum::<f64>() / prices.len() as f64;
        let estimated_spread = variance.sqrt() * 2.0; // Rough approximation
        
        // Build synthetic book
        let bid_price = (mid_price - estimated_spread / 2.0 / self.tick_size).round() * self.tick_size;
        let ask_price = (mid_price + estimated_spread / 2.0 / self.tick_size).round() * self.tick_size;
        
        // Estimate sizes from trade volumes
        let avg_trade_size = recent_trades.iter()
            .map(|t| t.size)
            .sum::<f64>() / recent_trades.len() as f64;
        
        let mut bids = Vec::with_capacity(self.num_levels);
        let mut asks = Vec::with_capacity(self.num_levels);
        
        for i in 0..self.num_levels {
            let level_size = avg_trade_size * 10.0 / (i + 1) as f64; // Decay with depth
            
            bids.push(PriceLevel {
                price: bid_price - (i as f64 * self.tick_size),
                size: level_size,
                order_count: Some(self.estimate_order_count(level_size)),
            });
            
            asks.push(PriceLevel {
                price: ask_price + (i as f64 * self.tick_size),
                size: level_size,
                order_count: Some(self.estimate_order_count(level_size)),
            });
        }
        
        (bids, asks)
    }
    
    fn calculate_level_size(&self, level: usize, _top_size: f64, model: &DepthModel) -> f64 {
        let size = match model {
            DepthModel::Linear { base_size, decay_rate } => {
                base_size - (level as f64 * decay_rate)
            }
            DepthModel::Exponential { base_size, decay_factor } => {
                base_size * decay_factor.powi(level as i32)
            }
            DepthModel::PowerLaw { alpha, x_min } => {
                x_min * ((level + 1) as f64).powf(-alpha)
            }
        };
        
        // Round to lot size and ensure positive
        ((size / self.lot_size).round() * self.lot_size).max(self.lot_size)
    }
    
    fn estimate_order_count(&self, size: f64) -> u32 {
        // Estimate based on typical order size distribution
        let avg_order_size = 100.0;
        let base_count = (size / avg_order_size).ceil() as u32;
        
        // Add some randomness
        let mut rng = thread_rng();
        let variation = rng.gen_range(0.8..1.2);
        ((base_count as f64 * variation) as u32).max(1)
    }
    
    fn estimate_market_conditions(
        &self,
        recent_trades: &[&Trade],
        quote: &Quote,
    ) -> (f64, f64) {
        if recent_trades.is_empty() {
            let spread_bps = ((quote.ask - quote.bid) / quote.bid) * 10000.0;
            return (spread_bps, 0.02); // Default 2% volatility
        }
        
        // Calculate spread in bps
        let mid_price = (quote.bid + quote.ask) / 2.0;
        let spread_bps = ((quote.ask - quote.bid) / mid_price) * 10000.0;
        
        // Estimate volatility from trade prices
        let prices: Vec<f64> = recent_trades.iter().map(|t| t.price).collect();
        let returns: Vec<f64> = prices.windows(2)
            .map(|w| (w[1] / w[0]).ln())
            .collect();
        
        let volatility = if !returns.is_empty() {
            let mean_return = returns.iter().sum::<f64>() / returns.len() as f64;
            let variance = returns.iter()
                .map(|r| (r - mean_return).powi(2))
                .sum::<f64>() / returns.len() as f64;
            variance.sqrt() * (252.0_f64).sqrt() // Annualize
        } else {
            0.02 // Default 2%
        };
        
        (spread_bps, volatility)
    }
}

// Market data synthesizer for generating realistic data
pub struct MarketDataSynthesizer {
    base_price: f64,
    tick_size: f64,
    base_spread_bps: f64,
    volatility: f64,
    mean_reversion_speed: f64,
    jump_intensity: f64,
    jump_size_dist: Normal<f64>,
    volume_dist: Pareto<f64>,
    intraday_pattern: Vec<f64>,
}

impl MarketDataSynthesizer {
    pub fn new(symbol_params: &MarketMicrostructure) -> Self {
        let jump_size_dist = Normal::new(0.0, symbol_params.volatility * 0.1).unwrap();
        let volume_dist = Pareto::new(1.0, 1.5).unwrap();
        
        Self {
            base_price: 100.0, // Will be updated with actual price
            tick_size: symbol_params.tick_size,
            base_spread_bps: symbol_params.avg_spread_bps,
            volatility: symbol_params.volatility,
            mean_reversion_speed: 0.1,
            jump_intensity: 0.05, // 5% chance of jump per time step
            jump_size_dist,
            volume_dist,
            intraday_pattern: symbol_params.intraday_volume_profile.clone(),
        }
    }
    
    pub fn generate_quote_sequence(
        &mut self,
        start_price: f64,
        start_time: DateTime<Utc>,
        end_time: DateTime<Utc>,
        interval_ms: i64,
    ) -> Vec<(DateTime<Utc>, Quote)> {
        self.base_price = start_price;
        let mut quotes = Vec::new();
        let mut current_time = start_time;
        let mut mid_price = start_price;
        let mut spread_factor;
        let mut rng = thread_rng();
        
        while current_time <= end_time {
            // Generate price movement
            let dt = interval_ms as f64 / 1000.0 / 86400.0; // Convert to days
            
            // Ornstein-Uhlenbeck process with jumps
            let drift = -self.mean_reversion_speed * (mid_price / self.base_price - 1.0).ln();
            let diffusion = self.volatility * (dt.sqrt()) * rng.gen::<f64>();
            
            // Add jump component
            let jump = if rng.gen::<f64>() < self.jump_intensity * dt {
                mid_price * self.jump_size_dist.sample(&mut rng)
            } else {
                0.0
            };
            
            mid_price *= 1.0 + drift * dt + diffusion + jump;
            mid_price = (mid_price / self.tick_size).round() * self.tick_size;
            
            // Dynamic spread based on volatility and time of day
            let hour_index = current_time.hour() as usize;
            let volume_factor = if hour_index < self.intraday_pattern.len() {
                self.intraday_pattern[hour_index]
            } else {
                0.04 // Default 4% of daily volume per hour
            };
            
            // Wider spreads during low volume periods
            spread_factor = 1.0 / volume_factor.sqrt();
            let spread_bps = self.base_spread_bps * spread_factor;
            let half_spread = mid_price * spread_bps / 20000.0;
            
            // Generate bid/ask
            let bid = ((mid_price - half_spread) / self.tick_size).floor() * self.tick_size;
            let ask = ((mid_price + half_spread) / self.tick_size).ceil() * self.tick_size;
            
            // Generate sizes with correlation to spread
            let size_multiplier = 1.0 / spread_factor; // Tighter spread = more size
            let bid_size = (self.volume_dist.sample(&mut rng) * 1000.0 * size_multiplier).round();
            let ask_size = (self.volume_dist.sample(&mut rng) * 1000.0 * size_multiplier).round();
            
            quotes.push((current_time, Quote {
                bid,
                ask,
                bid_size,
                ask_size,
            }));
            
            current_time = current_time + Duration::milliseconds(interval_ms);
        }
        
        quotes
    }
    
    pub fn generate_trade_sequence(
        &mut self,
        quotes: &[(DateTime<Utc>, Quote)],
        trade_intensity: f64,
    ) -> Vec<(DateTime<Utc>, Trade)> {
        let mut trades = Vec::new();
        let mut rng = thread_rng();
        let beta_dist = Beta::new(2.0, 5.0).unwrap(); // Skewed towards smaller trades
        
        for (time, quote) in quotes {
            // Poisson process for trade arrivals
            let num_trades = rng.gen_range(0..((trade_intensity * 10.0) as u32));
            
            for i in 0..num_trades {
                // Determine trade side (slight bias based on spread)
                let spread_ratio = (quote.ask - quote.bid) / quote.bid;
                let buy_prob = 0.5 - spread_ratio * 10.0; // More sells when spread is wide
                let side = if rng.gen::<f64>() < buy_prob {
                    Side::Buy
                } else {
                    Side::Sell
                };
                
                // Trade price (sometimes inside spread for large trades)
                let price = match side {
                    Side::Buy => {
                        if rng.gen::<f64>() < 0.9 {
                            quote.ask // Take liquidity
                        } else {
                            // Provide liquidity (inside spread)
                            quote.bid + (quote.ask - quote.bid) * rng.gen::<f64>()
                        }
                    }
                    Side::Sell => {
                        if rng.gen::<f64>() < 0.9 {
                            quote.bid // Take liquidity
                        } else {
                            // Provide liquidity (inside spread)
                            quote.bid + (quote.ask - quote.bid) * rng.gen::<f64>()
                        }
                    }
                };
                
                // Trade size (power law distribution)
                let size_percentile = beta_dist.sample(&mut rng);
                let base_size = match side {
                    Side::Buy => quote.ask_size,
                    Side::Sell => quote.bid_size,
                };
                let size = (base_size * size_percentile * 0.1).round().max(1.0);
                
                // Add small time offset for multiple trades
                let trade_time = *time + Duration::milliseconds(i as i64 * 100);
                
                trades.push((trade_time, Trade {
                    price,
                    size,
                    side,
                }));
            }
        }
        
        trades.sort_by_key(|(t, _)| *t);
        trades
    }
    
    pub fn aggregate_to_bars(
        &self,
        trades: &[(DateTime<Utc>, Trade)],
        bar_duration: Duration,
    ) -> Vec<(DateTime<Utc>, Bar)> {
        if trades.is_empty() {
            return Vec::new();
        }
        
        let mut bars = Vec::new();
        let mut current_bar_start = trades[0].0;
        let mut current_bar_end = current_bar_start + bar_duration;
        
        let mut open = 0.0;
        let mut high = 0.0;
        let mut low = f64::MAX;
        let mut close = 0.0;
        let mut volume = 0.0;
        let mut vwap_numerator = 0.0;
        let mut first_trade = true;
        
        for (time, trade) in trades {
            // Check if we need to start a new bar
            while *time >= current_bar_end {
                if volume > 0.0 {
                    bars.push((current_bar_start, Bar {
                        open,
                        high,
                        low,
                        close,
                        volume,
                        vwap: Some(vwap_numerator / volume),
                    }));
                }
                
                // Reset for new bar
                current_bar_start = current_bar_end;
                current_bar_end = current_bar_start + bar_duration;
                open = 0.0;
                high = 0.0;
                low = f64::MAX;
                close = 0.0;
                volume = 0.0;
                vwap_numerator = 0.0;
                first_trade = true;
            }
            
            // Update current bar
            if first_trade {
                open = trade.price;
                first_trade = false;
            }
            high = high.max(trade.price);
            low = low.min(trade.price);
            close = trade.price;
            volume += trade.size;
            vwap_numerator += trade.price * trade.size;
        }
        
        // Add final bar if it has data
        if volume > 0.0 {
            bars.push((current_bar_start, Bar {
                open,
                high,
                low,
                close,
                volume,
                vwap: Some(vwap_numerator / volume),
            }));
        }
        
        bars
    }
}