//! Self-improving enrichment module — 8 enrichment functions.
//!
//! `EnrichmentState` persists across epochs. `EvalTrade` captures per-trade
//! data from the validation backtest. `run_enrichments` dispatches all 8
//! enrichment functions and returns an `EnrichmentResult` summarising the
//! adjustments for the next training epoch.
//!
//! All pure Rust, no CUDA kernels.

use tracing::info;

// ---------------------------------------------------------------------------
// Data structures
// ---------------------------------------------------------------------------

/// Per-trade data extracted from the validation backtest.
#[derive(Debug, Clone)]
pub(crate) struct EvalTrade {
    /// Global bar index where the trade was entered.
    pub bar_index: usize,
    /// Realised P&L (percentage or dollar, consistent units).
    pub pnl: f32,
    /// Predicted Q-value at entry.
    pub predicted_q: f32,
    /// Direction branch output (0=Short, 1=Flat, 2=Long).
    pub direction: u8,
    /// Magnitude branch output (0=Small, 1=Half, 2=Full).
    pub magnitude: u8,
    /// How many bars the position was held.
    pub holding_bars: u32,
    /// Ensemble variance across heads at entry.
    pub ensemble_var: f32,
}

/// Synthetic experience from hindsight replay (E7).
#[derive(Debug, Clone)]
pub(crate) struct HindsightExperience {
    /// Bar index of the losing trade.
    pub bar_index: usize,
    /// Optimal direction (counterfactual).
    pub optimal_direction: u8,
    /// Optimal magnitude (counterfactual).
    pub optimal_magnitude: u8,
    /// Estimated counterfactual P&L.
    pub counterfactual_pnl: f32,
}

/// Output of one enrichment cycle.
#[derive(Debug, Clone)]
pub(crate) struct EnrichmentResult {
    /// E1: additive correction to Q-targets (clipped ±10.0).
    pub q_correction: f32,
    /// E2: next-epoch epsilon.
    pub epsilon: f32,
    /// E3: suggested gamma (None = keep current).
    pub gamma: Option<f32>,
    /// E4: per-branch LR scaling [dir, mag, order, urgency].
    pub branch_lr_scale: [f32; 4],
    /// E5: updated ensemble agreement threshold.
    pub agreement_threshold: f32,
    /// E6: bar indices of top-10% trades (priority replay).
    pub winner_indices: Vec<usize>,
    /// E7: hindsight labels for worst losers.
    pub hindsight: Vec<HindsightExperience>,
    /// E8: curriculum weights (one per segment).
    pub curriculum_weights: Vec<f32>,
}

/// Persistent state across epochs for the enrichment module.
#[derive(Debug, Clone)]
pub(crate) struct EnrichmentState {
    /// Current epsilon (updated each enrichment cycle).
    pub epsilon: f32,
    /// Current ensemble agreement threshold.
    pub agreement_threshold: f32,
    /// Number of enrichment cycles executed so far.
    pub cycle_count: u32,
    /// Number of curriculum segments.
    pub n_segments: usize,
}

impl EnrichmentState {
    pub(crate) fn new(initial_epsilon: f32) -> Self {
        Self {
            epsilon: initial_epsilon,
            agreement_threshold: 1.0,
            cycle_count: 0,
            n_segments: 8,
        }
    }
}

// ---------------------------------------------------------------------------
// Trade extraction from aggregate metrics
// ---------------------------------------------------------------------------

/// Extract approximate trade-level data from validation backtest metrics.
/// MVP: generates approximate trade distribution from aggregate WindowMetrics.
pub(crate) fn extract_eval_trades_from_metrics(
    _val_sharpe: f32,
    total_trades: f32,
    total_pnl: f32,
    win_rate: f32,
    avg_q_value: f32,
    n_bars: usize,
) -> Vec<EvalTrade> {
    let n_trades = total_trades as usize;
    if n_trades == 0 {
        return Vec::new();
    }

    let avg_pnl = total_pnl / n_trades as f32;
    let avg_holding = (n_bars as f32 / n_trades as f32).max(1.0) as u32;

    let mut trades = Vec::with_capacity(n_trades);
    for i in 0..n_trades {
        let frac = i as f32 / n_trades as f32;
        let is_winner = frac < win_rate;
        let pnl = if is_winner {
            avg_pnl.abs() * 1.5
        } else {
            -avg_pnl.abs() * 0.8
        };
        trades.push(EvalTrade {
            bar_index: i * (n_bars / n_trades).max(1),
            direction: if is_winner { 2 } else { 0 },
            magnitude: 1,
            holding_bars: avg_holding,
            pnl,
            predicted_q: avg_q_value,
            ensemble_var: 0.5,
        });
    }
    trades
}

// ---------------------------------------------------------------------------
// Enrichment functions (E1–E8)
// ---------------------------------------------------------------------------

/// E1: Q-value reality check — computes additive bias correction.
fn compute_q_correction(trades: &[EvalTrade]) -> f32 {
    if trades.is_empty() {
        return 0.0;
    }
    let bias: f32 = trades
        .iter()
        .map(|t| t.predicted_q - t.pnl)
        .sum::<f32>()
        / trades.len() as f32;
    (-bias).clamp(-10.0, 10.0)
}

/// E2: Adaptive epsilon — performance-driven exploration rate.
fn compute_adaptive_epsilon(current: f32, val_sharpe: f32) -> f32 {
    let scale = if val_sharpe > 2.0 {
        0.8
    } else if val_sharpe > 0.5 {
        0.95
    } else if val_sharpe < -0.5 {
        1.2
    } else {
        1.0
    };
    (current * scale).clamp(0.01, 0.15)
}

/// E3: Dynamic gamma from profitable trade durations.
fn compute_dynamic_gamma(trades: &[EvalTrade]) -> Option<f32> {
    let holds: Vec<f32> = trades
        .iter()
        .filter(|t| t.pnl > 0.0 && t.holding_bars > 0)
        .map(|t| t.holding_bars as f32)
        .collect();
    if holds.is_empty() {
        return None;
    }
    let mean = holds.iter().sum::<f32>() / holds.len() as f32;
    Some((1.0 - 1.0 / mean.max(5.0)).clamp(0.85, 0.98))
}

/// E4: Trade autopsy — per-branch error rates to LR scaling.
fn compute_branch_lr_scale(trades: &[EvalTrade]) -> [f32; 4] {
    let losers: Vec<&EvalTrade> = trades.iter().filter(|t| t.pnl < 0.0).collect();
    if losers.is_empty() {
        return [1.0; 4];
    }
    let total = losers.len() as f32;
    let mut errors = [0u32; 4];
    for t in &losers {
        // Direction errors: went Long and lost, or went Short and lost
        if (t.direction == 2 && t.pnl < 0.0) || (t.direction == 0 && t.pnl < 0.0) {
            errors[0] += 1;
        }
        // Magnitude errors: oversized and significant loss
        if t.magnitude >= 1 && t.pnl < -0.001 {
            errors[1] += 1;
        }
        // Urgency errors: very short hold and lost
        if t.holding_bars < 3 && t.pnl < 0.0 {
            errors[3] += 1;
        }
    }
    let rates: [f32; 4] = std::array::from_fn(|i| errors[i] as f32 / total);
    let mean_rate = rates.iter().sum::<f32>() / 4.0;
    std::array::from_fn(|d| (1.0 + 0.5 * (rates[d] - mean_rate)).clamp(0.5, 2.0))
}

/// E5: Ensemble agreement threshold tuning.
fn compute_agreement_threshold(
    trades: &[EvalTrade],
    current: f32,
    full_sharpe: f32,
) -> f32 {
    if trades.is_empty() {
        return current;
    }
    let agree: Vec<&EvalTrade> = trades.iter().filter(|t| t.ensemble_var < current).collect();
    if agree.is_empty() {
        return current * 1.1;
    }
    let pnls: Vec<f32> = agree.iter().map(|t| t.pnl).collect();
    let mean = pnls.iter().sum::<f32>() / pnls.len() as f32;
    let var = pnls
        .iter()
        .map(|p| (p - mean).powi(2))
        .sum::<f32>()
        / pnls.len() as f32;
    let agree_sharpe = if var > 1e-12 { mean / var.sqrt() } else { 0.0 };
    let new = if agree_sharpe > full_sharpe * 1.5 {
        current * 0.9
    } else if agree_sharpe < full_sharpe {
        current * 1.1
    } else {
        current
    };
    new.clamp(0.01, 10.0)
}

/// E6: Winner distillation — bar indices of top-10% trades by P&L.
fn compute_winner_indices(trades: &[EvalTrade]) -> Vec<usize> {
    if trades.is_empty() {
        return Vec::new();
    }
    let mut sorted: Vec<(usize, f32)> = trades
        .iter()
        .enumerate()
        .map(|(i, t)| (i, t.pnl))
        .collect();
    sorted.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    let top_n = (sorted.len() / 10).max(1);
    sorted[..top_n]
        .iter()
        .map(|(i, _)| trades[*i].bar_index)
        .collect()
}

/// E7: Hindsight optimal labels for losing trades.
fn compute_hindsight_labels(trades: &[EvalTrade]) -> Vec<HindsightExperience> {
    let take_n = (trades.len() / 10).max(1);
    trades
        .iter()
        .filter(|t| t.pnl < -0.001)
        .take(take_n)
        .map(|t| HindsightExperience {
            bar_index: t.bar_index,
            optimal_direction: if t.direction == 2 {
                0
            } else if t.direction == 0 {
                2
            } else {
                1
            },
            optimal_magnitude: 0,
            counterfactual_pnl: (-t.pnl).max(0.0),
        })
        .collect()
}

/// E8: Curriculum weights — inverse-Sharpe per segment.
fn compute_curriculum_weights(trades: &[EvalTrade], n_segments: usize) -> Vec<f32> {
    if trades.is_empty() || n_segments == 0 {
        return Vec::new();
    }
    let seg_size = (trades.len() / n_segments).max(1);
    let mut weights: Vec<f32> = (0..n_segments)
        .map(|seg| {
            let start = seg * seg_size;
            let end = ((seg + 1) * seg_size).min(trades.len());
            if start >= trades.len() {
                return 0.1;
            }
            let slice = &trades[start..end];
            let mean: f32 =
                slice.iter().map(|t| t.pnl).sum::<f32>() / slice.len().max(1) as f32;
            let var: f32 = slice
                .iter()
                .map(|t| (t.pnl - mean).powi(2))
                .sum::<f32>()
                / slice.len().max(1) as f32;
            let sharpe = mean / var.sqrt().max(1e-6);
            (1.0 / sharpe.max(0.1)).clamp(0.1, 10.0)
        })
        .collect();
    let total: f32 = weights.iter().sum();
    if total > 1e-6 {
        for w in &mut weights {
            *w /= total;
        }
    }
    weights
}

// ---------------------------------------------------------------------------
// Main dispatcher
// ---------------------------------------------------------------------------

/// Run all 8 enrichment functions, update persistent state, and return results.
pub(crate) fn run_enrichments(
    state: &mut EnrichmentState,
    eval_trades: &[EvalTrade],
    val_sharpe: f32,
) -> EnrichmentResult {
    state.cycle_count += 1;

    // E1: Q-value bias correction
    let q_correction = compute_q_correction(eval_trades);

    // E2: Adaptive epsilon
    let epsilon = compute_adaptive_epsilon(state.epsilon, val_sharpe);
    state.epsilon = epsilon;

    // E3: Dynamic gamma
    let gamma = compute_dynamic_gamma(eval_trades);

    // E4: Per-branch LR scaling
    let branch_lr_scale = compute_branch_lr_scale(eval_trades);

    // E5: Ensemble agreement threshold
    let agreement_threshold = compute_agreement_threshold(
        eval_trades,
        state.agreement_threshold,
        val_sharpe,
    );
    state.agreement_threshold = agreement_threshold;

    // E6: Winner distillation
    let winner_indices = compute_winner_indices(eval_trades);

    // E7: Hindsight labels
    let hindsight = compute_hindsight_labels(eval_trades);

    // E8: Curriculum weights
    let curriculum_weights = compute_curriculum_weights(eval_trades, state.n_segments);

    info!(
        "Enrichment cycle {} complete: q_corr={:.4}, eps={:.4}, gamma={:?}, \
         branch_lr={:?}, agree_thr={:.4}, winners={}, hindsight={}, \
         curriculum_segs={}",
        state.cycle_count,
        q_correction,
        epsilon,
        gamma,
        branch_lr_scale,
        agreement_threshold,
        winner_indices.len(),
        hindsight.len(),
        curriculum_weights.len(),
    );

    EnrichmentResult {
        q_correction,
        epsilon,
        gamma,
        branch_lr_scale,
        agreement_threshold,
        winner_indices,
        hindsight,
        curriculum_weights,
    }
}
