#![allow(unsafe_code)] // Intentional unsafe for HFT performance optimizations

//! # HFT Performance Optimizations for TFT
//!
//! Ultra-low latency optimizations for Temporal Fusion Transformer
//! targeting sub-50us inference latency for high-frequency trading.

// Price imported from crate root (lib.rs)
use std::collections::HashMap;
use std::sync::{
    atomic::{AtomicU64, Ordering},
    Arc, Mutex,
};
use std::time::{Duration, Instant};

use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use tracing::{info, instrument, warn};

use super::TemporalFusionTransformer;
use crate::liquid::FixedPoint;
use common::types::Price;
use ml_core::MLError;

/// HFT-specific configuration for ultra-low latency inference
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HFTOptimizationConfig {
    // Latency targets
    pub target_latency_us: u64,
    pub max_acceptable_latency_us: u64,
    pub latency_percentile_target: f64,

    // Memory optimizations
    pub use_memory_pool: bool,
    pub pool_size_mb: usize,
    pub enable_memory_prefetching: bool,
    pub cache_line_alignment: bool,

    // Compute optimizations
    pub use_simd_vectorization: bool,
    pub enable_kernel_fusion: bool,

    // Parallelization
    pub max_threads: usize,
    pub enable_thread_pinning: bool,
    pub numa_aware: bool,

    // Caching strategies
    pub enable_attention_caching: bool,
    pub enable_computation_graph_caching: bool,
    pub cache_size_mb: usize,

    // Batch processing
    pub micro_batch_size: usize,
    pub enable_dynamic_batching: bool,
    pub batch_timeout_us: u64,

    // Hardware utilization
    pub enable_cpu_affinity: bool,
    pub preferred_cpu_cores: Vec<usize>,
    pub enable_hyperthreading: bool,
}

impl Default for HFTOptimizationConfig {
    fn default() -> Self {
        Self {
            target_latency_us: 50,
            max_acceptable_latency_us: 100,
            latency_percentile_target: 99.9,
            use_memory_pool: true,
            pool_size_mb: 128,
            enable_memory_prefetching: true,
            cache_line_alignment: true,
            use_simd_vectorization: true,
            enable_kernel_fusion: true,
            max_threads: 4,
            enable_thread_pinning: true,
            numa_aware: true,
            enable_attention_caching: true,
            enable_computation_graph_caching: true,
            cache_size_mb: 64,
            micro_batch_size: 8,
            enable_dynamic_batching: true,
            batch_timeout_us: 10,
            enable_cpu_affinity: true,
            preferred_cpu_cores: vec![0, 1, 2, 3],
            enable_hyperthreading: false,
        }
    }
}

/// Memory pool for zero-allocation inference
#[derive(Debug)]
pub struct HFTMemoryPool {
    pool: Vec<u8>,
    allocations: Mutex<HashMap<usize, (usize, usize)>>,
    next_allocation_id: AtomicU64,
    current_offset: AtomicU64,
    pool_size: usize,
}

impl HFTMemoryPool {
    pub fn new(size_mb: usize) -> Self {
        let pool_size = size_mb * 1024 * 1024;
        let pool = vec![0_u8; pool_size];

        Self {
            pool,
            allocations: Mutex::new(HashMap::new()),
            next_allocation_id: AtomicU64::new(0),
            current_offset: AtomicU64::new(0),
            pool_size,
        }
    }

    pub fn allocate(&self, size: usize, alignment: usize) -> Option<*mut u8> {
        let current = self.current_offset.load(Ordering::Relaxed);
        let aligned_offset = (current + alignment as u64 - 1) & !(alignment as u64 - 1);

        if aligned_offset + size as u64 > self.pool_size as u64 {
            warn!(
                "Memory pool exhausted: requested {}, available {}",
                size,
                self.pool_size as u64 - aligned_offset
            );
            return None;
        }

        match self.current_offset.compare_exchange(
            current,
            aligned_offset + size as u64,
            Ordering::Relaxed,
            Ordering::Relaxed,
        ) {
            Ok(_) => {
                let allocation_id = self.next_allocation_id.fetch_add(1, Ordering::Relaxed);
                if let Ok(mut allocations) = self.allocations.lock() {
                    allocations.insert(allocation_id as usize, (aligned_offset as usize, size));
                }
                // SAFETY: aligned_offset + size <= pool_size (checked above), pointer within Vec bounds
                Some(unsafe { self.pool.as_ptr().add(aligned_offset as usize) as *mut u8 })
            },
            Err(_) => self.allocate(size, alignment),
        }
    }

    pub fn reset(&self) {
        self.current_offset.store(0, Ordering::Relaxed);
        if let Ok(mut allocations) = self.allocations.lock() {
            allocations.clear();
        }
    }

    pub fn usage_bytes(&self) -> usize {
        self.current_offset.load(Ordering::Relaxed) as usize
    }

    pub fn usage_percentage(&self) -> FixedPoint {
        let percentage = (self.usage_bytes() as f64 / self.pool_size as f64) * 100.0;
        FixedPoint::from_f64(percentage)
    }
}

/// SIMD-optimized matrix operations
#[derive(Debug)]
pub struct SIMDMatrixOps;

impl SIMDMatrixOps {
    #[cfg(target_arch = "x86_64")]
    #[allow(clippy::multiple_unsafe_ops_per_block)]
    pub fn vectorized_dot_product_f32(a: &[f32], b: &[f32]) -> f32 {
        use std::arch::x86_64::{_mm256_setzero_ps, _mm256_loadu_ps, _mm256_mul_ps, _mm256_add_ps};

        assert_eq!(a.len(), b.len());
        let len = a.len();

        // SAFETY: AVX2 SIMD dot product — unaligned loads, valid lengths
        unsafe {
            let chunks = len / 8;
            let mut sum_vec = _mm256_setzero_ps();

            for i in 0..chunks {
                let offset = i * 8;
                let a_vec = _mm256_loadu_ps(a.as_ptr().add(offset));
                let b_vec = _mm256_loadu_ps(b.as_ptr().add(offset));
                let mul_vec = _mm256_mul_ps(a_vec, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, mul_vec);
            }

            let sum_array: [f32; 8] = std::mem::transmute(sum_vec);
            let mut result = sum_array.iter().sum();

            for i in (chunks * 8)..len {
                result += a[i] * b[i];
            }

            result
        }
    }

    #[cfg(not(target_arch = "x86_64"))]
    pub fn vectorized_dot_product_f32(a: &[f32], b: &[f32]) -> f32 {
        a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
    }

    pub fn optimized_matrix_multiply(
        a: &[f32],
        a_rows: usize,
        a_cols: usize,
        b: &[f32],
        b_rows: usize,
        b_cols: usize,
        result: &mut [f32],
    ) {
        assert_eq!(a_cols, b_rows);
        assert_eq!(a.len(), a_rows * a_cols);
        assert_eq!(b.len(), b_rows * b_cols);
        assert_eq!(result.len(), a_rows * b_cols);

        result
            .par_chunks_mut(b_cols)
            .enumerate()
            .for_each(|(row_idx, result_row)| {
                for col_idx in 0..b_cols {
                    let a_row_start = row_idx * a_cols;
                    let a_row = &a[a_row_start..a_row_start + a_cols];

                    let mut b_col = Vec::with_capacity(b_rows);
                    for k in 0..b_rows {
                        b_col.push(b[k * b_cols + col_idx]);
                    }

                    let sum = Self::vectorized_dot_product_f32(a_row, &b_col);
                    result_row[col_idx] = sum;
                }
            });
    }
}

/// HFT-optimized `TFT` wrapper with all performance enhancements
#[derive(Debug)]
pub struct HFTOptimizedTFT {
    pub config: HFTOptimizationConfig,
    base_model: TemporalFusionTransformer,
    memory_pool: Arc<HFTMemoryPool>,

    // Performance metrics
    latency_samples: Mutex<Vec<u64>>,
    inference_count: AtomicU64,
    cache_hits: AtomicU64,
    cache_misses: AtomicU64,
}

impl HFTOptimizedTFT {
    pub fn new(
        model: TemporalFusionTransformer,
        config: HFTOptimizationConfig,
    ) -> Result<Self, MLError> {
        info!(
            "Creating HFT-optimized TFT with target latency {}us",
            config.target_latency_us
        );

        let memory_pool = Arc::new(HFTMemoryPool::new(config.pool_size_mb));

        if config.enable_cpu_affinity {
            Self::set_cpu_affinity(&config.preferred_cpu_cores)?;
        }

        Ok(Self {
            config,
            base_model: model,
            memory_pool,
            latency_samples: Mutex::new(Vec::with_capacity(10000)),
            inference_count: AtomicU64::new(0),
            cache_hits: AtomicU64::new(0),
            cache_misses: AtomicU64::new(0),
        })
    }

    /// Ultra-fast prediction with all optimizations enabled
    #[instrument(skip(self, static_features, historical_features, future_features))]
    pub fn predict_ultra_fast(
        &mut self,
        static_features: &[f32],
        historical_features: &[f32],
        future_features: &[f32],
    ) -> Result<Vec<Price>, MLError> {
        let start_time = Instant::now();

        let f32_predictions = self.base_model.predict_fast(
            static_features,
            historical_features,
            future_features,
        )?;
        let predictions = f32_predictions
            .into_iter()
            .map(|f| {
                Price::from_f64(f as f64)
                    .map_err(|e| MLError::InvalidInput(format!("Invalid price value: {e}")))
            })
            .collect::<Result<Vec<_>, _>>()?;

        let latency = start_time.elapsed();
        self.record_latency(latency);

        if latency.as_micros() as u64 > self.config.max_acceptable_latency_us {
            warn!(
                "Inference latency {}us exceeds maximum acceptable {}us",
                latency.as_micros(),
                self.config.max_acceptable_latency_us
            );
        }

        Ok(predictions)
    }

    fn record_latency(&self, latency: Duration) {
        let latency_us = latency.as_micros() as u64;
        self.inference_count.fetch_add(1, Ordering::Relaxed);

        if let Ok(mut samples) = self.latency_samples.lock() {
            samples.push(latency_us);
            if samples.len() > 10000 {
                samples.drain(0..1000);
            }
        }
    }

    #[allow(clippy::unnecessary_wraps, clippy::multiple_unsafe_ops_per_block)]
    fn set_cpu_affinity(preferred_cores: &[usize]) -> Result<(), MLError> {
        #[cfg(target_os = "linux")]
        {
            use libc::{cpu_set_t, sched_setaffinity, CPU_SET, CPU_ZERO};
            use std::mem::MaybeUninit;

            // SAFETY: CPU affinity via libc FFI — properly initialized MaybeUninit
            unsafe {
                let mut cpu_set: MaybeUninit<cpu_set_t> = MaybeUninit::uninit();
                let cpu_set = cpu_set.as_mut_ptr();

                CPU_ZERO(&mut *cpu_set);
                for &core in preferred_cores {
                    CPU_SET(core, &mut *cpu_set);
                }

                let result = sched_setaffinity(0, size_of::<cpu_set_t>(), cpu_set);
                if result != 0 {
                    warn!(
                        "Failed to set CPU affinity: {}",
                        std::io::Error::last_os_error()
                    );
                } else {
                    info!("Set CPU affinity to cores: {:?}", preferred_cores);
                }
            }
        }

        #[cfg(not(target_os = "linux"))]
        {
            info!("CPU affinity setting not supported on this platform");
        }

        Ok(())
    }

    /// Get comprehensive performance metrics
    pub fn get_performance_metrics(&self) -> HFTPerformanceMetrics {
        let inference_count = self.inference_count.load(Ordering::Relaxed);
        let cache_hits = self.cache_hits.load(Ordering::Relaxed);
        let cache_misses = self.cache_misses.load(Ordering::Relaxed);

        let (latency_stats, latency_percentiles) = if let Ok(samples) = self.latency_samples.lock()
        {
            if samples.is_empty() {
                (LatencyStatistics::default(), LatencyPercentiles::default())
            } else {
                let mut sorted_samples = samples.clone();
                sorted_samples.sort_unstable();

                let min = *sorted_samples.first().unwrap_or(&0);
                let max = *sorted_samples.last().unwrap_or(&0);
                let mean_f64 =
                    sorted_samples.iter().sum::<u64>() as f64 / sorted_samples.len() as f64;
                let mean = FixedPoint::from_f64(mean_f64);

                let p50_idx = sorted_samples.len() / 2;
                let p95_idx = (sorted_samples.len() * 95) / 100;
                let p99_idx = (sorted_samples.len() * 99) / 100;
                let p999_idx = (sorted_samples.len() * 999) / 1000;

                let stats = LatencyStatistics {
                    min_us: min,
                    max_us: max,
                    mean_us: mean,
                    samples_count: sorted_samples.len(),
                };

                let percentiles = LatencyPercentiles {
                    p50_us: sorted_samples.get(p50_idx).copied().unwrap_or(0),
                    p95_us: sorted_samples.get(p95_idx).copied().unwrap_or(0),
                    p99_us: sorted_samples.get(p99_idx).copied().unwrap_or(0),
                    p999_us: sorted_samples.get(p999_idx).copied().unwrap_or(0),
                };

                (stats, percentiles)
            }
        } else {
            (LatencyStatistics::default(), LatencyPercentiles::default())
        };

        let cache_hit_rate = if cache_hits + cache_misses > 0 {
            FixedPoint::from_f64(cache_hits as f64 / (cache_hits + cache_misses) as f64)
        } else {
            FixedPoint::zero()
        };

        let target_compliance_rate = if latency_stats.samples_count > 0 {
            let compliant_count = if let Ok(samples) = self.latency_samples.lock() {
                samples
                    .iter()
                    .filter(|&&latency| latency <= self.config.target_latency_us)
                    .count()
            } else {
                0
            };
            FixedPoint::from_f64(compliant_count as f64 / latency_stats.samples_count as f64)
        } else {
            FixedPoint::zero()
        };

        HFTPerformanceMetrics {
            inference_count,
            latency_stats,
            latency_percentiles,
            cache_hit_rate,
            memory_pool_usage_mb: FixedPoint::from_f64(
                self.memory_pool.usage_bytes() as f64 / (1024.0 * 1024.0),
            ),
            memory_pool_usage_percent: self.memory_pool.usage_percentage(),
            attention_cache_size: 0,
            target_compliance_rate,
        }
    }
}

/// Detailed performance metrics for HFT optimization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HFTPerformanceMetrics {
    pub inference_count: u64,
    pub latency_stats: LatencyStatistics,
    pub latency_percentiles: LatencyPercentiles,
    pub cache_hit_rate: FixedPoint,
    pub memory_pool_usage_mb: FixedPoint,
    pub memory_pool_usage_percent: FixedPoint,
    pub attention_cache_size: usize,
    pub target_compliance_rate: FixedPoint,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct LatencyStatistics {
    pub min_us: u64,
    pub max_us: u64,
    pub mean_us: FixedPoint,
    pub samples_count: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct LatencyPercentiles {
    pub p50_us: u64,
    pub p95_us: u64,
    pub p99_us: u64,
    pub p999_us: u64,
}

#[cfg(test)]
#[allow(dead_code, clippy::unnecessary_wraps, clippy::redundant_clone)]
mod tests {
    use super::*;

    #[test]
    fn test_hft_config_creation() {
        let config = HFTOptimizationConfig::default();

        assert_eq!(config.target_latency_us, 50);
        assert!(config.use_memory_pool);
        assert!(config.use_simd_vectorization);
        assert!(config.enable_attention_caching);
    }

    #[test]
    fn test_simd_dot_product() -> Result<(), MLError> {
        let a = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let b = vec![2.0_f32, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0];

        let result = SIMDMatrixOps::vectorized_dot_product_f32(&a, &b);
        let expected: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();

        assert!((result - expected).abs() < 1e-6);
        Ok(())
    }
}
