rustorch 0.6.29

//! Simplified WebGPU implementation for Chrome browser
//! Chrome ブラウザ用簡素化WebGPU実装

#[cfg(feature = "webgpu")]
use wasm_bindgen::prelude::*;

#[cfg(feature = "webgpu")]
#[wasm_bindgen]
extern "C" {
    #[wasm_bindgen(js_namespace = console)]
    fn log(s: &str);
}

#[cfg(feature = "webgpu")]
macro_rules! console_log {
    ($($t:tt)*) => (log(&format_args!($($t)*).to_string()))
}

// Simplified WebGPU wrapper for basic tensor operations
#[cfg(feature = "webgpu")]
#[wasm_bindgen]
pub struct WebGPUSimple {
    initialized: bool,
}

#[cfg(feature = "webgpu")]
#[wasm_bindgen]
impl WebGPUSimple {
    #[wasm_bindgen(constructor)]
    pub fn new() -> WebGPUSimple {
        console_error_panic_hook::set_once();
        console_log!("🎮 WebGPU Simple engine created");

        WebGPUSimple { initialized: false }
    }

    #[wasm_bindgen]
    pub async fn initialize(&mut self) -> Result<String, JsValue> {
        // Check WebGPU support
        if !self.check_webgpu_support().await {
            return Err(JsValue::from_str("WebGPU not supported"));
        }

        self.initialized = true;
        let message = "✅ WebGPU initialized successfully for Chrome";
        console_log!("{}", message);
        Ok(message.to_string())
    }

    #[wasm_bindgen]
    pub async fn check_webgpu_support(&self) -> bool {
        // Use JavaScript to check WebGPU support
        let check_result = js_sys::eval(
            r#"
            (async () => {
                if (!navigator.gpu) return false;
                try {
                    const adapter = await navigator.gpu.requestAdapter();
                    return adapter !== null;
                } catch (e) {
                    return false;
                }
            })()
        "#,
        );

        match check_result {
            Ok(promise) => {
                match wasm_bindgen_futures::JsFuture::from(js_sys::Promise::from(promise)).await {
                    Ok(result) => result.as_bool().unwrap_or(false),
                    Err(_) => false,
                }
            }
            Err(_) => false,
        }
    }

    #[wasm_bindgen]
    pub fn tensor_add_cpu(&self, a: Vec<f32>, b: Vec<f32>) -> Result<Vec<f32>, JsValue> {
        if a.len() != b.len() {
            return Err(JsValue::from_str("Tensor dimensions must match"));
        }

        let result: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x + y).collect();
        Ok(result)
    }

    #[wasm_bindgen]
    pub fn tensor_mul_cpu(&self, a: Vec<f32>, b: Vec<f32>) -> Result<Vec<f32>, JsValue> {
        if a.len() != b.len() {
            return Err(JsValue::from_str("Tensor dimensions must match"));
        }

        let result: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x * y).collect();
        Ok(result)
    }

    #[wasm_bindgen]
    pub fn matrix_multiply_cpu(
        &self,
        a: Vec<f32>,
        b: Vec<f32>,
        m: u32,
        n: u32,
        k: u32,
    ) -> Result<Vec<f32>, JsValue> {
        if a.len() != (m * k) as usize || b.len() != (k * n) as usize {
            return Err(JsValue::from_str("Matrix dimensions invalid"));
        }

        let mut result = vec![0.0f32; (m * n) as usize];

        for i in 0..m {
            for j in 0..n {
                let mut sum = 0.0;
                for p in 0..k {
                    let a_idx = (i * k + p) as usize;
                    let b_idx = (p * n + j) as usize;
                    sum += a[a_idx] * b[b_idx];
                }
                result[(i * n + j) as usize] = sum;
            }
        }

        Ok(result)
    }

    #[wasm_bindgen]
    pub fn relu_cpu(&self, input: Vec<f32>) -> Vec<f32> {
        input.iter().map(|&x| x.max(0.0)).collect()
    }

    #[wasm_bindgen]
    pub fn sigmoid_cpu(&self, input: Vec<f32>) -> Vec<f32> {
        input.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect()
    }

    #[wasm_bindgen]
    pub fn get_status(&self) -> String {
        if self.initialized {
            "WebGPU Ready".to_string()
        } else {
            "Not Initialized".to_string()
        }
    }

    #[wasm_bindgen]
    pub fn get_chrome_info(&self) -> String {
        format!(
            "Chrome WebGPU Demo\nTarget: Chrome 113+\nStatus: {}\nRecommendation: Enable hardware acceleration",
            self.get_status()
        )
    }
}

// Utility functions
#[cfg(feature = "webgpu")]
#[wasm_bindgen]
pub fn get_browser_webgpu_info() -> String {
    "WebGPU Support Detection:\n1. Chrome 113+ required\n2. Enable chrome://flags/#enable-unsafe-webgpu\n3. Hardware GPU acceleration recommended".to_string()
}

#[cfg(feature = "webgpu")]
#[wasm_bindgen]
pub fn calculate_performance_estimate(operation: &str, size: u32) -> f32 {
    match operation {
        "add" | "mul" => {
            if size > 1000 {
                2.0
            } else {
                1.2
            }
        }
        "matmul" => {
            if size > 256 {
                10.0
            } else {
                1.5
            }
        }
        "relu" | "sigmoid" => {
            if size > 500 {
                3.0
            } else {
                1.5
            }
        }
        _ => 1.0,
    }
}

/// WebGPU Simple Demo struct for browser demonstration
/// ブラウザデモ用WebGPUシンプルデモ構造体
#[cfg(feature = "webgpu")]
#[wasm_bindgen]
pub struct WebGPUSimpleDemo {
    engine: Option<WebGPUSimple>,
    results: Vec<String>,
}

#[cfg(feature = "webgpu")]
#[wasm_bindgen]
impl WebGPUSimpleDemo {
    #[wasm_bindgen(constructor)]
    pub fn new() -> WebGPUSimpleDemo {
        console_log!("🚀 WebGPU Simple Demo initialized");

        WebGPUSimpleDemo {
            engine: None,
            results: Vec::new(),
        }
    }

    #[wasm_bindgen]
    pub async fn initialize(&mut self) -> Result<String, JsValue> {
        let mut engine = WebGPUSimple::new();

        match engine.initialize().await {
            Ok(message) => {
                console_log!("✅ Engine initialized: {}", message);
                self.engine = Some(engine);
                Ok(message)
            }
            Err(e) => {
                console_log!("⚠️ WebGPU failed, using CPU fallback: {:?}", e);
                self.engine = Some(engine); // Still use engine for CPU operations
                Ok("⚠️ WebGPU unavailable, using CPU fallback".to_string())
            }
        }
    }

    #[wasm_bindgen]
    pub fn run_tensor_addition_demo(&mut self) -> Result<String, JsValue> {
        let engine = self
            .engine
            .as_ref()
            .ok_or_else(|| JsValue::from_str("Engine not initialized"))?;

        console_log!("🧮 Running tensor addition demo...");

        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let b = vec![0.5, 1.5, 2.5, 3.5, 4.5];

        let result = engine.tensor_add_cpu(a.clone(), b.clone())?;
        let expected = vec![1.5, 3.5, 5.5, 7.5, 9.5];

        let message = format!(
            "✅ Tensor Addition:\n  A: {:?}\n  B: {:?}\n  Result: {:?}\n  Expected: {:?}\n  ✅ Match: {}",
            a, b, result, expected,
            result.iter().zip(expected.iter()).all(|(x, y)| (x - y).abs() < 1e-6)
        );

        console_log!("{}", message);
        self.results.push(message.clone());
        Ok(message)
    }

    #[wasm_bindgen]
    pub fn run_matrix_multiplication_demo(&mut self) -> Result<String, JsValue> {
        let engine = self
            .engine
            .as_ref()
            .ok_or_else(|| JsValue::from_str("Engine not initialized"))?;

        console_log!("🧮 Running matrix multiplication demo...");

        // 2x3 × 3x2 = 2x2
        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; // [[1,2,3], [4,5,6]]
        let b = vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0]; // [[7,8], [9,10], [11,12]]

        let result = engine.matrix_multiply_cpu(a.clone(), b.clone(), 2, 2, 3)?;
        let expected = vec![58.0, 64.0, 139.0, 154.0]; // Manually calculated

        let message = format!(
            "✅ Matrix Multiplication (2x3 × 3x2):\n  A: {:?}\n  B: {:?}\n  Result (2x2): {:?}\n  Expected: {:?}\n  ✅ Match: {}",
            a, b, result, expected,
            result.iter().zip(expected.iter()).all(|(x, y)| (x - y).abs() < 1e-6)
        );

        console_log!("{}", message);
        self.results.push(message.clone());
        Ok(message)
    }

    #[wasm_bindgen]
    pub fn run_activation_functions_demo(&mut self) -> Result<String, JsValue> {
        let engine = self
            .engine
            .as_ref()
            .ok_or_else(|| JsValue::from_str("Engine not initialized"))?;

        console_log!("🧮 Running activation functions demo...");

        let input = vec![-2.0, -1.0, 0.0, 1.0, 2.0];

        let relu_result = engine.relu_cpu(input.clone());
        let sigmoid_result = engine.sigmoid_cpu(input.clone());

        let expected_relu = vec![0.0, 0.0, 0.0, 1.0, 2.0];
        let expected_sigmoid: Vec<f32> = input.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();

        let message = format!(
            "✅ Activation Functions:\n  Input: {:?}\n  ReLU: {:?}\n  Expected ReLU: {:?}\n  Sigmoid: {:?}\n  ✅ ReLU Match: {}\n  ✅ Sigmoid Match: {}",
            input, relu_result, expected_relu, sigmoid_result,
            relu_result.iter().zip(expected_relu.iter()).all(|(x, y)| (x - y).abs() < 1e-6),
            sigmoid_result.iter().zip(expected_sigmoid.iter()).all(|(x, y)| (x - y).abs() < 1e-6)
        );

        console_log!("{}", message);
        self.results.push(message.clone());
        Ok(message)
    }

    #[wasm_bindgen]
    pub fn run_performance_benchmark(&mut self) -> Result<String, JsValue> {
        console_log!("📊 Running performance benchmark...");

        let test_sizes = vec![100, 1000, 10000];
        let mut benchmark_results = Vec::new();

        for size in test_sizes {
            console_log!("🔬 Testing size: {}", size);

            let estimate_add = calculate_performance_estimate("add", size as u32);
            let estimate_matmul = calculate_performance_estimate("matmul", size as u32);

            let result = format!(
                "  Size {}: Add {}x, MatMul {}x speedup estimate",
                size, estimate_add, estimate_matmul
            );
            console_log!("{}", result);
            benchmark_results.push(result);
        }

        let message = format!(
            "📊 Performance Estimates (WebGPU vs CPU):\n{}",
            benchmark_results.join("\n")
        );

        console_log!("{}", message);
        self.results.push(message.clone());
        Ok(message)
    }

    #[wasm_bindgen]
    pub async fn run_comprehensive_demo(&mut self) -> Result<String, JsValue> {
        console_log!("🎯 Starting comprehensive WebGPU demo...");

        let init_result = self.initialize().await?;
        let add_result = self.run_tensor_addition_demo()?;
        let matmul_result = self.run_matrix_multiplication_demo()?;
        let activation_result = self.run_activation_functions_demo()?;
        let benchmark_result = self.run_performance_benchmark()?;

        let info = self.engine.as_ref().unwrap().get_chrome_info();
        let browser_info = get_browser_webgpu_info();

        let summary = format!(
            "🎉 WebGPU Chrome Demo Complete!\n\n{}\n\n{}\n\n{}\n\n{}\n\n{}\n\n{}\n\n{}",
            info,
            browser_info,
            init_result,
            add_result,
            matmul_result,
            activation_result,
            benchmark_result
        );

        console_log!("{}", summary);
        Ok(summary)
    }

    #[wasm_bindgen]
    pub fn get_all_results(&self) -> Vec<String> {
        self.results.clone()
    }

    #[wasm_bindgen]
    pub fn cleanup(&mut self) {
        self.results.clear();
        self.engine = None;
        console_log!("🧹 Demo cleanup completed");
    }
}