//! CSV export utilities for Brain-Harmony embeddings.

use crate::inference::EmbeddingResult;
use std::io::Write;

/// Write embeddings to a CSV file.
pub fn save_embeddings_csv(result: &EmbeddingResult, path: &str) -> anyhow::Result<()> {
    let embed_dim = result.embed_dim();
    let n_patches = result.n_patches();

    let mut file = std::fs::File::create(path)?;

    let header: Vec<String> = (0..embed_dim).map(|i| format!("dim_{i}")).collect();
    writeln!(file, "{}", header.join(","))?;

    for row in 0..n_patches {
        let start = row * embed_dim;
        let end = start + embed_dim;
        let vals: Vec<String> = result.embeddings[start..end]
            .iter()
            .map(|v| v.to_string())
            .collect();
        writeln!(file, "{}", vals.join(","))?;
    }

    Ok(())
}

/// Write embeddings to a CSV file with additional `roi_idx` and `time_idx` columns.
pub fn save_embeddings_csv_with_metadata(
    result: &EmbeddingResult,
    path: &str,
    n_rois: usize,
    n_time_patches: usize,
) -> anyhow::Result<()> {
    let embed_dim = result.embed_dim();

    let mut file = std::fs::File::create(path)?;

    let dim_cols: Vec<String> = (0..embed_dim).map(|i| format!("dim_{i}")).collect();
    writeln!(file, "roi_idx,time_idx,{}", dim_cols.join(","))?;

    let mut row = 0;
    for roi in 0..n_rois {
        for t in 0..n_time_patches {
            let start = row * embed_dim;
            let end = start + embed_dim;
            let vals: Vec<String> = result.embeddings[start..end]
                .iter()
                .map(|v| v.to_string())
                .collect();
            writeln!(file, "{roi},{t},{}", vals.join(","))?;
            row += 1;
        }
    }

    Ok(())
}
