//! P&L Attribution Calculator
//!
//! Decomposes realized trade P&L into per-model contributions.
//! Each model gets credit proportional to:
//! - Its ensemble weight at trade time
//! - Whether its signal aligned with the realized direction

use ml::ensemble::decision::EnsembleDecision;

/// Per-model attribution for a single trade
#[derive(Debug, Clone)]
pub struct TradeAttribution {
    /// Model identifier
    pub model_id: String,
    /// Model's ensemble weight at trade time
    pub model_weight: f64,
    /// Model's raw signal (-1.0 to 1.0)
    pub model_signal: f64,
    /// 1.0 if signal direction matched realized direction, -1.0 otherwise
    pub signal_alignment: f64,
    /// Attributed P&L contribution
    pub pnl_contribution: f64,
}

/// Full attribution result for a closed trade
#[derive(Debug, Clone)]
pub struct AttributionResult {
    /// Per-model attributions
    pub attributions: Vec<TradeAttribution>,
    /// Total realized P&L (should equal sum of contributions)
    pub realized_pnl: f64,
    /// Residual (rounding error, should be near zero)
    pub residual: f64,
}

/// Calculate per-model P&L attribution from an ensemble decision and realized P&L.
///
/// For each model that voted:
/// - `signal_alignment = 1.0` if sign(model_signal) == sign(realized_pnl), else `-1.0`
/// - `pnl_contribution = model_weight × signal_alignment × |realized_pnl|`
///
/// Models with zero signal are treated as neutral (alignment = 0.0).
pub fn attribute(decision: &EnsembleDecision, realized_pnl: f64) -> AttributionResult {
    let votes = &decision.model_votes;

    if votes.is_empty() || realized_pnl.abs() < 1e-12 {
        return AttributionResult {
            attributions: votes
                .iter()
                .map(|(id, v)| TradeAttribution {
                    model_id: id.clone(),
                    model_weight: v.weight,
                    model_signal: v.signal,
                    signal_alignment: 0.0,
                    pnl_contribution: 0.0,
                })
                .collect(),
            realized_pnl,
            residual: realized_pnl,
        };
    }

    let realized_direction = realized_pnl.signum();

    let attributions: Vec<TradeAttribution> = votes
        .iter()
        .map(|(id, vote)| {
            let alignment = compute_alignment(vote.signal, realized_direction);
            let contribution = vote.weight * alignment * realized_pnl.abs();

            TradeAttribution {
                model_id: id.clone(),
                model_weight: vote.weight,
                model_signal: vote.signal,
                signal_alignment: alignment,
                pnl_contribution: contribution,
            }
        })
        .collect();

    let total_attributed: f64 = attributions.iter().map(|a| a.pnl_contribution).sum();
    let residual = realized_pnl - total_attributed;

    AttributionResult {
        attributions,
        realized_pnl,
        residual,
    }
}

/// Compute signal alignment: does the model's signal direction match the realized direction?
///
/// - Zero signal → neutral (0.0)
/// - Same sign → aligned (1.0)
/// - Opposite sign → misaligned (-1.0)
fn compute_alignment(model_signal: f64, realized_direction: f64) -> f64 {
    if model_signal.abs() < 1e-12 {
        return 0.0;
    }
    if model_signal.signum() == realized_direction {
        1.0
    } else {
        -1.0
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ml::ensemble::decision::{ModelVote, TradingAction};
    use std::collections::HashMap;

    fn make_decision(votes: Vec<(&str, f64, f64)>) -> EnsembleDecision {
        let mut model_votes = HashMap::new();
        let mut total_signal = 0.0;

        for (id, signal, weight) in &votes {
            total_signal += signal * weight;
            model_votes.insert(
                id.to_string(),
                ModelVote::new(id.to_string(), *signal, signal.abs(), *weight)
                    .with_model_type("DQN".into()),
            );
        }

        EnsembleDecision::new(
            if total_signal > 0.0 {
                TradingAction::Buy
            } else if total_signal < 0.0 {
                TradingAction::Sell
            } else {
                TradingAction::Hold
            },
            0.8,
            total_signal,
            0.0,
            model_votes,
        )
    }

    #[test]
    fn test_correct_attribution_positive_pnl() {
        // Two models, both bullish, trade was profitable
        let decision = make_decision(vec![
            ("dqn", 0.8, 0.6), // 60% weight, bullish
            ("ppo", 0.5, 0.4), // 40% weight, bullish
        ]);

        let result = attribute(&decision, 100.0);
        assert_eq!(result.attributions.len(), 2);

        // DQN: 0.6 × 1.0 × 100.0 = 60.0
        let dqn = result.attributions.iter().find(|a| a.model_id == "dqn");
        assert!(dqn.is_some());
        if let Some(dqn) = dqn {
            assert!((dqn.pnl_contribution - 60.0).abs() < 1e-10);
            assert!((dqn.signal_alignment - 1.0).abs() < 1e-10);
        }

        // PPO: 0.4 × 1.0 × 100.0 = 40.0
        let ppo = result.attributions.iter().find(|a| a.model_id == "ppo");
        assert!(ppo.is_some());
        if let Some(ppo) = ppo {
            assert!((ppo.pnl_contribution - 40.0).abs() < 1e-10);
        }

        // Sum should equal realized PnL
        assert!(result.residual.abs() < 1e-10);
    }

    #[test]
    fn test_attribution_with_disagreement() {
        // DQN bullish, PPO bearish, trade was profitable (bullish correct)
        let decision = make_decision(vec![
            ("dqn", 0.8, 0.6),  // bullish, correct
            ("ppo", -0.5, 0.4), // bearish, wrong
        ]);

        let result = attribute(&decision, 100.0);

        // DQN: 0.6 × 1.0 × 100.0 = 60.0 (aligned)
        if let Some(dqn) = result.attributions.iter().find(|a| a.model_id == "dqn") {
            assert!((dqn.pnl_contribution - 60.0).abs() < 1e-10);
        }

        // PPO: 0.4 × (-1.0) × 100.0 = -40.0 (misaligned)
        if let Some(ppo) = result.attributions.iter().find(|a| a.model_id == "ppo") {
            assert!((ppo.pnl_contribution - (-40.0)).abs() < 1e-10);
        }
    }

    #[test]
    fn test_zero_pnl() {
        let decision = make_decision(vec![("dqn", 0.8, 0.5), ("ppo", 0.5, 0.5)]);

        let result = attribute(&decision, 0.0);

        // All contributions should be zero
        for attr in &result.attributions {
            assert!(attr.pnl_contribution.abs() < 1e-12);
            assert!(attr.signal_alignment.abs() < 1e-12);
        }
    }

    #[test]
    fn test_all_models_wrong() {
        // All bearish, but market went up → they were wrong
        let decision = make_decision(vec![("dqn", -0.7, 0.5), ("ppo", -0.6, 0.5)]);

        let result = attribute(&decision, 50.0);

        // Both bearish but pnl positive → misaligned
        for attr in &result.attributions {
            assert!((attr.signal_alignment - (-1.0)).abs() < 1e-10);
            assert!(attr.pnl_contribution < 0.0);
        }
    }

    #[test]
    fn test_all_models_correct_bearish() {
        // All bearish, PnL negative → bearish direction correct
        let decision = make_decision(vec![("dqn", -0.7, 0.5), ("ppo", -0.6, 0.5)]);

        let result = attribute(&decision, -100.0);

        // Both bearish, PnL negative → aligned
        for attr in &result.attributions {
            assert!((attr.signal_alignment - 1.0).abs() < 1e-10);
            assert!(attr.pnl_contribution > 0.0);
        }
    }

    #[test]
    fn test_neutral_model_gets_zero() {
        let decision = make_decision(vec![
            ("dqn", 0.8, 0.5),
            ("neutral", 0.0, 0.5), // Zero signal = neutral
        ]);

        let result = attribute(&decision, 100.0);

        if let Some(neutral) = result.attributions.iter().find(|a| a.model_id == "neutral") {
            assert!(neutral.signal_alignment.abs() < 1e-12);
            assert!(neutral.pnl_contribution.abs() < 1e-12);
        }
    }

    #[test]
    fn test_empty_votes() {
        let decision = EnsembleDecision::new(
            TradingAction::Hold,
            0.0,
            0.0,
            0.0,
            HashMap::new(),
        );

        let result = attribute(&decision, 100.0);
        assert!(result.attributions.is_empty());
        assert!((result.residual - 100.0).abs() < 1e-10);
    }
}
