//! Multi-timeframe LSTM encoder and feature fusion.
//!
//! Resamples 1-minute OHLCV bars into 5m / 15m / 1h bars, encodes each
//! timeframe independently with a single-layer LSTM, and fuses the four
//! hidden-state embeddings through a linear projection.
//!
//! Architecture:
//! ```text
//! 1m bars ──> [LSTM 6->64] ──> emb_1m  (64)  ─┐
//!   ├─ resample 5m  ──> [LSTM 6->64] ──> emb_5m  (64)  ─┤
//!   ├─ resample 15m ──> [LSTM 6->64] ──> emb_15m (64) ─┼─> [Concat 256] ─> [Linear 128] ─> macro_state
//!   └─ resample 1h  ──> [LSTM 6->64] ──> emb_1h  (64)  ─┘
//! ```

use std::collections::VecDeque;
use std::sync::Arc;

use ml_core::cuda_autograd::GpuVarStore;
use ml_core::device::MlDevice;
use ml_core::native_types::NativeDevice;

use cudarc::driver::CudaStream;
use ml_supervised::gpu_tensor::{
    GpuLinear, GpuTensor, gpu_sigmoid, gpu_tanh, gpu_mul, gpu_add,
    gpu_cat_dim1, gpu_narrow_2d, gpu_matmul, gpu_transpose,
};

use super::bar_resampler::BarResampler;
use crate::types::OHLCVBar;
use crate::MLError;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/// Configuration for the multi-timeframe encoder.
#[derive(Debug, Clone)]
pub struct MultiTimeframeConfig {
    /// Number of input features per bar (default: 6 = O/H/L/C/V/returns).
    pub input_dim: usize,
    /// Hidden dimension of each per-timeframe LSTM (default: 64).
    pub hidden_dim: usize,
    /// Output dimension after fusion projection (default: 128).
    pub output_dim: usize,
    /// Number of recent bars to keep per timeframe (default: 20).
    pub history_len: usize,
}

impl Default for MultiTimeframeConfig {
    fn default() -> Self {
        Self {
            input_dim: 6,
            hidden_dim: 64,
            output_dim: 128,
            history_len: 20,
        }
    }
}

// ---------------------------------------------------------------------------
// LstmEncoder -- single-layer LSTM cell (manual implementation)
// ---------------------------------------------------------------------------

/// A single-layer LSTM encoder that processes a sequence and returns the
/// final hidden state as an embedding vector.
///
/// Uses cuBLAS-backed matmul via GpuTensor for gate computations.
///
/// LSTM equations:
/// ```text
/// gates = W_ih * x_t + b_ih + W_hh * h_{t-1} + b_hh
/// (i, f, g, o) = split(gates, 4)
/// i_t = sigmoid(i),  f_t = sigmoid(f),  g_t = tanh(g),  o_t = sigmoid(o)
/// c_t = f_t * c_{t-1} + i_t * g_t
/// h_t = o_t * tanh(c_t)
/// ```
#[derive(Debug)]
pub struct LstmEncoder {
    /// Input-to-hidden weights [4*hidden, input] — GpuTensor.
    w_ih: GpuTensor,
    /// Hidden-to-hidden weights [4*hidden, hidden] — GpuTensor.
    w_hh: GpuTensor,
    /// Input-to-hidden bias [4*hidden] — GpuTensor.
    b_ih: GpuTensor,
    /// Hidden-to-hidden bias [4*hidden] — GpuTensor.
    b_hh: GpuTensor,
    hidden_dim: usize,
    stream: Arc<CudaStream>,
}

impl LstmEncoder {
    /// Build a new LSTM encoder with Xavier-initialized weights on the CUDA stream.
    pub fn new(input_dim: usize, hidden_dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        use rand::Rng;
        let gate_dim = 4 * hidden_dim;
        let limit_ih = (6.0 / (input_dim + hidden_dim) as f64).sqrt() as f32;
        let limit_hh = (6.0 / (hidden_dim + hidden_dim) as f64).sqrt() as f32;

        let mut rng = rand::thread_rng();
        let w_ih_data: Vec<f32> = (0..gate_dim * input_dim).map(|_| rng.gen_range(-limit_ih..limit_ih)).collect();
        let w_hh_data: Vec<f32> = (0..gate_dim * hidden_dim).map(|_| rng.gen_range(-limit_hh..limit_hh)).collect();

        let w_ih = GpuTensor::from_vec(w_ih_data, &[gate_dim, input_dim], stream)?;
        let w_hh = GpuTensor::from_vec(w_hh_data, &[gate_dim, hidden_dim], stream)?;
        let b_ih = GpuTensor::zeros(&[1, gate_dim], stream)?;
        let b_hh = GpuTensor::zeros(&[1, gate_dim], stream)?;

        Ok(Self {
            w_ih,
            w_hh,
            b_ih,
            b_hh,
            hidden_dim,
            stream: Arc::clone(stream),
        })
    }

    /// Run the LSTM over a sequence and return the final hidden state.
    ///
    /// * `seq` -- GpuTensor of shape `(seq_len, input_dim)`
    ///
    /// Returns a GpuTensor of shape `(1, hidden_dim)` (the final h).
    pub fn forward(&self, seq: &GpuTensor) -> Result<GpuTensor, MLError> {
        let seq_len = seq.dim(0)?;

        let mut h = GpuTensor::zeros(&[1, self.hidden_dim], &self.stream)?;
        let mut c = GpuTensor::zeros(&[1, self.hidden_dim], &self.stream)?;

        // Transpose weight matrices once for efficient matmul: W^T
        let w_ih_t = gpu_transpose(&self.w_ih)?;
        let w_hh_t = gpu_transpose(&self.w_hh)?;

        for t in 0..seq_len {
            // x_t: (1, input_dim) -- narrow row t
            let x_t = gpu_narrow_2d(seq, 0, t, 1)?;

            // gates = x_t @ W_ih^T + h @ W_hh^T + b_ih + b_hh
            let xw = gpu_matmul(&x_t, &w_ih_t)?;
            let hw = gpu_matmul(&h, &w_hh_t)?;
            let gates = gpu_add(&gpu_add(&gpu_add(&xw, &self.b_ih)?, &hw)?, &self.b_hh)?;

            let hd = self.hidden_dim;

            // Split gates into 4 chunks along dim 1
            let i_gate = gpu_narrow_2d(&gates, 1, 0, hd)?;
            let f_gate = gpu_narrow_2d(&gates, 1, hd, hd)?;
            let g_gate = gpu_narrow_2d(&gates, 1, 2 * hd, hd)?;
            let o_gate = gpu_narrow_2d(&gates, 1, 3 * hd, hd)?;

            let i_sig = gpu_sigmoid(&i_gate)?;
            let f_sig = gpu_sigmoid(&f_gate)?;
            let g_tanh = gpu_tanh(&g_gate)?;
            let o_sig = gpu_sigmoid(&o_gate)?;

            // c_t = f_t * c_{t-1} + i_t * g_t
            let fc = gpu_mul(&f_sig, &c)?;
            let ig = gpu_mul(&i_sig, &g_tanh)?;
            c = gpu_add(&fc, &ig)?;

            // h_t = o_t * tanh(c_t)
            let c_tanh = gpu_tanh(&c)?;
            h = gpu_mul(&o_sig, &c_tanh)?;
        }

        Ok(h)
    }
}

// ---------------------------------------------------------------------------
// MultiTimeframeEncoder
// ---------------------------------------------------------------------------

/// Encoder that processes OHLCV bars at four timeframes (1m, 5m, 15m, 1h)
/// and fuses the per-timeframe LSTM embeddings into a single macro-state vector.
///
/// All compute uses cuBLAS-backed GpuTensor. Output is converted to Candle Tensor
/// at the API boundary.
#[derive(Debug)]
pub struct MultiTimeframeEncoder {
    lstm_1m: LstmEncoder,
    lstm_5m: LstmEncoder,
    lstm_15m: LstmEncoder,
    lstm_1h: LstmEncoder,
    projection: GpuLinear,
    resampler: BarResampler,
    history_1m: VecDeque<OHLCVBar>,
    history_5m: VecDeque<OHLCVBar>,
    history_15m: VecDeque<OHLCVBar>,
    history_1h: VecDeque<OHLCVBar>,
    config: MultiTimeframeConfig,
    device: NativeDevice,
    stream: Arc<CudaStream>,
}

impl MultiTimeframeEncoder {
    /// Build a new encoder with GPU-native LSTM encoders and cuBLAS projection.
    pub fn new(config: MultiTimeframeConfig, stream: &Arc<CudaStream>, device: &NativeDevice) -> Result<Self, MLError> {
        let concat_dim = config.hidden_dim * 4; // 4 timeframes

        let lstm_1m = LstmEncoder::new(config.input_dim, config.hidden_dim, stream)?;
        let lstm_5m = LstmEncoder::new(config.input_dim, config.hidden_dim, stream)?;
        let lstm_15m = LstmEncoder::new(config.input_dim, config.hidden_dim, stream)?;
        let lstm_1h = LstmEncoder::new(config.input_dim, config.hidden_dim, stream)?;

        let projection = GpuLinear::new(concat_dim, config.output_dim, stream)?;

        let history_len = config.history_len;

        Ok(Self {
            lstm_1m,
            lstm_5m,
            lstm_15m,
            lstm_1h,
            projection,
            resampler: BarResampler::new(),
            history_1m: VecDeque::with_capacity(history_len),
            history_5m: VecDeque::with_capacity(history_len),
            history_15m: VecDeque::with_capacity(history_len),
            history_1h: VecDeque::with_capacity(history_len),
            config,
            device: device.clone(),
            stream: Arc::clone(stream),
        })
    }

    /// Build with a fresh `GpuVarStore` on the specified device (convenience).
    pub fn with_device(
        config: MultiTimeframeConfig,
        device: &NativeDevice,
    ) -> Result<(Self, GpuVarStore), MLError> {
        let ordinal = device.cuda_ordinal().ok_or_else(|| {
            MLError::ConfigError("MultiTimeframeEncoder requires CUDA device".to_owned())
        })?;
        let ml_dev = MlDevice::cuda(ordinal)?;
        let stream = ml_dev.cuda_stream()?;
        let vars = GpuVarStore::new(Arc::clone(stream));
        let encoder = Self::new(config, stream, device)?;
        Ok((encoder, vars))
    }

    /// Ingest a 1-minute bar, update internal ring buffers, and return
    /// the fused macro-state embedding of shape `(1, output_dim)`.
    ///
    /// The resampler converts the 1m bar into higher-timeframe bars when
    /// enough constituent bars have been collected. Ring buffers are capped
    /// at `history_len`.
    pub fn push_bar(&mut self, bar: OHLCVBar) -> Result<GpuTensor, MLError> {
        // Update resampler
        let (bar_5m, bar_15m, bar_1h) = self.resampler.push(bar);

        // Update ring buffers
        push_ring(&mut self.history_1m, bar, self.config.history_len);
        if let Some(b) = bar_5m {
            push_ring(&mut self.history_5m, b, self.config.history_len);
        }
        if let Some(b) = bar_15m {
            push_ring(&mut self.history_15m, b, self.config.history_len);
        }
        if let Some(b) = bar_1h {
            push_ring(&mut self.history_1h, b, self.config.history_len);
        }

        self.encode()
    }

    /// Encode the current ring-buffer contents and return the fused state.
    ///
    /// For any timeframe with no history yet, a zero embedding is used.
    /// Returns a Candle Tensor at the API boundary.
    pub fn encode(&self) -> Result<GpuTensor, MLError> {
        let emb_1m = self.encode_timeframe(&self.lstm_1m, &self.history_1m)?;
        let emb_5m = self.encode_timeframe(&self.lstm_5m, &self.history_5m)?;
        let emb_15m = self.encode_timeframe(&self.lstm_15m, &self.history_15m)?;
        let emb_1h = self.encode_timeframe(&self.lstm_1h, &self.history_1h)?;

        // Concat along feature dim: (1, 4*hidden_dim) using GpuTensor ops
        let cat_12 = gpu_cat_dim1(&emb_1m, &emb_5m)?;
        let cat_123 = gpu_cat_dim1(&cat_12, &emb_15m)?;
        let concat = gpu_cat_dim1(&cat_123, &emb_1h)?;

        // Project to output_dim (cuBLAS sgemm)
        self.projection.forward(&concat)
    }

    /// Reset the resampler and all history buffers.
    pub fn reset(&mut self) {
        self.resampler.reset();
        self.history_1m.clear();
        self.history_5m.clear();
        self.history_15m.clear();
        self.history_1h.clear();
    }

    /// Return a reference to the config.
    pub fn config(&self) -> &MultiTimeframeConfig {
        &self.config
    }

    // -----------------------------------------------------------------------
    // Internal helpers
    // -----------------------------------------------------------------------

    /// Encode a single timeframe's history through the given LSTM.
    /// Returns `(1, hidden_dim)` GpuTensor.
    fn encode_timeframe(
        &self,
        lstm: &LstmEncoder,
        history: &VecDeque<OHLCVBar>,
    ) -> Result<GpuTensor, MLError> {
        if history.is_empty() {
            // No data yet -- return zeros
            return GpuTensor::zeros(&[1, self.config.hidden_dim], &self.stream);
        }

        let seq = bars_to_gpu_tensor(history, &self.stream)?;
        lstm.forward(&seq)
    }
}

// ---------------------------------------------------------------------------
// Utility functions
// ---------------------------------------------------------------------------

/// Convert an OHLCV bar to a 6-dim feature vector:
/// `[open, high, low, close, volume, returns]`.
///
/// `prev_close` is used to compute returns (`(close - prev_close) / prev_close`).
/// If `prev_close` is `None` or zero, returns is set to 0.0.
pub fn bar_to_features(bar: &OHLCVBar, prev_close: Option<f64>) -> [f64; 6] {
    let returns = match prev_close {
        Some(pc) if pc.abs() > f64::EPSILON => (bar.close - pc) / pc,
        _ => 0.0,
    };
    [bar.open, bar.high, bar.low, bar.close, bar.volume, returns]
}

/// Convert a sequence of OHLCV bars to a `(seq_len, 6)` GpuTensor.
fn bars_to_gpu_tensor(bars: &VecDeque<OHLCVBar>, stream: &Arc<CudaStream>) -> Result<GpuTensor, MLError> {
    let len = bars.len();
    let mut data = Vec::with_capacity(len * 6);

    let mut prev_close: Option<f64> = None;
    for bar in bars {
        let feats = bar_to_features(bar, prev_close);
        for &f in &feats {
            data.push(f as f32);
        }
        prev_close = Some(bar.close);
    }

    GpuTensor::from_vec(data, &[len, 6], stream)
}

/// Push a bar into a ring buffer, popping the oldest if at capacity.
fn push_ring(buf: &mut VecDeque<OHLCVBar>, bar: OHLCVBar, max_len: usize) {
    if buf.len() >= max_len {
        buf.pop_front();
    }
    buf.push_back(bar);
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::{TimeZone, Utc};
    use ml_core::cuda_autograd::stream_ops::{gpu_sub, gpu_abs};

    fn make_bar(minute: u32, close: f64) -> OHLCVBar {
        OHLCVBar {
            timestamp: Utc
                .with_ymd_and_hms(2026, 1, 1, 10, minute % 60, 0)
                .single()
                .unwrap_or_else(Utc::now),
            open: close - 1.0,
            high: close + 2.0,
            low: close - 2.0,
            close,
            volume: 1000.0,
        }
    }

    fn make_config() -> MultiTimeframeConfig {
        MultiTimeframeConfig {
            input_dim: 6,
            hidden_dim: 64,
            output_dim: 128,
            history_len: 20,
        }
    }

    #[test]
    fn test_encoder_output_shape() {
        let config = make_config();
        let (mut encoder, _vars) =
            MultiTimeframeEncoder::with_device(config.clone(), &NativeDevice::Cuda(0))
                .expect("encoder creation should succeed");

        // Push some bars
        for i in 0..10 {
            let result = encoder.push_bar(make_bar(i, 100.0 + f64::from(i)));
            assert!(result.is_ok(), "push_bar should succeed: {:?}", result.err());
        }

        let output = encoder.encode().expect("encode should succeed");
        let dims = &output.shape;
        assert_eq!(dims.len(), 2, "output should be 2D");
        assert_eq!(
            dims.first().copied().unwrap_or(0),
            1,
            "batch dim should be 1"
        );
        assert_eq!(
            dims.last().copied().unwrap_or(0),
            config.output_dim,
            "feature dim should be output_dim={}",
            config.output_dim
        );
    }

    #[test]
    fn test_encoder_deterministic() {
        let config = make_config();
        let (mut encoder, _vars) =
            MultiTimeframeEncoder::with_device(config, &NativeDevice::Cuda(0))
                .expect("encoder creation should succeed");

        // Push bars
        for i in 0..5 {
            encoder
                .push_bar(make_bar(i, 100.0 + f64::from(i)))
                .expect("push should work");
        }

        let out1 = encoder.encode().expect("encode 1");
        let out2 = encoder.encode().expect("encode 2");

        // Same state -> same output: compute |out1 - out2| and sum on host
        let diff_tensor = gpu_sub(&out1, &out2).expect("sub");
        let abs_tensor = gpu_abs(&diff_tensor).expect("abs");
        let abs_vec = abs_tensor.to_vec().expect("to_vec"); // test-only readback
        let diff: f32 = abs_vec.iter().sum();

        assert!(
            diff < 1e-6,
            "same input should produce same output, diff={}",
            diff
        );
    }

    #[test]
    fn test_config_defaults() {
        let config = MultiTimeframeConfig::default();
        assert_eq!(config.input_dim, 6);
        assert_eq!(config.hidden_dim, 64);
        assert_eq!(config.output_dim, 128);
        assert_eq!(config.history_len, 20);
    }

    #[test]
    fn test_bar_to_features() {
        let bar = OHLCVBar {
            timestamp: Utc
                .with_ymd_and_hms(2026, 1, 1, 10, 0, 0)
                .single()
                .unwrap_or_else(Utc::now),
            open: 100.0,
            high: 105.0,
            low: 95.0,
            close: 102.0,
            volume: 5000.0,
        };

        // No previous close -> returns = 0
        let feats = bar_to_features(&bar, None);
        assert!((feats[0] - 100.0).abs() < f64::EPSILON, "open");
        assert!((feats[1] - 105.0).abs() < f64::EPSILON, "high");
        assert!((feats[2] - 95.0).abs() < f64::EPSILON, "low");
        assert!((feats[3] - 102.0).abs() < f64::EPSILON, "close");
        assert!((feats[4] - 5000.0).abs() < f64::EPSILON, "volume");
        assert!((feats[5] - 0.0).abs() < f64::EPSILON, "returns with no prev");

        // With previous close: returns = (102 - 100) / 100 = 0.02
        let feats = bar_to_features(&bar, Some(100.0));
        assert!(
            (feats[5] - 0.02).abs() < 1e-10,
            "returns should be 0.02, got {}",
            feats[5]
        );
    }

    #[test]
    fn test_bar_to_features_zero_prev_close() {
        let bar = make_bar(0, 102.0);
        let feats = bar_to_features(&bar, Some(0.0));
        assert!(
            feats[5].abs() < f64::EPSILON,
            "returns should be 0.0 when prev_close is 0"
        );
    }

    #[test]
    fn test_empty_history_produces_output() {
        // Even with no bars pushed, encode() should succeed (zero embeddings)
        let config = make_config();
        let (encoder, _vars) =
            MultiTimeframeEncoder::with_device(config.clone(), &NativeDevice::Cuda(0))
                .expect("encoder creation should succeed");

        let output = encoder.encode().expect("encode on empty history should work");
        let dims = &output.shape;
        assert_eq!(dims.last().copied().unwrap_or(0), config.output_dim);
    }

    #[test]
    fn test_lstm_encoder_single_step() {
        let ml_dev = MlDevice::cuda(0).expect("CUDA required");
        let stream = std::sync::Arc::clone(ml_dev.cuda_stream().expect("stream"));
        let lstm = LstmEncoder::new(6, 32, &stream).expect("lstm creation");

        // Single timestep: (1, 6)
        let input = GpuTensor::randn(&[1, 6], 1.0, &stream).expect("input tensor");
        let out = lstm.forward(&input).expect("lstm forward");
        assert_eq!(out.shape.len(), 2);
        assert_eq!(out.dim(0).unwrap_or(0), 1);
        assert_eq!(out.dim(1).unwrap_or(0), 32);
    }

    #[test]
    fn test_lstm_encoder_multi_step() {
        let ml_dev = MlDevice::cuda(0).expect("CUDA required");
        let stream = std::sync::Arc::clone(ml_dev.cuda_stream().expect("stream"));
        let lstm = LstmEncoder::new(6, 64, &stream).expect("lstm creation");

        // 10 timesteps: (10, 6)
        let input = GpuTensor::randn(&[10, 6], 1.0, &stream).expect("input tensor");
        let out = lstm.forward(&input).expect("lstm forward");
        assert_eq!(out.dim(1).unwrap_or(0), 64);
    }

    #[test]
    fn test_push_ring_eviction() {
        let mut buf = VecDeque::new();
        let max_len = 3;

        for i in 0..5 {
            push_ring(&mut buf, make_bar(i, 100.0 + f64::from(i)), max_len);
        }

        assert_eq!(buf.len(), 3, "ring buffer should cap at max_len");
        // Should contain bars for minutes 2, 3, 4
        let front = buf.front().expect("front exists");
        assert!((front.close - 102.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_reset_clears_state() {
        let config = make_config();
        let (mut encoder, _vars) =
            MultiTimeframeEncoder::with_device(config, &NativeDevice::Cuda(0))
                .expect("encoder creation");

        for i in 0..10 {
            encoder
                .push_bar(make_bar(i, 100.0 + f64::from(i)))
                .expect("push");
        }

        encoder.reset();

        // After reset, encode should return zero-based output
        let output = encoder.encode().expect("encode after reset");
        let abs_out = gpu_abs(&output).expect("abs");
        let vals = abs_out.to_vec().expect("to_vec"); // test-only readback
        let sum: f32 = vals.iter().sum();

        // With all-zero embeddings going through the projection (which has bias),
        // the output is just the projection bias. That's fine.
        assert!(sum.is_finite(), "output should be finite after reset");
    }
}
