//! Garleanu-Pedersen (2013) Optimal Execution.
//!
//! Implements the dynamic trading model from "Dynamic Trading with Predictable
//! Returns and Transaction Costs" (Garleanu & Pedersen, 2013).
//!
//! The key insight: with a decaying alpha signal and both temporary and
//! permanent market impact, the optimal strategy trades toward an "aim
//! portfolio" at a rate proportional to an "urgency" parameter.
//!
//! - Urgency captures the trade-off between signal decay and trading costs
//! - The aim portfolio balances the current target with cost considerations
//! - The optimal trajectory is the solution to the resulting ODE
//!
//! All math from scratch in Rust -- no external crate dependencies beyond PyO3.

use pyo3::prelude::*;

// ===========================================================================
// Frozen result types
// ===========================================================================

/// Optimal execution trajectory from the Garleanu-Pedersen model.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct GPTrajectory {
    /// Time points in the trajectory.
    #[pyo3(get)]
    pub times: Vec<f64>,
    /// Position at each time point.
    #[pyo3(get)]
    pub positions: Vec<f64>,
    /// Target position (the final desired holding).
    #[pyo3(get)]
    pub target: f64,
    /// Trading rate (urgency-adjusted).
    #[pyo3(get)]
    pub rate: f64,
    /// Total estimated execution cost.
    #[pyo3(get)]
    pub total_cost: f64,
}

#[pymethods]
impl GPTrajectory {
    fn __repr__(&self) -> String {
        format!(
            "GPTrajectory(steps={}, target={:.4}, rate={:.4}, cost={:.6})",
            self.times.len(),
            self.target,
            self.rate,
            self.total_cost
        )
    }
}

// ===========================================================================
// Private helpers
// ===========================================================================

/// Clamp a value to finite, returning 0.0 for NaN/Inf.
#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

// ===========================================================================
// Public functions
// ===========================================================================

/// Compute the Garleanu-Pedersen urgency parameter.
///
/// Urgency determines how fast to trade toward the aim portfolio.
/// Higher urgency means faster trading (alpha decays quickly, so act now).
///
/// urgency = sqrt(alpha_decay_rate * risk_aversion * volatility^2 / temporary_impact)
///
/// Args:
///     alpha_decay_rate: Rate at which the alpha signal decays (positive).
///     temporary_impact: Temporary market impact coefficient (positive).
///     risk_aversion: Risk aversion parameter (positive).
///     volatility: Asset volatility (positive).
///
/// Returns:
///     Urgency parameter (positive scalar).
#[pyfunction]
pub fn gp_urgency(
    alpha_decay_rate: f64,
    temporary_impact: f64,
    risk_aversion: f64,
    volatility: f64,
) -> PyResult<f64> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if !alpha_decay_rate.is_finite() || alpha_decay_rate <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "alpha_decay_rate must be a positive finite number",
        ));
    }
    if !temporary_impact.is_finite() || temporary_impact <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temporary_impact must be a positive finite number",
        ));
    }
    if !risk_aversion.is_finite() || risk_aversion <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "risk_aversion must be a positive finite number",
        ));
    }
    if !volatility.is_finite() || volatility <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "volatility must be a positive finite number",
        ));
    }

    let urgency = (alpha_decay_rate * risk_aversion * volatility * volatility / temporary_impact).sqrt();
    Ok(finite_or_zero(urgency))
}

/// Compute the Garleanu-Pedersen optimal execution trajectory.
///
/// Solves the G-P (2013) optimal trading problem with exponentially decaying
/// alpha signal. The optimal strategy trades toward an "aim portfolio" at a
/// rate proportional to the urgency parameter.
///
/// The aim portfolio at time t is:
///     aim(t) = target * alpha_weight(t)
/// where alpha_weight(t) = alpha_decay / (alpha_decay + urgency)
///
/// The position evolves as:
///     dx/dt = urgency * (aim(t) - x(t))
///
/// Args:
///     current_position: Starting position.
///     target_position: Desired final position (alpha-implied).
///     alpha_decay_rate: Rate of alpha signal decay (positive).
///     temporary_impact: Temporary market impact coefficient (positive).
///     permanent_impact: Permanent market impact coefficient (non-negative).
///     risk_aversion: Risk aversion parameter (positive).
///     n_steps: Number of discrete time steps in the trajectory.
///     dt: Time step size.
///
/// Returns:
///     GPTrajectory with optimal positions and cost estimate.
#[pyfunction]
#[pyo3(signature = (current_position, target_position, alpha_decay_rate, temporary_impact, permanent_impact, risk_aversion, n_steps, dt))]
pub fn gp_optimal_trajectory(
    current_position: f64,
    target_position: f64,
    alpha_decay_rate: f64,
    temporary_impact: f64,
    permanent_impact: f64,
    risk_aversion: f64,
    n_steps: usize,
    dt: f64,
) -> PyResult<GPTrajectory> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    // Validate inputs
    if !current_position.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "current_position must be finite",
        ));
    }
    if !target_position.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "target_position must be finite",
        ));
    }
    if !alpha_decay_rate.is_finite() || alpha_decay_rate <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "alpha_decay_rate must be a positive finite number",
        ));
    }
    if !temporary_impact.is_finite() || temporary_impact <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temporary_impact must be a positive finite number",
        ));
    }
    if !permanent_impact.is_finite() || permanent_impact < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "permanent_impact must be a non-negative finite number",
        ));
    }
    if !risk_aversion.is_finite() || risk_aversion <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "risk_aversion must be a positive finite number",
        ));
    }
    if n_steps == 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "n_steps must be at least 1",
        ));
    }
    if !dt.is_finite() || dt <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "dt must be a positive finite number",
        ));
    }

    let n_steps = n_steps.min(100_000); // cap for safety

    // Compute urgency parameter
    // Use a proxy volatility of 1.0 for the urgency computation in the trajectory
    // (the risk_aversion already encodes the vol-scaled preference)
    let urgency = (alpha_decay_rate * risk_aversion / temporary_impact).sqrt();
    let urgency = if urgency.is_finite() && urgency > 0.0 {
        urgency
    } else {
        1.0
    };

    // Alpha weight: fraction of target to aim for, accounting for decay vs urgency
    let alpha_weight = alpha_decay_rate / (alpha_decay_rate + urgency);
    let alpha_weight = finite_or_zero(alpha_weight.clamp(0.0, 1.0));

    // Build trajectory by forward Euler integration
    let mut times = Vec::with_capacity(n_steps + 1);
    let mut positions = Vec::with_capacity(n_steps + 1);

    let mut x = current_position;
    let mut t = 0.0;
    times.push(t);
    positions.push(x);

    let mut total_cost = 0.0;
    let total_time = n_steps as f64 * dt;

    for step in 0..n_steps {
        t = (step + 1) as f64 * dt;
        let remaining = total_time - t;

        // Alpha signal decays exponentially
        let alpha_signal = (-alpha_decay_rate * t).exp();
        let alpha_signal = finite_or_zero(alpha_signal);

        // Aim portfolio: blend target with decay
        // aim = target * alpha_weight * alpha_signal
        // The aim shrinks as the alpha decays
        let aim = target_position * alpha_weight * alpha_signal;
        let aim = finite_or_zero(aim);

        // Also account for terminal urgency: as we near the end,
        // aim shifts based on remaining time
        let time_weight = if remaining > 0.0 {
            1.0 - (-urgency * remaining).exp()
        } else {
            1.0
        };
        let time_weight = finite_or_zero(time_weight.clamp(0.0, 1.0));

        let effective_aim = aim + (target_position - aim) * time_weight;
        let effective_aim = finite_or_zero(effective_aim);

        // Trading rate: trade toward aim at speed proportional to urgency
        let trade_rate = urgency * (effective_aim - x);
        let trade_size = trade_rate * dt;
        let trade_size = finite_or_zero(trade_size);

        // Accumulate costs
        // Temporary cost: proportional to trade rate squared
        let temp_cost = 0.5 * temporary_impact * trade_size * trade_size;
        // Permanent cost: proportional to absolute trade size
        let perm_cost = permanent_impact * trade_size.abs() * x.abs();

        total_cost += finite_or_zero(temp_cost + perm_cost);

        x += trade_size;
        x = finite_or_zero(x);

        times.push(t);
        positions.push(x);
    }

    Ok(GPTrajectory {
        times,
        positions,
        target: target_position,
        rate: finite_or_zero(urgency),
        total_cost: finite_or_zero(total_cost),
    })
}

/// Compute the execution cost of a given GP trajectory.
///
/// Args:
///     trajectory: A GPTrajectory object.
///     temporary_impact: Temporary market impact coefficient.
///     permanent_impact: Permanent market impact coefficient.
///
/// Returns:
///     Total execution cost.
#[pyfunction]
pub fn gp_execution_cost(
    trajectory: &GPTrajectory,
    temporary_impact: f64,
    permanent_impact: f64,
) -> PyResult<f64> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if !temporary_impact.is_finite() || temporary_impact < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temporary_impact must be a non-negative finite number",
        ));
    }
    if !permanent_impact.is_finite() || permanent_impact < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "permanent_impact must be a non-negative finite number",
        ));
    }

    let positions = &trajectory.positions;
    let times = &trajectory.times;
    if positions.len() < 2 {
        return Ok(0.0);
    }

    let mut cost = 0.0;

    for i in 1..positions.len() {
        let dt = times[i] - times[i - 1];
        if dt <= 0.0 {
            continue;
        }
        let trade = positions[i] - positions[i - 1];
        let trade_rate = trade / dt;

        // Temporary impact: 0.5 * eta * (dx/dt)^2 * dt
        let temp = 0.5 * temporary_impact * trade_rate * trade_rate * dt;
        // Permanent impact: gamma * |trade| * |position|
        let perm = permanent_impact * trade.abs() * positions[i - 1].abs();

        cost += finite_or_zero(temp + perm);
    }

    Ok(finite_or_zero(cost))
}

/// Register Garleanu-Pedersen types and functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<GPTrajectory>()?;
    m.add_function(wrap_pyfunction!(gp_urgency, m)?)?;
    m.add_function(wrap_pyfunction!(gp_optimal_trajectory, m)?)?;
    m.add_function(wrap_pyfunction!(gp_execution_cost, m)?)?;
    Ok(())
}

// ===========================================================================
// Tests
// ===========================================================================

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

    #[test]
    fn test_finite_or_zero() {
        assert_eq!(finite_or_zero(3.14), 3.14);
        assert_eq!(finite_or_zero(f64::NAN), 0.0);
        assert_eq!(finite_or_zero(f64::INFINITY), 0.0);
    }

    #[test]
    fn test_gp_urgency_basic() {
        let urgency = gp_urgency(0.5, 0.01, 1.0, 0.2).unwrap();
        assert!(urgency > 0.0);
        assert!(urgency.is_finite());
    }

    #[test]
    fn test_gp_urgency_higher_decay_means_higher_urgency() {
        let u1 = gp_urgency(0.1, 0.01, 1.0, 0.2).unwrap();
        let u2 = gp_urgency(1.0, 0.01, 1.0, 0.2).unwrap();
        assert!(u2 > u1, "higher decay rate should increase urgency");
    }

    #[test]
    fn test_gp_urgency_invalid_inputs() {
        assert!(gp_urgency(-1.0, 0.01, 1.0, 0.2).is_err());
        assert!(gp_urgency(0.5, -0.01, 1.0, 0.2).is_err());
        assert!(gp_urgency(0.5, 0.01, -1.0, 0.2).is_err());
        assert!(gp_urgency(0.5, 0.01, 1.0, -0.2).is_err());
        assert!(gp_urgency(0.5, 0.0, 1.0, 0.2).is_err());
        assert!(gp_urgency(f64::NAN, 0.01, 1.0, 0.2).is_err());
    }

    #[test]
    fn test_gp_trajectory_basic() {
        let traj = gp_optimal_trajectory(0.0, 100.0, 0.5, 0.01, 0.001, 1.0, 50, 0.1).unwrap();
        assert_eq!(traj.times.len(), 51); // n_steps + 1
        assert_eq!(traj.positions.len(), 51);
        assert!((traj.positions[0] - 0.0).abs() < 1e-10);
        assert!(traj.target == 100.0);
        assert!(traj.rate > 0.0);
        assert!(traj.total_cost >= 0.0);
    }

    #[test]
    fn test_gp_trajectory_moves_toward_target() {
        let traj = gp_optimal_trajectory(0.0, 100.0, 0.5, 0.01, 0.0, 1.0, 100, 0.1).unwrap();
        // Position should increase over time (moving toward target of 100)
        let final_pos = *traj.positions.last().unwrap();
        assert!(final_pos > 0.0, "should move toward target");
        // Should be closer to target than start
        assert!(
            (final_pos - 100.0).abs() < (0.0f64 - 100.0).abs(),
            "should be closer to target at the end"
        );
    }

    #[test]
    fn test_gp_trajectory_sell() {
        // Start at 100, target 0 (liquidation)
        let traj = gp_optimal_trajectory(100.0, 0.0, 0.5, 0.01, 0.0, 1.0, 50, 0.1).unwrap();
        let final_pos = *traj.positions.last().unwrap();
        assert!(final_pos < 100.0, "should decrease toward target");
    }

    #[test]
    fn test_gp_trajectory_already_at_target() {
        let traj = gp_optimal_trajectory(50.0, 50.0, 0.5, 0.01, 0.0, 1.0, 20, 0.1).unwrap();
        // Should stay near target with minimal trading
        for pos in &traj.positions {
            assert!(
                (pos - 50.0).abs() < 10.0,
                "should stay near target, got {}",
                pos
            );
        }
    }

    #[test]
    fn test_gp_trajectory_invalid_inputs() {
        assert!(gp_optimal_trajectory(0.0, 100.0, -0.5, 0.01, 0.0, 1.0, 50, 0.1).is_err());
        assert!(gp_optimal_trajectory(0.0, 100.0, 0.5, 0.0, 0.0, 1.0, 50, 0.1).is_err());
        assert!(gp_optimal_trajectory(0.0, 100.0, 0.5, 0.01, -1.0, 1.0, 50, 0.1).is_err());
        assert!(gp_optimal_trajectory(0.0, 100.0, 0.5, 0.01, 0.0, 0.0, 50, 0.1).is_err());
        assert!(gp_optimal_trajectory(0.0, 100.0, 0.5, 0.01, 0.0, 1.0, 0, 0.1).is_err());
        assert!(gp_optimal_trajectory(0.0, 100.0, 0.5, 0.01, 0.0, 1.0, 50, 0.0).is_err());
        assert!(gp_optimal_trajectory(f64::NAN, 100.0, 0.5, 0.01, 0.0, 1.0, 50, 0.1).is_err());
    }

    #[test]
    fn test_gp_execution_cost() {
        let traj = gp_optimal_trajectory(0.0, 100.0, 0.5, 0.01, 0.001, 1.0, 50, 0.1).unwrap();
        let cost = gp_execution_cost(&traj, 0.01, 0.001).unwrap();
        assert!(cost >= 0.0);
        assert!(cost.is_finite());
    }

    #[test]
    fn test_gp_execution_cost_no_trade() {
        // If position doesn't change, cost should be near zero
        let traj = GPTrajectory {
            times: vec![0.0, 1.0, 2.0],
            positions: vec![50.0, 50.0, 50.0],
            target: 50.0,
            rate: 1.0,
            total_cost: 0.0,
        };
        let cost = gp_execution_cost(&traj, 0.01, 0.001).unwrap();
        assert!(cost < 1e-10);
    }

    #[test]
    fn test_gp_execution_cost_invalid() {
        let traj = GPTrajectory {
            times: vec![0.0, 1.0],
            positions: vec![0.0, 10.0],
            target: 10.0,
            rate: 1.0,
            total_cost: 0.0,
        };
        assert!(gp_execution_cost(&traj, -1.0, 0.0).is_err());
        assert!(gp_execution_cost(&traj, 0.0, -1.0).is_err());
    }

    #[test]
    fn test_gp_higher_impact_means_higher_cost() {
        let traj = gp_optimal_trajectory(0.0, 100.0, 0.5, 0.01, 0.0, 1.0, 50, 0.1).unwrap();
        let cost_low = gp_execution_cost(&traj, 0.001, 0.0).unwrap();
        let cost_high = gp_execution_cost(&traj, 0.1, 0.0).unwrap();
        assert!(cost_high > cost_low);
    }

    #[test]
    fn test_gp_trajectory_monotone_buying() {
        // When buying (target > current), positions should generally increase
        let traj = gp_optimal_trajectory(0.0, 100.0, 0.5, 0.01, 0.0, 1.0, 20, 0.1).unwrap();
        for i in 1..traj.positions.len() {
            assert!(
                traj.positions[i] >= traj.positions[i - 1] - 1e-10,
                "position should not decrease when buying"
            );
        }
    }
}
