use crate::{Side, MarketMicrostructure, PriceLevel};
use chrono::{DateTime, Utc, Timelike};

#[derive(Debug, Clone)]
pub struct MarketImpactEstimate {
    pub temporary_impact: f64,
    pub permanent_impact: f64,
    pub total_impact: f64,
    pub expected_cost: f64,
    pub impact_decay_ms: i64,
}

#[derive(Debug, Clone, Copy)]
pub enum ImpactModelType {
    Linear,
    SquareRoot,
    PowerLaw { exponent: f64 },
    AlmgrenChriss,
    IStarModel,
}

pub struct MarketImpactModel {
    model_type: ImpactModelType,
    // Model parameters
    temporary_impact_coef: f64,
    permanent_impact_coef: f64,
    spread_impact_weight: f64,
    volatility_adjustment: bool,
}

impl MarketImpactModel {
    pub fn new(model_type: ImpactModelType) -> Self {
        match model_type {
            ImpactModelType::Linear => Self {
                model_type,
                temporary_impact_coef: 0.1,
                permanent_impact_coef: 0.05,
                spread_impact_weight: 0.5,
                volatility_adjustment: true,
            },
            ImpactModelType::SquareRoot => Self {
                model_type,
                temporary_impact_coef: 0.142, // Empirical from literature
                permanent_impact_coef: 0.0625,
                spread_impact_weight: 0.5,
                volatility_adjustment: true,
            },
            ImpactModelType::AlmgrenChriss => Self {
                model_type,
                temporary_impact_coef: 0.314,
                permanent_impact_coef: 0.142,
                spread_impact_weight: 0.7,
                volatility_adjustment: true,
            },
            ImpactModelType::PowerLaw { .. } => Self {
                model_type,
                temporary_impact_coef: 0.2,
                permanent_impact_coef: 0.1,
                spread_impact_weight: 0.5,
                volatility_adjustment: true,
            },
            ImpactModelType::IStarModel => Self {
                model_type,
                temporary_impact_coef: 1.0,
                permanent_impact_coef: 0.5,
                spread_impact_weight: 0.8,
                volatility_adjustment: true,
            },
        }
    }

    pub fn estimate_impact(
        &self,
        order_size: f64,
        side: Side,
        microstructure: &MarketMicrostructure,
        orderbook: &[PriceLevel],
        current_time: DateTime<Utc>,
    ) -> MarketImpactEstimate {
        // Calculate participation rate
        let intraday_volume = self.get_expected_volume(microstructure, current_time);
        let participation_rate = order_size / intraday_volume.max(1.0);
        
        // Calculate spread in basis points
        let spread_bps = microstructure.avg_spread_bps;
        
        // Calculate volatility adjustment
        let vol_adjustment = if self.volatility_adjustment {
            (microstructure.volatility / 0.02).sqrt() // Normalize to 2% daily vol
        } else {
            1.0
        };
        
        // Calculate temporary impact based on model type
        let temp_impact_bps = match self.model_type {
            ImpactModelType::Linear => {
                self.temporary_impact_coef * participation_rate * 10000.0
            },
            ImpactModelType::SquareRoot => {
                self.temporary_impact_coef * participation_rate.sqrt() * 10000.0
            },
            ImpactModelType::PowerLaw { exponent } => {
                self.temporary_impact_coef * participation_rate.powf(exponent) * 10000.0
            },
            ImpactModelType::AlmgrenChriss => {
                self.calculate_almgren_chriss_impact(
                    participation_rate,
                    spread_bps,
                    microstructure.volatility,
                    order_size,
                    microstructure.avg_trade_size,
                )
            },
            ImpactModelType::IStarModel => {
                self.calculate_istar_impact(
                    order_size,
                    microstructure,
                    orderbook,
                    side,
                )
            },
        };
        
        // Calculate permanent impact (usually smaller)
        let perm_impact_bps = self.permanent_impact_coef * participation_rate.sqrt() * 10000.0;
        
        // Add spread cost
        let spread_cost_bps = spread_bps * self.spread_impact_weight;
        
        // Apply volatility adjustment
        let adjusted_temp_impact = temp_impact_bps * vol_adjustment;
        let adjusted_perm_impact = perm_impact_bps * vol_adjustment;
        
        // Calculate total impact
        let total_impact_bps = adjusted_temp_impact + adjusted_perm_impact + spread_cost_bps;
        
        // Calculate impact decay time (how long temporary impact lasts)
        let impact_decay_ms = self.calculate_impact_decay_time(
            order_size,
            microstructure.daily_volume,
            microstructure.avg_trade_size,
        );
        
        // Calculate expected cost
        let mid_price = if !orderbook.is_empty() {
            orderbook[0].price
        } else {
            100.0 // Default if no orderbook
        };
        
        let direction_multiplier = match side {
            Side::Buy => 1.0,
            Side::Sell => -1.0,
        };
        
        let expected_cost = mid_price * order_size * total_impact_bps / 10000.0 * direction_multiplier;
        
        MarketImpactEstimate {
            temporary_impact: adjusted_temp_impact,
            permanent_impact: adjusted_perm_impact,
            total_impact: total_impact_bps,
            expected_cost: expected_cost.abs(),
            impact_decay_ms,
        }
    }
    
    fn calculate_almgren_chriss_impact(
        &self,
        participation_rate: f64,
        spread_bps: f64,
        volatility: f64,
        order_size: f64,
        avg_trade_size: f64,
    ) -> f64 {
        // Almgren-Chriss model parameters
        let eta = self.temporary_impact_coef; // Temporary impact coefficient
        let gamma = self.permanent_impact_coef; // Permanent impact coefficient
        let trading_rate = order_size / avg_trade_size;
        
        // Temporary impact: eta * (v/V)^alpha * sigma
        let temp_component = eta * participation_rate.sqrt() * volatility * 10000.0;
        
        // Permanent impact: gamma * (X/V)
        let perm_component = gamma * trading_rate * 10000.0;
        
        // Add half spread
        let spread_component = spread_bps * 0.5;
        
        temp_component + perm_component + spread_component
    }
    
    fn calculate_istar_impact(
        &self,
        order_size: f64,
        microstructure: &MarketMicrostructure,
        orderbook: &[PriceLevel],
        _side: Side,
    ) -> f64 {
        // I* model - uses order book shape
        if orderbook.is_empty() {
            return self.temporary_impact_coef * 100.0; // Fallback
        }
        
        // Calculate order book imbalance
        let mut cumulative_size = 0.0;
        let mut impact_bps = 0.0;
        
        // Walk through the book until we've "consumed" our order
        for (_i, level) in orderbook.iter().enumerate() {
            cumulative_size += level.size;
            
            if cumulative_size >= order_size {
                // Calculate average price impact to this level
                let ref_price = orderbook[0].price;
                let exec_price = level.price;
                impact_bps = ((exec_price - ref_price).abs() / ref_price) * 10000.0;
                break;
            }
        }
        
        // Add participation rate impact
        let participation_impact = self.temporary_impact_coef * 
            (order_size / microstructure.daily_volume).sqrt() * 10000.0;
        
        impact_bps + participation_impact
    }
    
    fn get_expected_volume(
        &self,
        microstructure: &MarketMicrostructure,
        current_time: DateTime<Utc>,
    ) -> f64 {
        // Use intraday volume profile if available
        if microstructure.intraday_volume_profile.len() == 24 {
            let hour = current_time.hour() as usize;
            let hour_pct = microstructure.intraday_volume_profile[hour];
            microstructure.daily_volume * hour_pct
        } else {
            // Simple assumption: 1/6.5 of daily volume per hour (6.5 hour trading day)
            microstructure.daily_volume / 6.5
        }
    }
    
    fn calculate_impact_decay_time(
        &self,
        order_size: f64,
        daily_volume: f64,
        avg_trade_size: f64,
    ) -> i64 {
        // Empirical formula for impact decay
        // Larger orders relative to volume decay slower
        let volume_ratio = order_size / daily_volume;
        let trade_ratio = order_size / avg_trade_size;
        
        // Base decay time in milliseconds
        let base_decay_ms = 60_000; // 1 minute base
        
        // Adjust based on order characteristics
        let decay_multiplier = 1.0 + volume_ratio * 10.0 + trade_ratio.ln().max(0.0);
        
        (base_decay_ms as f64 * decay_multiplier) as i64
    }
    
    pub fn calculate_optimal_execution_schedule(
        &self,
        total_size: f64,
        time_horizon_minutes: f64,
        microstructure: &MarketMicrostructure,
        risk_aversion: f64,
    ) -> Vec<(f64, f64)> {
        // Almgren-Chriss optimal execution trajectory
        let n_slices = (time_horizon_minutes / 5.0).ceil() as usize; // 5-minute buckets
        let tau = time_horizon_minutes / n_slices as f64;
        
        let mut schedule = Vec::with_capacity(n_slices);
        
        // Parameters
        let volatility = microstructure.volatility;
        let _daily_volume = microstructure.daily_volume;
        let eta = self.temporary_impact_coef;
        let _gamma = self.permanent_impact_coef;
        let lambda = risk_aversion;
        
        // Calculate optimal trading rate
        let kappa = lambda * volatility.powi(2) / eta;
        let alpha = (kappa / tau).sqrt();
        
        for i in 0..n_slices {
            let t = i as f64 * tau;
            let t_next = (i + 1) as f64 * tau;
            
            // Optimal trajectory: x(t) = X * sinh(alpha * (T - t)) / sinh(alpha * T)
            let remaining_start = total_size * (alpha * (time_horizon_minutes - t)).sinh() 
                / (alpha * time_horizon_minutes).sinh();
            let remaining_end = total_size * (alpha * (time_horizon_minutes - t_next)).sinh() 
                / (alpha * time_horizon_minutes).sinh();
            
            let slice_size = remaining_start - remaining_end;
            let slice_time = t + tau / 2.0; // Midpoint
            
            schedule.push((slice_time, slice_size));
        }
        
        schedule
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_market_impact_models() {
        let microstructure = MarketMicrostructure {
            symbol: "TEST".to_string(),
            avg_spread_bps: 2.0,
            daily_volume: 10_000_000.0,
            avg_trade_size: 100.0,
            volatility: 0.02,
            tick_size: 0.01,
            lot_size: 1.0,
            intraday_volume_profile: vec![0.04; 24], // Flat profile
        };
        
        let orderbook = vec![
            PriceLevel { price: 100.0, size: 1000.0, order_count: Some(10) },
            PriceLevel { price: 100.01, size: 2000.0, order_count: Some(15) },
        ];
        
        let models = vec![
            ImpactModelType::Linear,
            ImpactModelType::SquareRoot,
            ImpactModelType::AlmgrenChriss,
        ];
        
        for model_type in models {
            let model = MarketImpactModel::new(model_type);
            let impact = model.estimate_impact(
                1000.0,
                Side::Buy,
                &microstructure,
                &orderbook,
                Utc::now(),
            );
            
            assert!(impact.total_impact > 0.0);
            assert!(impact.temporary_impact >= 0.0);
            assert!(impact.permanent_impact >= 0.0);
            assert!(impact.expected_cost > 0.0);
            assert!(impact.impact_decay_ms > 0);
        }
    }
}