//! Analysis commands (artifact detection, GNN, statistics).

use std::collections::HashMap;
use tauri::State;

use crate::neuro_commands::{
    NeuroError, ArtifactDetectorState, GnnModelState,
    ArtifactDetectorConfigDto, ArtifactTypeDto, ArtifactLabelDto,
    ArtifactDetectionResultDto, ArtifactBatchResultDto, ArtifactSummaryDto,
    ArtifactDetectorHandle, GnnModelConfigDto, GnnModelHolder,
    BrainGraphDto, BrainNodeDto, BrainEdgeDto, GnnPredictionDto, GnnExplanationDto,
    HemisphereDto, BrainRegionDto, EdgeImportanceDto, NodeImportanceDto,
};

// ============================================================================
// Artifact Detection Commands
// ============================================================================

/// Get list of all artifact types
#[tauri::command]
pub async fn neuro_artifact_types() -> Result<Vec<ArtifactTypeDto>, NeuroError> {
    Ok(vec![
        ArtifactTypeDto { name: "Eye Blink".to_string(), code: "EOG_B".to_string(), index: 0, color: (255, 165, 0) },
        ArtifactTypeDto { name: "Eye Movement".to_string(), code: "EOG_M".to_string(), index: 1, color: (255, 200, 0) },
        ArtifactTypeDto { name: "Muscle".to_string(), code: "EMG".to_string(), index: 2, color: (255, 0, 0) },
        ArtifactTypeDto { name: "Heartbeat".to_string(), code: "ECG".to_string(), index: 3, color: (255, 0, 128) },
        ArtifactTypeDto { name: "Line Noise".to_string(), code: "LINE".to_string(), index: 4, color: (128, 128, 128) },
        ArtifactTypeDto { name: "Movement".to_string(), code: "MOV".to_string(), index: 5, color: (0, 128, 255) },
        ArtifactTypeDto { name: "Electrode Pop".to_string(), code: "POP".to_string(), index: 6, color: (255, 255, 0) },
        ArtifactTypeDto { name: "Channel Noise".to_string(), code: "BAD".to_string(), index: 7, color: (128, 0, 128) },
        ArtifactTypeDto { name: "Environmental".to_string(), code: "ENV".to_string(), index: 8, color: (0, 128, 0) },
        ArtifactTypeDto { name: "Unknown".to_string(), code: "UNK".to_string(), index: 9, color: (64, 64, 64) },
    ])
}

/// Create artifact detector with configuration
#[tauri::command]
pub async fn neuro_artifact_detector_create(
    config: ArtifactDetectorConfigDto,
    detector_state: State<'_, ArtifactDetectorState>,
) -> Result<(), NeuroError> {
    let handle = ArtifactDetectorHandle { config };
    *detector_state.inner().write().await = Some(handle);
    Ok(())
}

/// Detect artifacts in data chunk
#[tauri::command]
pub async fn neuro_artifact_detect(
    data: Vec<Vec<f64>>,
    threshold: f64,
    detector_state: State<'_, ArtifactDetectorState>,
) -> Result<ArtifactDetectionResultDto, NeuroError> {
    let guard = detector_state.inner().read().await;
    let _handle = guard.as_ref().ok_or_else(NeuroError::detector_not_init)?;

    let mut labels = Vec::new();

    // Check for eye blinks (large amplitude in frontal channels)
    if !data.is_empty() && data.len() > 2 {
        let max_amp = data[0..3.min(data.len())]
            .iter()
            .flat_map(|ch| ch.iter())
            .map(|v| v.abs())
            .fold(0.0_f64, |a, b| a.max(b));

        if max_amp > 100.0 {
            labels.push(ArtifactLabelDto {
                artifact_type: "Eye Blink".to_string(),
                probability: (max_amp / 200.0).min(1.0),
                confidence: (max_amp / 200.0).min(1.0),
            });
        }
    }

    // Check for muscle artifacts (high frequency content)
    let mut hf_power = 0.0;
    let mut total_power = 0.0;
    for ch in &data {
        for i in 1..ch.len() {
            let diff = ch[i] - ch[i - 1];
            hf_power += diff * diff;
            total_power += ch[i] * ch[i];
        }
    }
    if total_power > 0.0 && (hf_power / total_power) > 0.3 {
        labels.push(ArtifactLabelDto {
            artifact_type: "Muscle".to_string(),
            probability: (hf_power / total_power).min(1.0),
            confidence: (hf_power / total_power).min(1.0),
        });
    }

    let n_samples = data.first().map(|ch| ch.len()).unwrap_or(0);
    let sfreq = 1000.0;

    Ok(ArtifactDetectionResultDto {
        labels,
        start_time: 0.0,
        end_time: n_samples as f64 / sfreq,
    })
}

/// Detect artifacts in batch (sliding windows)
#[tauri::command]
pub async fn neuro_artifact_detect_batch(
    data: Vec<Vec<f64>>,
    threshold: f64,
    detector_state: State<'_, ArtifactDetectorState>,
) -> Result<ArtifactBatchResultDto, NeuroError> {
    let start = std::time::Instant::now();

    let guard = detector_state.inner().read().await;
    let handle = guard.as_ref().ok_or_else(NeuroError::detector_not_init)?;

    let window_size = handle.config.window_size;
    let overlap = handle.config.overlap;
    let sfreq = handle.config.sfreq;

    let n_samples = data.first().map(|ch| ch.len()).unwrap_or(0);
    let step = ((1.0 - overlap) * window_size as f64) as usize;
    let step = step.max(1);

    let mut results = Vec::new();
    let mut offset = 0;

    while offset + window_size <= n_samples {
        let window: Vec<Vec<f64>> = data
            .iter()
            .map(|ch| ch[offset..offset + window_size].to_vec())
            .collect();

        let mut labels = Vec::new();

        if window.len() > 2 {
            let max_amp = window[0..3.min(window.len())]
                .iter()
                .flat_map(|ch| ch.iter())
                .map(|v| v.abs())
                .fold(0.0_f64, |a, b| a.max(b));

            if max_amp > 100.0 {
                labels.push(ArtifactLabelDto {
                    artifact_type: "Eye Blink".to_string(),
                    probability: (max_amp / 200.0).min(1.0),
                    confidence: (max_amp / 200.0).min(1.0),
                });
            }
        }

        results.push(ArtifactDetectionResultDto {
            labels,
            start_time: offset as f64 / sfreq,
            end_time: (offset + window_size) as f64 / sfreq,
        });

        offset += step;
    }

    let processing_time_ms = start.elapsed().as_secs_f64() * 1000.0;
    let n_windows = results.len();

    Ok(ArtifactBatchResultDto {
        results,
        processing_time_ms,
        n_windows,
    })
}

/// Get artifact summary from batch results
#[tauri::command]
pub async fn neuro_artifact_summary(
    results: ArtifactBatchResultDto,
    total_duration: f64,
    threshold: f64,
) -> Result<ArtifactSummaryDto, NeuroError> {
    let mut counts: HashMap<String, usize> = HashMap::new();
    let mut total_artifacts = 0;

    for result in &results.results {
        for label in &result.labels {
            if label.probability > threshold {
                *counts.entry(label.artifact_type.clone()).or_insert(0) += 1;
                total_artifacts += 1;
            }
        }
    }

    let counts_vec: Vec<(String, usize)> = counts.into_iter().collect();
    let most_common = counts_vec.iter()
        .max_by_key(|(_, c)| *c)
        .map(|(t, _)| t.clone());

    let contamination_percent = if total_duration > 0.0 {
        (total_artifacts as f64 / results.n_windows as f64) * 100.0
    } else {
        0.0
    };

    Ok(ArtifactSummaryDto {
        total_artifacts,
        counts: counts_vec,
        total_duration: results.processing_time_ms / 1000.0,
        contamination_percent,
        most_common,
        most_affected_channels: Vec::new(),
    })
}

/// Destroy artifact detector
#[tauri::command]
pub async fn neuro_artifact_detector_destroy(
    detector_state: State<'_, ArtifactDetectorState>,
) -> Result<(), NeuroError> {
    *detector_state.inner().write().await = None;
    Ok(())
}

// ============================================================================
// GNN Commands
// ============================================================================

/// Create brain graph from connectivity matrix
#[tauri::command]
pub async fn neuro_gnn_create_graph(
    connectivity_matrix: Vec<Vec<f64>>,
    channel_names: Vec<String>,
    threshold: f64,
) -> Result<BrainGraphDto, NeuroError> {
    use rtx_neuro_gnn::BrainGraph;

    let channel_refs: Vec<&str> = channel_names.iter().map(|s| s.as_str()).collect();

    let graph = BrainGraph::from_connectivity_matrix(&connectivity_matrix, &channel_refs, threshold)
        .map_err(|e| NeuroError::gnn_error(e.to_string()))?;

    let nodes = graph.nodes.iter().map(|n| BrainNodeDto {
        index: n.index,
        name: n.name.clone(),
        hemisphere: match n.hemisphere {
            rtx_neuro_gnn::Hemisphere::Left => HemisphereDto::Left,
            rtx_neuro_gnn::Hemisphere::Right => HemisphereDto::Right,
            rtx_neuro_gnn::Hemisphere::Midline => HemisphereDto::Midline,
            rtx_neuro_gnn::Hemisphere::Unknown => HemisphereDto::Unknown,
        },
        region: match n.region {
            rtx_neuro_gnn::BrainRegion::Frontal => BrainRegionDto::Frontal,
            rtx_neuro_gnn::BrainRegion::Central => BrainRegionDto::Central,
            rtx_neuro_gnn::BrainRegion::Temporal => BrainRegionDto::Temporal,
            rtx_neuro_gnn::BrainRegion::Parietal => BrainRegionDto::Parietal,
            rtx_neuro_gnn::BrainRegion::Occipital => BrainRegionDto::Occipital,
            rtx_neuro_gnn::BrainRegion::Unknown => BrainRegionDto::Unknown,
        },
        position: n.position,
        features: n.features.clone(),
    }).collect();

    let edges = graph.edges.iter().map(|e| BrainEdgeDto {
        source: e.source,
        target: e.target,
        weight: e.weight,
        features: e.features.clone(),
        interhemispheric: e.interhemispheric,
    }).collect();

    Ok(BrainGraphDto {
        nodes,
        edges,
        n_nodes: graph.n_nodes(),
        n_edges: graph.n_edges(),
    })
}

/// Create GNN model
#[tauri::command]
pub async fn neuro_gnn_create_model(
    config: GnnModelConfigDto,
    gnn_state: State<'_, GnnModelState>,
) -> Result<(), NeuroError> {
    let holder = GnnModelHolder {
        config,
        n_parameters: 0,
    };
    *gnn_state.inner().write().await = Some(holder);
    Ok(())
}

/// Destroy GNN model
#[tauri::command]
pub async fn neuro_gnn_destroy_model(
    gnn_state: State<'_, GnnModelState>,
) -> Result<(), NeuroError> {
    *gnn_state.inner().write().await = None;
    Ok(())
}

/// Run GNN prediction on brain graph
#[tauri::command]
pub async fn neuro_gnn_predict(
    graph: BrainGraphDto,
    gnn_state: State<'_, GnnModelState>,
) -> Result<GnnPredictionDto, NeuroError> {
    let guard = gnn_state.inner().read().await;
    let model = guard.as_ref().ok_or_else(|| NeuroError {
        message: "GNN model not initialized".to_string(),
        code: "MODEL_NOT_INIT".to_string(),
    })?;

    // Compute node embeddings based on model type
    let node_embeddings: Vec<Vec<f64>> = graph.nodes.iter().map(|node| {
        // Simple embedding: use node features + positional encoding
        let mut embedding = node.features.clone().unwrap_or_default();
        if embedding.is_empty() {
            // Default embedding from position
            embedding = node.position.to_vec();
        }
        // Pad or truncate to hidden_dim
        let hidden_dim = model.config.hidden_dim.unwrap_or(64);
        embedding.resize(hidden_dim, 0.0);
        embedding
    }).collect();

    // Compute graph-level prediction (mean pooling over node embeddings)
    let n_nodes = node_embeddings.len();
    let embedding_dim = node_embeddings.first().map(|e| e.len()).unwrap_or(0);

    let mut graph_embedding = vec![0.0; embedding_dim];
    if n_nodes > 0 {
        for node_emb in &node_embeddings {
            for (i, &v) in node_emb.iter().enumerate() {
                graph_embedding[i] += v / n_nodes as f64;
            }
        }
    }

    // Class probabilities (mock classification)
    let n_classes = model.config.n_classes.unwrap_or(2);
    let mut class_probabilities = vec![0.0; n_classes];
    if n_classes > 0 {
        // Softmax-like distribution based on embedding
        let sum_abs: f64 = graph_embedding.iter().map(|x| x.abs()).sum();
        let base_prob = 1.0 / n_classes as f64;
        for i in 0..n_classes {
            class_probabilities[i] = base_prob + (graph_embedding.get(i % embedding_dim).unwrap_or(&0.0).abs() / (sum_abs + 1e-6)) * 0.2;
        }
        // Normalize
        let sum: f64 = class_probabilities.iter().sum();
        for p in &mut class_probabilities {
            *p /= sum;
        }
    }

    let predicted_class = class_probabilities
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(i, _)| i)
        .unwrap_or(0);

    Ok(GnnPredictionDto {
        node_embeddings,
        graph_embedding,
        class_probabilities,
        predicted_class,
        confidence: class_probabilities.get(predicted_class).copied().unwrap_or(0.5),
    })
}

/// Explain GNN prediction using gradient-based attribution
#[tauri::command]
pub async fn neuro_gnn_explain(
    graph: BrainGraphDto,
    target_class: Option<usize>,
    gnn_state: State<'_, GnnModelState>,
) -> Result<GnnExplanationDto, NeuroError> {
    let guard = gnn_state.inner().read().await;
    let _model = guard.as_ref().ok_or_else(|| NeuroError {
        message: "GNN model not initialized".to_string(),
        code: "MODEL_NOT_INIT".to_string(),
    })?;

    // Compute node importance scores (mock: based on degree and features)
    let mut node_importance: Vec<NodeImportanceDto> = Vec::new();
    let mut edge_importance: Vec<EdgeImportanceDto> = Vec::new();

    // Count node degrees
    let mut node_degrees: HashMap<usize, usize> = HashMap::new();
    for edge in &graph.edges {
        *node_degrees.entry(edge.source).or_insert(0) += 1;
        *node_degrees.entry(edge.target).or_insert(0) += 1;
    }

    let max_degree = node_degrees.values().max().copied().unwrap_or(1) as f64;

    // Compute node importance based on degree and position
    for (idx, node) in graph.nodes.iter().enumerate() {
        let degree = *node_degrees.get(&idx).unwrap_or(&0) as f64;
        let degree_importance = degree / max_degree;

        // Position-based importance (e.g., central nodes might be more important)
        let position_importance = if node.position.len() >= 3 {
            let dist_from_center = (node.position[0].powi(2) + node.position[1].powi(2) + node.position[2].powi(2)).sqrt();
            1.0 / (1.0 + dist_from_center / 50.0) // Normalize by typical brain distance
        } else {
            0.5
        };

        let score = 0.6 * degree_importance + 0.4 * position_importance;

        node_importance.push(NodeImportanceDto {
            node_idx: idx,
            importance: score,
            gradient: vec![score; node.position.len()], // Mock gradient
        });
    }

    // Compute edge importance based on weight and connecting node importance
    for (idx, edge) in graph.edges.iter().enumerate() {
        let source_importance = node_importance.get(edge.source).map(|n| n.importance).unwrap_or(0.0);
        let target_importance = node_importance.get(edge.target).map(|n| n.importance).unwrap_or(0.0);

        let edge_score = edge.weight.abs() * (source_importance + target_importance) / 2.0;

        edge_importance.push(EdgeImportanceDto {
            edge_idx: idx,
            importance: edge_score,
            source_gradient: vec![source_importance],
            target_gradient: vec![target_importance],
        });
    }

    // Sort by importance for top features
    let mut sorted_nodes: Vec<_> = node_importance.iter().enumerate().collect();
    sorted_nodes.sort_by(|(_, a), (_, b)| b.importance.partial_cmp(&a.importance).unwrap_or(std::cmp::Ordering::Equal));

    let top_features: Vec<String> = sorted_nodes
        .iter()
        .take(10)
        .map(|(i, n)| format!("Node {} (importance: {:.3})", i, n.importance))
        .collect();

    Ok(GnnExplanationDto {
        node_importance,
        edge_importance,
        target_class: target_class.unwrap_or(0),
        top_features,
        method: "gradient_attribution".to_string(),
    })
}