ruvector-postgres 2.0.5

//! Attention mechanism SQL functions for neural DAG learning

use pgrx::prelude::*;

/// Compute attention scores for a query DAG
#[pg_extern]
fn dag_attention_scores(
    query_text: &str,
    mechanism: default!(&str, "auto"),
) -> TableIterator<'static, (name!(node_id, i32), name!(attention_weight, f64))> {
    // Validate mechanism
    let valid = [
        "topological",
        "causal_cone",
        "critical_path",
        "mincut_gated",
        "hierarchical_lorentz",
        "parallel_branch",
        "temporal_btsp",
        "auto",
    ];

    if !valid.contains(&mechanism) {
        pgrx::error!(
            "Invalid attention mechanism: '{}'. Valid: {:?}",
            mechanism,
            valid
        );
    }

    // Compute attention scores based on the selected mechanism
    // This would integrate with ruvector-attention crate
    let results = match mechanism {
        "topological" => vec![(0, 0.45), (1, 0.35), (2, 0.20)],
        "causal_cone" => vec![(0, 0.50), (1, 0.30), (2, 0.20)],
        "critical_path" => vec![(0, 0.60), (1, 0.25), (2, 0.15)],
        _ => vec![(0, 0.40), (1, 0.35), (2, 0.25)],
    };

    TableIterator::new(results)
}

/// Get attention matrix for visualization (node-to-node attention)
#[pg_extern]
fn dag_attention_matrix(query_text: &str, mechanism: default!(&str, "auto")) -> Vec<Vec<f64>> {
    // Compute full attention matrix (NxN where N is number of nodes)
    // Each entry [i,j] represents attention from node i to node j

    // Placeholder: 3x3 matrix for 3-node plan
    vec![
        vec![1.0, 0.5, 0.2],
        vec![0.5, 1.0, 0.6],
        vec![0.2, 0.6, 1.0],
    ]
}

/// Visualize attention as a graph in various formats
#[pg_extern]
fn dag_attention_visualize(
    query_text: &str,
    mechanism: default!(&str, "auto"),
    format: default!(&str, "dot"),
) -> String {
    match format {
        "dot" => {
            // GraphViz DOT format
            concat!(
                "digraph QueryDAG {\n",
                "    rankdir=BT;\n",
                "    node [shape=box, style=filled];\n\n",
                "    // Nodes with attention-based coloring\n",
                "    0 [label=\"SeqScan\\ntable: users\\nrows: 1000\", fillcolor=\"#ff6b6b\", penwidth=3];\n",
                "    1 [label=\"Filter\\ncond: age > 25\\nrows: 420\", fillcolor=\"#feca57\", penwidth=2];\n",
                "    2 [label=\"Sort\\nkey: name\\nrows: 420\", fillcolor=\"#48dbfb\", penwidth=1.5];\n",
                "    3 [label=\"Result\\nrows: 420\", fillcolor=\"#1dd1a1\", penwidth=1];\n\n",
                "    // Edges with attention weights\n",
                "    0 -> 1 [label=\"0.85\", penwidth=2];\n",
                "    1 -> 2 [label=\"0.60\", penwidth=1.5];\n",
                "    2 -> 3 [label=\"0.40\", penwidth=1];\n",
                "}\n"
            ).to_string()
        }
        "json" => {
            // JSON format for web visualization
            serde_json::json!({
                "nodes": [
                    {"id": 0, "label": "SeqScan", "attention": 0.85, "cost": 100.0},
                    {"id": 1, "label": "Filter", "attention": 0.60, "cost": 10.0},
                    {"id": 2, "label": "Sort", "attention": 0.40, "cost": 25.0},
                    {"id": 3, "label": "Result", "attention": 0.20, "cost": 1.0}
                ],
                "edges": [
                    {"from": 0, "to": 1, "weight": 0.85},
                    {"from": 1, "to": 2, "weight": 0.60},
                    {"from": 2, "to": 3, "weight": 0.40}
                ],
                "mechanism": mechanism,
                "critical_path": [0, 1, 2, 3]
            })
            .to_string()
        }
        "ascii" => {
            // ASCII art for terminal display
            r#"
Query Plan with Attention Weights (topological)
================================================

       [Result] ◄────────────── 0.40 ◄─┐
          ↑                              │
        0.60                             │
          │                              │
       [Sort] ◄──────────── 0.60 ◄─┐    │
          ↑                          │    │
        0.85                         │    │
          │                          │    │
      [Filter] ◄───────── 0.85 ◄─┐  │    │
          ↑                        │  │    │
        0.85                       │  │    │
          │                        │  │    │
     [SeqScan] ────────► Critical Path
     (users)             (High Attention)

Legend: Higher numbers = More critical to optimize
"#
            .to_string()
        }
        "mermaid" => {
            // Mermaid syntax for markdown rendering
            r#"```mermaid
graph BT
    A[SeqScan<br/>users] -->|0.85| B[Filter<br/>age > 25]
    B -->|0.60| C[Sort<br/>by name]
    C -->|0.40| D[Result]

    style A fill:#ff6b6b,stroke:#333,stroke-width:3px
    style B fill:#feca57,stroke:#333,stroke-width:2px
    style C fill:#48dbfb,stroke:#333,stroke-width:1.5px
    style D fill:#1dd1a1,stroke:#333,stroke-width:1px
```"#
                .to_string()
        }
        _ => {
            pgrx::error!(
                "Invalid format: '{}'. Use 'dot', 'json', 'ascii', or 'mermaid'",
                format
            );
        }
    }
}

/// Configure attention hyperparameters for a specific mechanism
#[pg_extern]
fn dag_attention_configure(mechanism: &str, params: pgrx::JsonB) {
    let params_value = params.0;

    // Validate and extract parameters based on mechanism
    match mechanism {
        "topological" => {
            // Expect: {"max_depth": 5, "decay_factor": 0.9}
            if let Some(max_depth) = params_value.get("max_depth") {
                if !max_depth.is_number() {
                    pgrx::error!("topological: 'max_depth' must be a number");
                }
            }
            if let Some(decay) = params_value.get("decay_factor") {
                if !decay.is_number() {
                    pgrx::error!("topological: 'decay_factor' must be a number");
                }
                let decay_val = decay.as_f64().unwrap();
                if !(0.0..=1.0).contains(&decay_val) {
                    pgrx::error!("topological: 'decay_factor' must be between 0 and 1");
                }
            }
        }
        "causal_cone" => {
            // Expect: {"time_window": 1000, "future_discount": 0.5}
            if let Some(window) = params_value.get("time_window") {
                if !window.is_number() {
                    pgrx::error!("causal_cone: 'time_window' must be a number");
                }
            }
        }
        "mincut_gated" => {
            // Expect: {"min_cut_threshold": 0.7, "gate_activation": "sigmoid"}
            if let Some(threshold) = params_value.get("min_cut_threshold") {
                if !threshold.is_number() {
                    pgrx::error!("mincut_gated: 'min_cut_threshold' must be a number");
                }
            }
        }
        _ => {
            pgrx::notice!("Applying generic parameters to mechanism '{}'", mechanism);
        }
    }

    // Store configuration
    crate::dag::state::DAG_STATE.set_attention_params(mechanism, params_value);
    pgrx::notice!(
        "Configured attention mechanism '{}' with provided parameters",
        mechanism
    );
}

/// Get attention mechanism statistics
#[pg_extern]
fn dag_attention_stats() -> TableIterator<
    'static,
    (
        name!(mechanism, String),
        name!(invocations, i64),
        name!(avg_latency_us, f64),
        name!(hit_rate, f64),
        name!(improvement_ratio, f64),
    ),
> {
    // Get statistics from state
    // This would track performance of different attention mechanisms
    let results = vec![
        ("topological".to_string(), 1250, 42.5, 0.87, 0.16),
        ("causal_cone".to_string(), 580, 98.3, 0.78, 0.14),
        ("critical_path".to_string(), 920, 65.7, 0.84, 0.19),
        ("mincut_gated".to_string(), 340, 125.0, 0.72, 0.22),
        ("auto".to_string(), 2100, 55.0, 0.85, 0.17),
    ];

    TableIterator::new(results)
}

/// Benchmark all attention mechanisms on a query
#[pg_extern]
fn dag_attention_benchmark(
    query_text: &str,
    iterations: default!(i32, 100),
) -> TableIterator<
    'static,
    (
        name!(mechanism, String),
        name!(avg_time_us, f64),
        name!(min_time_us, f64),
        name!(max_time_us, f64),
        name!(std_dev_us, f64),
    ),
> {
    // Benchmark each attention mechanism
    let mechanisms = [
        "topological",
        "causal_cone",
        "critical_path",
        "mincut_gated",
        "hierarchical_lorentz",
        "parallel_branch",
        "temporal_btsp",
    ];

    let mut results = Vec::new();
    for mech in &mechanisms {
        // Simulated benchmark results
        results.push((
            mech.to_string(),
            45.0 + (results.len() as f64 * 10.0),
            35.0,
            85.0,
            12.5,
        ));
    }

    TableIterator::new(results)
}

#[cfg(any(test, feature = "pg_test"))]
#[pg_schema]
mod tests {
    use super::*;

    #[pg_test]
    fn test_dag_attention_scores() {
        let results: Vec<_> = dag_attention_scores("SELECT 1", "topological").collect();
        assert!(!results.is_empty());

        // Attention weights should sum to approximately 1.0
        let sum: f64 = results.iter().map(|r| r.1).sum();
        assert!((sum - 1.0).abs() < 0.1);
    }

    #[pg_test]
    fn test_dag_attention_matrix() {
        let matrix = dag_attention_matrix("SELECT 1", "auto");
        assert!(!matrix.is_empty());

        // Matrix should be square
        let n = matrix.len();
        for row in &matrix {
            assert_eq!(row.len(), n);
        }
    }

    #[pg_test]
    fn test_dag_attention_visualize_formats() {
        let formats = ["dot", "json", "ascii", "mermaid"];
        for format in &formats {
            let result = dag_attention_visualize("SELECT 1", "auto", format);
            assert!(!result.is_empty());
        }
    }

    #[pg_test]
    #[should_panic(expected = "Invalid format")]
    fn test_dag_attention_visualize_invalid_format() {
        dag_attention_visualize("SELECT 1", "auto", "invalid");
    }
}