//! Source localization (inverse problem) module

use pyo3::prelude::*;
use numpy::{PyArray1, PyArray2, PyReadonlyArray1, PyReadonlyArray2, IntoPyArray};
use rtx_neuro_inverse::{mne, beamformer, dipole, loreta};

/// Create the inverse submodule
pub fn create_module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
    let m = PyModule::new(py, "inverse")?;

    // MNE methods
    m.add_function(wrap_pyfunction!(compute_mne, &m)?)?;
    m.add_function(wrap_pyfunction!(compute_dspm, &m)?)?;
    m.add_function(wrap_pyfunction!(compute_sloreta, &m)?)?;
    m.add_function(wrap_pyfunction!(compute_eloreta, &m)?)?;

    // Beamformers
    m.add_function(wrap_pyfunction!(compute_lcmv, &m)?)?;

    // Dipole fitting
    m.add_function(wrap_pyfunction!(fit_dipole, &m)?)?;

    Ok(m)
}

/// Compute Minimum Norm Estimate (MNE)
///
/// # Arguments
/// * `data` - Sensor data [n_channels x n_times]
/// * `gain` - Forward model gain matrix [n_channels x n_sources]
/// * `noise_cov` - Noise covariance matrix [n_channels x n_channels]
/// * `lambda2` - Regularization parameter (default: 1.0)
///
/// # Returns
/// Source estimate [n_sources x n_times]
#[pyfunction]
#[pyo3(signature = (data, gain, noise_cov, lambda2=1.0))]
fn compute_mne<'py>(
    py: Python<'py>,
    data: PyReadonlyArray2<f64>,
    gain: PyReadonlyArray2<f64>,
    noise_cov: PyReadonlyArray2<f64>,
    lambda2: f64,
) -> PyResult<Bound<'py, PyArray2<f64>>> {
    let data_vec = array2_to_matrix(&data)?;
    let gain_vec = array2_to_matrix(&gain)?;
    let noise_vec = array2_to_matrix(&noise_cov)?;

    let config = mne::MneConfig {
        regularization: lambda2,
        method: mne::MneMethod::Mne,
    };

    let source = mne::apply_inverse(&data_vec, &gain_vec, &noise_vec, &config)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    let n_sources = source.len();
    let n_times = if n_sources > 0 { source[0].len() } else { 0 };

    let flat: Vec<f64> = source.into_iter().flatten().collect();
    let array = ndarray::Array2::from_shape_vec((n_sources, n_times), flat)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    Ok(array.into_pyarray(py))
}

/// Compute dynamic Statistical Parametric Mapping (dSPM)
///
/// Noise-normalized MNE providing z-score-like source estimates.
///
/// # Arguments
/// * `data` - Sensor data [n_channels x n_times]
/// * `gain` - Forward model gain matrix [n_channels x n_sources]
/// * `noise_cov` - Noise covariance matrix [n_channels x n_channels]
/// * `lambda2` - Regularization parameter (default: 1.0)
///
/// # Returns
/// Source estimate [n_sources x n_times]
#[pyfunction]
#[pyo3(signature = (data, gain, noise_cov, lambda2=1.0))]
fn compute_dspm<'py>(
    py: Python<'py>,
    data: PyReadonlyArray2<f64>,
    gain: PyReadonlyArray2<f64>,
    noise_cov: PyReadonlyArray2<f64>,
    lambda2: f64,
) -> PyResult<Bound<'py, PyArray2<f64>>> {
    let data_vec = array2_to_matrix(&data)?;
    let gain_vec = array2_to_matrix(&gain)?;
    let noise_vec = array2_to_matrix(&noise_cov)?;

    let config = mne::MneConfig {
        regularization: lambda2,
        method: mne::MneMethod::Dspm,
    };

    let source = mne::apply_inverse(&data_vec, &gain_vec, &noise_vec, &config)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    let n_sources = source.len();
    let n_times = if n_sources > 0 { source[0].len() } else { 0 };

    let flat: Vec<f64> = source.into_iter().flatten().collect();
    let array = ndarray::Array2::from_shape_vec((n_sources, n_times), flat)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    Ok(array.into_pyarray(py))
}

/// Compute standardized LOw Resolution brain Electromagnetic Tomography (sLORETA)
///
/// # Arguments
/// * `data` - Sensor data [n_channels x n_times]
/// * `gain` - Forward model gain matrix [n_channels x n_sources]
/// * `noise_cov` - Noise covariance matrix [n_channels x n_channels]
/// * `lambda2` - Regularization parameter (default: 1.0)
///
/// # Returns
/// Source estimate [n_sources x n_times]
#[pyfunction]
#[pyo3(signature = (data, gain, noise_cov, lambda2=1.0))]
fn compute_sloreta<'py>(
    py: Python<'py>,
    data: PyReadonlyArray2<f64>,
    gain: PyReadonlyArray2<f64>,
    noise_cov: PyReadonlyArray2<f64>,
    lambda2: f64,
) -> PyResult<Bound<'py, PyArray2<f64>>> {
    let data_vec = array2_to_matrix(&data)?;
    let gain_vec = array2_to_matrix(&gain)?;
    let noise_vec = array2_to_matrix(&noise_cov)?;

    let config = mne::MneConfig {
        regularization: lambda2,
        method: mne::MneMethod::Sloreta,
    };

    let source = mne::apply_inverse(&data_vec, &gain_vec, &noise_vec, &config)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    let n_sources = source.len();
    let n_times = if n_sources > 0 { source[0].len() } else { 0 };

    let flat: Vec<f64> = source.into_iter().flatten().collect();
    let array = ndarray::Array2::from_shape_vec((n_sources, n_times), flat)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    Ok(array.into_pyarray(py))
}

/// Compute exact LOw Resolution brain Electromagnetic Tomography (eLORETA)
///
/// # Arguments
/// * `data` - Sensor data [n_channels x n_times]
/// * `gain` - Forward model gain matrix [n_channels x n_sources]
/// * `noise_cov` - Noise covariance matrix [n_channels x n_channels]
/// * `lambda2` - Regularization parameter (default: 1e-6)
/// * `max_iter` - Maximum iterations (default: 100)
///
/// # Returns
/// Source estimate [n_sources x n_times]
#[pyfunction]
#[pyo3(signature = (data, gain, noise_cov, lambda2=1e-6, max_iter=100))]
fn compute_eloreta<'py>(
    py: Python<'py>,
    data: PyReadonlyArray2<f64>,
    gain: PyReadonlyArray2<f64>,
    noise_cov: PyReadonlyArray2<f64>,
    lambda2: f64,
    max_iter: usize,
) -> PyResult<Bound<'py, PyArray2<f64>>> {
    let data_vec = array2_to_matrix(&data)?;
    let gain_vec = array2_to_matrix(&gain)?;
    let noise_vec = array2_to_matrix(&noise_cov)?;

    let config = loreta::EloretaConfig {
        regularization: lambda2,
        max_iter,
        tol: 1e-6,
    };

    let source = loreta::eloreta(&data_vec, &gain_vec, &noise_vec, &config)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    let n_sources = source.len();
    let n_times = if n_sources > 0 { source[0].len() } else { 0 };

    let flat: Vec<f64> = source.into_iter().flatten().collect();
    let array = ndarray::Array2::from_shape_vec((n_sources, n_times), flat)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    Ok(array.into_pyarray(py))
}

/// Compute LCMV (Linearly Constrained Minimum Variance) beamformer
///
/// # Arguments
/// * `data` - Sensor data [n_channels x n_times]
/// * `gain` - Forward model gain matrix [n_channels x n_sources x 3]
/// * `data_cov` - Data covariance matrix [n_channels x n_channels]
/// * `regularization` - Regularization parameter (default: 0.05)
///
/// # Returns
/// Source power [n_sources]
#[pyfunction]
#[pyo3(signature = (data, gain, data_cov, regularization=0.05))]
fn compute_lcmv<'py>(
    py: Python<'py>,
    data: PyReadonlyArray2<f64>,
    gain: PyReadonlyArray2<f64>,
    data_cov: PyReadonlyArray2<f64>,
    regularization: f64,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let data_vec = array2_to_matrix(&data)?;
    let gain_vec = array2_to_matrix(&gain)?;
    let cov_vec = array2_to_matrix(&data_cov)?;

    let config = beamformer::LcmvConfig {
        regularization,
        pick_ori: beamformer::PickOri::MaxPower,
    };

    let power = beamformer::lcmv_source_power(&data_vec, &gain_vec, &cov_vec, &config)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    Ok(power.into_pyarray(py))
}

/// Fit a single equivalent current dipole
///
/// # Arguments
/// * `data` - Sensor data at one time point [n_channels]
/// * `gain_func` - Function to compute gain at any position
/// * `initial_pos` - Initial dipole position [x, y, z] in meters
/// * `fixed_orientation` - If True, fit only position (default: False)
/// * `max_iter` - Maximum iterations (default: 100)
///
/// # Returns
/// Dict with position, moment, gof (goodness of fit)
#[pyfunction]
#[pyo3(signature = (data, gain, initial_pos, fixed_orientation=false, max_iter=100))]
fn fit_dipole<'py>(
    py: Python<'py>,
    data: PyReadonlyArray1<f64>,
    gain: PyReadonlyArray2<f64>,
    initial_pos: [f64; 3],
    fixed_orientation: bool,
    max_iter: usize,
) -> PyResult<PyObject> {
    let data_vec = data.as_slice()?.to_vec();
    let gain_vec = array2_to_matrix(&gain)?;

    let config = dipole::DipoleConfig {
        n_dipoles: 1,
        fixed_orientation,
        max_iter,
        tol: 1e-6,
    };

    let result = dipole::fit_dipole(&data_vec, &gain_vec, &initial_pos, &config)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

    // Create Python dict with results
    let dict = pyo3::types::PyDict::new(py);
    dict.set_item("position", result.position.to_vec())?;
    dict.set_item("moment", result.moment.to_vec())?;
    dict.set_item("gof", result.gof)?;
    dict.set_item("residual_variance", result.residual_variance)?;

    Ok(dict.into())
}

// Helper function to convert 2D numpy array to Vec<Vec<f64>>
fn array2_to_matrix(arr: &PyReadonlyArray2<f64>) -> PyResult<Vec<Vec<f64>>> {
    let shape = arr.shape();
    let n_rows = shape[0];
    let n_cols = shape[1];

    let flat = arr.as_slice()?;
    let mut result = Vec::with_capacity(n_rows);

    for i in 0..n_rows {
        let start = i * n_cols;
        let end = start + n_cols;
        result.push(flat[start..end].to_vec());
    }

    Ok(result)
}