/// Export embeddings to CSV for downstream analysis in R/pandas.
use brainharmony::prelude::*;
use burn::backend::NdArray;

type B = NdArray;

fn main() -> anyhow::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 6 {
        eprintln!("usage: csv_export <weights> <gradient.csv> <geoh.csv> <input> <output.csv>");
        std::process::exit(1);
    }

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    brainharmony::init_threads(None);

    let (encoder, _) = BrainHarmonyEncoder::<B>::from_weights(
        &args[1], &args[2], &args[3],
        &ModelConfig::default(),
        &DataConfig::default(),
        &device,
    )?;

    let result = encoder.encode_safetensors(&args[4])?;
    println!("{} patches x {} dims", result.n_patches(), result.embed_dim());

    brainharmony::csv_export::save_embeddings_csv_with_metadata(
        &result, &args[5], result.n_rois, result.n_time_patches,
    )?;
    println!("Saved: {}", args[5]);

    Ok(())
}
