//! Kelly Criterion Optimizer for Position Sizing
//!
//! Implements Kelly Criterion and enhanced Kelly strategies for optimal position sizing
//! using canonical types from the types crate.

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

// Import types from crate root (which imports from common)
use crate::{MLError, MLResult as Result};
use common::types::Price;

/// Kelly position recommendation using canonical types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KellyPositionRecommendation {
    pub asset_id: String, // Using String until AssetId is implemented
    pub recommended_fraction: f64,
    pub expected_return: f64,
    pub volatility: f64,
    pub win_probability: f64,
    pub avg_win: Price,
    pub avg_loss: Price,
    pub max_fraction: f64,
    pub confidence: f64,
    pub timestamp: DateTime<Utc>,
}

/// Kelly optimizer configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KellyOptimizerConfig {
    pub max_fraction: f64,
    pub min_fraction: f64,
    pub lookback_period: usize,
    pub confidence_threshold: f64,
    pub volatility_adjustment: bool,
    pub drawdown_protection: bool,
}

impl Default for KellyOptimizerConfig {
    fn default() -> Self {
        Self {
            max_fraction: 0.25,        // Maximum 25% of portfolio
            min_fraction: 0.01,        // Minimum 1% of portfolio
            lookback_period: 252,      // 1 year of daily data
            confidence_threshold: 0.6, // 60% minimum confidence
            volatility_adjustment: true,
            drawdown_protection: true,
        }
    }
}

/// Kelly Criterion optimizer
#[derive(Debug)]
pub struct KellyCriterionOptimizer {
    config: KellyOptimizerConfig,
}

impl KellyCriterionOptimizer {
    pub fn new(config: KellyOptimizerConfig) -> Result<Self> {
        Ok(Self { config })
    }

    /// Calculate basic Kelly fraction: f = (bp - q) / b
    /// where:
    /// - b = odds (avg_win / avg_loss)
    /// - p = probability of winning
    /// - q = probability of losing (1 - p)
    pub fn calculate_basic_kelly(
        &self,
        win_probability: f64,
        avg_win: f64,
        avg_loss: f64,
    ) -> Result<f64> {
        if win_probability <= 0.0 || win_probability >= 1.0 {
            return Err(MLError::InvalidInput(
                "Win probability must be between 0 and 1".to_owned(),
            ));
        }

        if avg_win <= 0.0 || avg_loss <= 0.0 {
            return Err(MLError::InvalidInput(
                "Average win and loss must be positive".to_owned(),
            ));
        }

        let odds = avg_win / avg_loss;
        let q = 1.0 - win_probability;

        let kelly_fraction = (odds * win_probability - q) / odds;

        // Apply safety constraints
        Ok(kelly_fraction.clamp(self.config.min_fraction, self.config.max_fraction))
    }

    /// Calculate enhanced Kelly with volatility adjustment
    pub fn calculate_enhanced_kelly(
        &self,
        expected_return: f64,
        variance: f64,
        win_probability: f64,
        avg_win: f64,
        avg_loss: f64,
    ) -> Result<f64> {
        if variance <= 0.0 {
            return Err(MLError::InvalidInput(
                "Variance must be positive".to_owned(),
            ));
        }

        // Standard Kelly: f = μ / σ²
        let standard_kelly = expected_return / variance;

        // Basic Kelly from win/loss statistics
        let basic_kelly = self.calculate_basic_kelly(win_probability, avg_win, avg_loss)?;

        // Combine both approaches with weighting
        let combined_kelly = if self.config.volatility_adjustment {
            0.7 * standard_kelly + 0.3 * basic_kelly
        } else {
            basic_kelly
        };

        // Apply safety constraints
        Ok(combined_kelly.clamp(self.config.min_fraction, self.config.max_fraction))
    }

    /// Generate position recommendation for an asset
    pub fn recommend_position(
        &self,
        asset_id: String,
        historical_returns: &[f64],
    ) -> Result<KellyPositionRecommendation> {
        if historical_returns.is_empty() {
            return Err(MLError::InvalidInput(
                "Empty historical returns".to_owned(),
            ));
        }

        // Calculate statistics
        let mean_return = historical_returns.iter().sum::<f64>() / historical_returns.len() as f64;
        let variance = historical_returns
            .iter()
            .map(|r| (r - mean_return).powi(2))
            .sum::<f64>()
            / historical_returns.len() as f64;
        let volatility = variance.sqrt();

        // Calculate win/loss statistics
        let wins: Vec<f64> = historical_returns
            .iter()
            .filter(|&&r| r > 0.0)
            .copied()
            .collect();
        let losses: Vec<f64> = historical_returns
            .iter()
            .filter(|&&r| r < 0.0)
            .map(|r| -r)
            .collect();

        let win_probability = wins.len() as f64 / historical_returns.len() as f64;
        let avg_win = if wins.is_empty() {
            0.01
        } else {
            wins.iter().sum::<f64>() / wins.len() as f64
        };
        let avg_loss = if losses.is_empty() {
            0.01
        } else {
            losses.iter().sum::<f64>() / losses.len() as f64
        };

        // Calculate recommended fraction
        let recommended_fraction = self.calculate_enhanced_kelly(
            mean_return,
            variance,
            win_probability,
            avg_win,
            avg_loss,
        )?;

        // Calculate confidence based on sample size and consistency
        let confidence = (historical_returns.len() as f64 / self.config.lookback_period as f64)
            .min(1.0)
            * win_probability;

        Ok(KellyPositionRecommendation {
            asset_id,
            recommended_fraction,
            expected_return: mean_return,
            volatility,
            win_probability,
            avg_win: Price::from_f64(avg_win).map_err(|e| {
                MLError::InvalidInput(format!("Failed to convert avg_win to Price: {}", e))
            })?,
            avg_loss: Price::from_f64(avg_loss).map_err(|e| {
                MLError::InvalidInput(format!("Failed to convert avg_loss to Price: {}", e))
            })?,
            max_fraction: self.config.max_fraction,
            confidence,
            timestamp: Utc::now(),
        })
    }

    /// Calculate fractional Kelly to reduce risk
    pub fn calculate_fractional_kelly(&self, kelly_fraction: f64, fraction: f64) -> f64 {
        (kelly_fraction * fraction).clamp(self.config.min_fraction, self.config.max_fraction)
    }
}

#[cfg(test)]
#[allow(clippy::assertions_on_result_states)]
mod tests {
    use super::*;


    #[test]
    fn test_basic_kelly_calculation() -> Result<()> {
        let config = KellyOptimizerConfig::default();
        let optimizer = KellyCriterionOptimizer::new(config)?;

        // Test with favorable odds
        let kelly = optimizer.calculate_basic_kelly(0.6, 2.0, 1.0)?;
        assert!(kelly > 0.0);
        assert!(kelly <= 0.25); // Should be capped at max_fraction

        // Test with unfavorable odds
        let kelly = optimizer.calculate_basic_kelly(0.4, 1.0, 2.0)?;
        assert!(kelly <= 0.01); // Should be at min_fraction or near zero
        Ok(())
    }

    #[test]
    fn test_enhanced_kelly_calculation() -> Result<()> {
        let config = KellyOptimizerConfig::default();
        let optimizer = KellyCriterionOptimizer::new(config)?;

        let kelly = optimizer.calculate_enhanced_kelly(
            0.1,  // 10% expected return
            0.04, // 4% variance (20% volatility)
            0.6,  // 60% win probability
            0.15, // 15% average win
            0.1,  // 10% average loss
        )?;

        assert!(kelly > 0.0);
        assert!(kelly <= 0.25); // Should respect max_fraction
        Ok(())
    }

    #[test]
    fn test_position_recommendation() -> Result<()> {
        let config = KellyOptimizerConfig::default();
        let optimizer = KellyCriterionOptimizer::new(config.clone())?;

        // Generate some sample returns
        let returns = vec![
            0.1, -0.05, 0.08, -0.03, 0.12, -0.02, 0.06, -0.04, 0.09, -0.01,
        ];
        let asset_id = "AAPL".to_owned();

        let recommendation = optimizer.recommend_position(asset_id, &returns);
        assert!(recommendation.is_ok());

        let rec = recommendation?;
        assert!(rec.recommended_fraction >= 0.0);
        assert!(rec.recommended_fraction <= config.max_fraction);
        assert!(rec.confidence >= 0.0 && rec.confidence <= 1.0);
        assert!(rec.win_probability >= 0.0 && rec.win_probability <= 1.0);
        Ok(())
    }

    #[test]
    fn test_fractional_kelly() -> Result<()> {
        let config = KellyOptimizerConfig::default();
        let optimizer = KellyCriterionOptimizer::new(config.clone())?;

        let full_kelly = 0.2;
        let half_kelly = optimizer.calculate_fractional_kelly(full_kelly, 0.5);

        assert_eq!(half_kelly, 0.1);

        // Test that it respects limits
        let excessive_kelly = optimizer.calculate_fractional_kelly(1.0, 0.5);
        assert!(excessive_kelly <= config.max_fraction);
        Ok(())
    }

    #[test]
    fn test_invalid_inputs() -> Result<()> {
        let config = KellyOptimizerConfig::default();
        let optimizer = KellyCriterionOptimizer::new(config)?;

        // Invalid win probability
        let result = optimizer.calculate_basic_kelly(1.1, 2.0, 1.0);
        assert!(result.is_err());

        // Negative average win
        let result = optimizer.calculate_basic_kelly(0.6, -1.0, 1.0);
        assert!(result.is_err());

        // Zero variance
        let result = optimizer.calculate_enhanced_kelly(0.1, 0.0, 0.6, 0.15, 0.1);
        assert!(result.is_err());
        Ok(())
    }
}
