//! Multi-strategy composition tree + flat bytecode program for the
//! decision-policy kernel.
//!
//! See docs/superpowers/specs/2026-05-18-real-lob-integration-design.md §6.

pub mod composition;
pub mod latency;
pub mod sizing;

pub use composition::{ConflictPolicy, EnsembleAggregator, Regime};
pub use latency::LatencyConfig;
pub use sizing::{IsvKellyStateHost, SizingPolicyId};

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Re-export of `ml_alpha::heads::N_HORIZONS` — single source of truth.
pub const N_HORIZONS: usize = ml_alpha::heads::N_HORIZONS;

/// One horizon-anchored leaf strategy. See spec §5 (per-horizon ISV-Kelly) + §6.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StrategyConfig {
    pub horizon_idx: u8, // 0..N_HORIZONS
    pub sizing_policy: SizingPolicyId,
    pub max_concurrent_lots: u16,
}

/// Recursive composition tree. Default constructor `default_for(max_lots)`
/// emits the adaptive 5-horizon Ensemble with WeightedByRealizedSharpe.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Strategy {
    Leaf(StrategyConfig),
    Ensemble {
        children: Vec<Strategy>,
        aggregator: EnsembleAggregator,
    },
    Portfolio {
        children: Vec<Strategy>,
        capital_allocation: Vec<f32>,
        conflict_policy: ConflictPolicy,
    },
    RegimeSwitch {
        regime_to_child: HashMap<Regime, Strategy>,
    },
}

impl Strategy {
    /// Default v1 policy: per-horizon adaptive Ensemble. No static
    /// horizon mask; WeightedByRealizedSharpe auto-shifts capital from
    /// in-run evidence. See spec §6 `default_strategy(...)`.
    pub fn default_for(max_lots: u16) -> Self {
        let children = (0..N_HORIZONS as u8)
            .map(|h| {
                Strategy::Leaf(StrategyConfig {
                    horizon_idx: h,
                    sizing_policy: SizingPolicyId::IsvKelly,
                    max_concurrent_lots: max_lots,
                })
            })
            .collect();
        Strategy::Ensemble {
            children,
            aggregator: EnsembleAggregator::WeightedByRealizedSharpe,
        }
    }

    /// Walk the tree and emit a flat bytecode Program. See spec §6
    /// host-side flattening; layout must stay in sync with
    /// cuda/decision_policy.cu (added in C7).
    pub fn flatten(&self) -> Program {
        let mut prog = Program::default();
        flatten_into(self, &mut prog);
        prog.finalize();
        prog
    }
}

fn flatten_into(s: &Strategy, prog: &mut Program) {
    match s {
        Strategy::Leaf(cfg) => {
            prog.instructions.push(Instruction {
                op: OpCode::EmitPerHorizonSize as u8,
                arg0: cfg.horizon_idx,
                arg1: cfg.max_concurrent_lots,
                arg2: 0,
            });
        }
        Strategy::Ensemble { children, aggregator } => {
            for c in children {
                flatten_into(c, prog);
            }
            let agg_op = match aggregator {
                EnsembleAggregator::Mean => OpCode::AggMean,
                EnsembleAggregator::MaxConfidence => OpCode::AggMaxConfidence,
                // v1: treat ValAuc as Mean (no static-weights upload yet).
                EnsembleAggregator::WeightedByValAuc => OpCode::AggMean,
                EnsembleAggregator::WeightedByRealizedSharpe => OpCode::AggWeightedSharpe,
            };
            prog.instructions.push(Instruction {
                op: agg_op as u8,
                arg0: children.len() as u8,
                arg1: 0,
                arg2: 0,
            });
        }
        Strategy::Portfolio { children, capital_allocation: _, conflict_policy } => {
            for c in children {
                flatten_into(c, prog);
            }
            prog.instructions.push(Instruction {
                op: OpCode::AggMean as u8,
                arg0: children.len() as u8,
                arg1: 0,
                arg2: 0,
            });
            prog.instructions.push(Instruction {
                op: OpCode::ApplyConflict as u8,
                arg0: *conflict_policy as u8,
                arg1: 0,
                arg2: 0,
            });
        }
        Strategy::RegimeSwitch { regime_to_child } => {
            let mut branch_indices: Vec<usize> = Vec::with_capacity(regime_to_child.len());
            for (regime, child) in regime_to_child {
                let branch_idx = prog.instructions.len();
                prog.instructions.push(Instruction {
                    op: OpCode::BranchIfRegime as u8,
                    arg0: *regime as u8,
                    arg1: 0,
                    arg2: 0, // patched below
                });
                flatten_into(child, prog);
                branch_indices.push(branch_idx);
            }
            let end = prog.instructions.len() as u32;
            for branch_idx in branch_indices {
                prog.instructions[branch_idx].arg2 = end;
            }
        }
    }
}

/// Bytecode opcodes. The decision-policy kernel (C7) interprets these
/// against current alpha probs + per-horizon ISV-Kelly state to mutate
/// shared-mem orders[]. Layout MUST stay in sync with
/// `cuda/decision_policy.cu`. See spec §6.
#[repr(u8)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpCode {
    NoOp = 0,
    EvalRegime = 1,
    EmitPerHorizonSize = 2,
    AggWeightedSharpe = 3,
    AggMean = 4,
    AggMaxConfidence = 5,
    ApplyConflict = 6,
    WriteOrder = 7,
    PushScalar = 8,
    BranchIfRegime = 9,
}

#[repr(C)]
#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct Instruction {
    pub op: u8,
    pub arg0: u8,
    pub arg1: u16,
    pub arg2: u32,
}

/// Flat device-uploadable program. One per backtest cell; indexed by
/// blockIdx.x in device-global memory (NOT __constant__, see spec §6
/// "Why not __constant__").
#[derive(Clone, Debug, Default)]
pub struct Program {
    pub instructions: Vec<Instruction>,
}

impl Program {
    /// Append the top-level WriteOrder. Called once by Strategy::flatten().
    pub(crate) fn finalize(&mut self) {
        self.instructions.push(Instruction {
            op: OpCode::WriteOrder as u8,
            arg0: 0,
            arg1: 0,
            arg2: 0,
        });
    }
}

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

    #[test]
    fn default_strategy_has_one_leaf_per_horizon() {
        let s = Strategy::default_for(3);
        match s {
            Strategy::Ensemble { children, aggregator } => {
                assert_eq!(children.len(), N_HORIZONS);
                assert!(matches!(
                    aggregator,
                    EnsembleAggregator::WeightedByRealizedSharpe
                ));
                for (i, c) in children.iter().enumerate() {
                    match c {
                        Strategy::Leaf(cfg) => {
                            assert_eq!(cfg.horizon_idx as usize, i);
                            assert_eq!(cfg.max_concurrent_lots, 3);
                        }
                        _ => panic!("expected Leaf at index {i}"),
                    }
                }
            }
            _ => panic!("expected top-level Ensemble"),
        }
    }

    #[test]
    fn default_strategy_flattens_to_emits_plus_agg_plus_write() {
        let s = Strategy::default_for(3);
        let p = s.flatten();
        // N_HORIZONS EmitPerHorizonSize + 1 AggWeightedSharpe + 1 WriteOrder.
        assert_eq!(p.instructions.len(), N_HORIZONS + 2);
        for i in 0..N_HORIZONS {
            assert_eq!(p.instructions[i].op, OpCode::EmitPerHorizonSize as u8);
            assert_eq!(p.instructions[i].arg0, i as u8);
            assert_eq!(p.instructions[i].arg1, 3);
        }
        assert_eq!(p.instructions[N_HORIZONS].op, OpCode::AggWeightedSharpe as u8);
        assert_eq!(p.instructions[N_HORIZONS].arg0, N_HORIZONS as u8);
        assert_eq!(p.instructions[N_HORIZONS + 1].op, OpCode::WriteOrder as u8);
    }

    #[test]
    fn single_leaf_flattens_to_emit_plus_write() {
        let leaf = Strategy::Leaf(StrategyConfig {
            horizon_idx: 4,
            sizing_policy: SizingPolicyId::IsvKelly,
            max_concurrent_lots: 1,
        });
        let p = leaf.flatten();
        assert_eq!(p.instructions.len(), 2);
        assert_eq!(p.instructions[0].op, OpCode::EmitPerHorizonSize as u8);
        assert_eq!(p.instructions[0].arg0, 4);
        assert_eq!(p.instructions[1].op, OpCode::WriteOrder as u8);
    }

    #[test]
    fn portfolio_flattens_with_conflict_op() {
        let p = Strategy::Portfolio {
            children: vec![
                Strategy::Leaf(StrategyConfig {
                    horizon_idx: 0,
                    sizing_policy: SizingPolicyId::IsvKelly,
                    max_concurrent_lots: 1,
                }),
                Strategy::Leaf(StrategyConfig {
                    horizon_idx: 1,
                    sizing_policy: SizingPolicyId::IsvKelly,
                    max_concurrent_lots: 1,
                }),
            ],
            capital_allocation: vec![0.6, 0.4],
            conflict_policy: ConflictPolicy::NetOut,
        }
        .flatten();
        // 2 × EmitPerHorizonSize + AggMean + ApplyConflict + WriteOrder
        assert_eq!(p.instructions.len(), 5);
        assert_eq!(p.instructions[2].op, OpCode::AggMean as u8);
        assert_eq!(p.instructions[3].op, OpCode::ApplyConflict as u8);
        assert_eq!(p.instructions[3].arg0, ConflictPolicy::NetOut as u8);
        assert_eq!(p.instructions[4].op, OpCode::WriteOrder as u8);
    }

    #[test]
    fn instruction_is_8_bytes_pod() {
        assert_eq!(std::mem::size_of::<Instruction>(), 8);
    }
}
