data-beans 0.10.13

use crate::handlers::explore::{off_screen, open_session, pick_entries};
use crate::hdf5_io::*;
use crate::interactive::list_tui::{browse_in, ListView};
use crate::interactive::picks;
use crate::interactive::stat_tui::Side;
use crate::interactive::ui::Tui;
use crate::sparse_io::*;
use crate::utilities::io_helpers::{
    parse_index_spec_within, read_col_names, MAX_COLUMN_NAME_IDX, MAX_ROW_NAME_IDX,
};
use crate::utilities::name_matching::match_by_substring;
use clap::Args;
use legume_numeric::matrix::common_io::*;
use legume_numeric::matrix::traits::IoOps;
use std::io::Write;

/// A quick information of the underlying matrix of a backend file.
#[derive(Args, Debug)]
pub struct InfoArgs {
    /// data file -- .zarr or .h5 file
    pub data_file: Box<str>,

    /// file header for {output}.{rows.gz,columns.gz}
    #[arg(short, long, default_value = "")]
    pub output: Box<str>,
}

#[derive(Args, Debug)]
pub struct TakeColumnsArgs {
    /// data file -- either `.zarr` or `.h5`
    pub data_file: Box<str>,

    /// column indices to take: singles and inclusive ranges, e.g. `0,1,2,3`, `1-10`, or `1-20,50-55`
    #[arg(short = 'i', long)]
    pub column_indices: Option<Box<str>>,

    /// column name file where each line is a column name
    #[arg(short = 'f', long)]
    pub name_file: Option<Box<str>>,

    /// column names to take: e.g., `col1,col2,col3` (supports substring matching)
    #[arg(short = 'n', long, value_delimiter = ',')]
    pub column_names: Option<Vec<Box<str>>>,

    /// output `parquet` file
    #[arg(short, long, default_value = "stdout")]
    pub output: Box<str>,

    /// never open the explorer: what to take is given on the command line
    #[arg(long, default_value_t = false)]
    pub batch_process: bool,

    #[command(flatten)]
    pub session: legume_tui::cli::SessionArgs,
}

#[derive(Args, Debug)]
pub struct TakeRowsArgs {
    /// data file -- either `.zarr` or `.h5`
    pub data_file: Box<str>,

    /// row indices to take: singles and inclusive ranges, e.g. `0,1,2,3`, `1-10`, or `1-20,50-55`
    #[arg(short = 'i', long)]
    pub row_indices: Option<Box<str>>,

    /// row name file where each line is a row name
    #[arg(short = 'f', long)]
    pub name_file: Option<Box<str>>,

    /// row names to take: e.g., `gene1,gene2,gene3` (supports substring matching)
    #[arg(short = 'n', long, value_delimiter = ',')]
    pub row_names: Option<Vec<Box<str>>>,

    /// output `parquet` file
    #[arg(short, long, default_value = "stdout")]
    pub output: Box<str>,

    /// never open the explorer: what to take is given on the command line
    #[arg(long, default_value_t = false)]
    pub batch_process: bool,

    #[command(flatten)]
    pub session: legume_tui::cli::SessionArgs,
}

#[derive(Args, Debug)]
pub struct TakeColumnNamesArgs {
    /// data file -- either `.zarr` or `.h5`
    pub data_file: Box<str>,

    /// output file
    #[arg(short, long, default_value = "stdout")]
    pub output: Box<str>,

    #[arg(
        long,
        default_value_t = false,
        help = "Print the names and stop: no browser, even on a terminal",
        long_help = "On a terminal, the names open in a browser instead of being printed: `/`\n\
                     filters, Space marks, `w` writes the marked (or shown) names, `y` copies\n\
                     them. A file --output is still written. --batch-process (or no terminal)\n\
                     prints or writes them and stops."
    )]
    pub batch_process: bool,

    #[command(flatten)]
    pub session: legume_tui::cli::SessionArgs,
}

#[derive(Args, Debug)]
pub struct TakeRowNamesArgs {
    /// data file -- either `.zarr` or `.h5`
    pub data_file: Box<str>,

    /// output file
    #[arg(short, long, default_value = "stdout")]
    pub output: Box<str>,

    #[arg(
        long,
        default_value_t = false,
        help = "Print the names and stop: no browser, even on a terminal",
        long_help = "On a terminal, the names open in a browser instead of being printed: `/`\n\
                     filters, Space marks, `w` writes the marked (or shown) names, `y` copies\n\
                     them. A file --output is still written. --batch-process (or no terminal)\n\
                     prints or writes them and stops."
    )]
    pub batch_process: bool,

    #[command(flatten)]
    pub session: legume_tui::cli::SessionArgs,
}

/// Write a matrix as TSV with optional row/column names.
/// Arguments mirror `to_parquet_with_names`:
/// - `row_names`: `(Option<&[Box<str>]>, Option<&str>)` — names and label for the row dimension
/// - `column_names`: `Option<&[Box<str>]>` — column headers
fn write_named_tsv(
    output: &str,
    data: &DMatrix<f32>,
    row_names: (Option<&[Box<str>]>, Option<&str>),
    column_names: Option<&[Box<str>]>,
) -> anyhow::Result<()> {
    let nrows = data.nrows();
    let row_label = row_names.1.unwrap_or("row");
    let row_names: Vec<Box<str>> = match row_names.0 {
        Some(names) => names.to_vec(),
        None => (0..nrows).map(|i| format!("{}", i).into()).collect(),
    };
    let mut w = open_buf_writer(output)?;
    // Header
    write!(w, "{}", row_label)?;
    if let Some(col_names) = column_names {
        for name in col_names {
            write!(w, "\t{}", name)?;
        }
    } else {
        for j in 0..data.ncols() {
            write!(w, "\t{}", j)?;
        }
    }
    writeln!(w)?;
    // Data rows
    for (i, row) in data.row_iter().enumerate() {
        write!(w, "{}", row_names[i])?;
        for val in row.iter() {
            write!(w, "\t{}", val)?;
        }
        writeln!(w)?;
    }
    Ok(())
}

pub fn take_column_names(cmd_args: &TakeColumnNamesArgs) -> anyhow::Result<()> {
    let names = Names {
        command: "column-names",
        what: "columns",
        data_file: &cmd_args.data_file,
        output: &cmd_args.output,
        batch: cmd_args.batch_process,
        session: &cmd_args.session,
    };
    names.run(|data| data.column_names())
}

pub fn take_row_names(cmd_args: &TakeRowNamesArgs) -> anyhow::Result<()> {
    let names = Names {
        command: "row-names",
        what: "rows",
        data_file: &cmd_args.data_file,
        output: &cmd_args.output,
        batch: cmd_args.batch_process,
        session: &cmd_args.session,
    };
    names.run(|data| data.row_names())
}

/// `row-names` or `column-names`: the names to `output`, then, on a
/// terminal, in a browser (which replaces printing them to stdout).
struct Names<'a> {
    command: &'static str,
    what: &'static str,
    data_file: &'a str,
    output: &'a str,
    batch: bool,
    session: &'a legume_tui::cli::SessionArgs,
}

impl Names<'_> {
    fn run(
        &self,
        names_of: impl FnOnce(&dyn SparseIo<IndexIter = Vec<usize>>) -> anyhow::Result<Vec<Box<str>>>,
    ) -> anyhow::Result<()> {
        let view = crate::interactive::opens_view(self.command, self.batch);
        // The session first: hosted, the names are read there.
        let session = view
            .then(|| open_session(self.command, self.session))
            .transpose()?;
        let (backend, input) = resolve_backend_file(self.data_file, None)?;
        let data = open_sparse_matrix(&input, &backend)?;
        let names = names_of(data.as_ref())?;
        let to_stdout = self.output.eq_ignore_ascii_case("stdout");
        if !(view && to_stdout) {
            write_lines(&names, self.output)?;
        }
        let Some(session) = session else {
            return Ok(());
        };
        let stem = basename(self.data_file).unwrap_or_else(|_| self.data_file.into());
        let write_to = format!("{stem}.{}.txt", self.what);
        let title = format!("{} · {}", self.data_file, self.what);
        let mut tui = Tui::new(&session);
        browse_in(
            &mut tui,
            ListView::names("names", &title, &names, &write_to),
        )
    }
}

pub fn show_info(cmd_args: &InfoArgs) -> anyhow::Result<()> {
    let output = cmd_args.output.clone();

    let (backend, input) = resolve_backend_file(&cmd_args.data_file, None)?;

    let data = open_sparse_matrix(&input, &backend)?;

    if let (Some(nrow), Some(ncol), Some(nnz)) =
        (data.num_rows(), data.num_columns(), data.num_non_zeros())
    {
        println!("number_of_rows:\t{}", nrow);
        println!("number_of_columns:\t{}", ncol);
        println!("number_of_nonzeros:\t{}", nnz);
    }
    for (key, value) in data.metadata() {
        println!("meta.{key}:\t{value}");
    }

    if !output.is_empty() {
        use {mkdir, write_lines};
        mkdir(&output)?;
        let row_names = data.row_names()?;
        let col_names = data.column_names()?;
        write_lines(&row_names, &(output.to_string() + ".rows.gz"))?;
        write_lines(&col_names, &(output.to_string() + ".columns.gz"))?;
    }

    Ok(())
}

pub fn take_columns(args: &TakeColumnsArgs) -> anyhow::Result<()> {
    // The picker's session first: hosted, the rest runs in the session.
    // Nothing selected: picked in the explorer.
    let selected =
        args.column_indices.is_some() || args.name_file.is_some() || args.column_names.is_some();
    let options = "--column-indices, --name-file or --column-names";
    let session = picks("columns", args.batch_process, selected, options)?
        .then(|| open_session("columns", &args.session))
        .transpose()?;
    let mut tui = session.as_ref().map(Tui::new);
    let columns = args.column_indices.clone();
    let column_name_file = args.name_file.clone();
    let column_names_arg = args.column_names.clone();

    let (backend, data_file) = resolve_backend_file(&args.data_file, None)?;

    let output = args.output.clone();

    let picked = match tui.as_mut() {
        Some(tui) => match pick_entries(tui, &args.data_file, Side::Columns, "take")? {
            Some(picked) => Some(picked),
            None => return Ok(()),
        },
        None => None,
    };

    // Written off the screen: a `kill` meanwhile lets it finish.
    off_screen(tui.as_mut(), || {
        let data = open_sparse_matrix(&data_file, &backend)?;
        let row_names = data.row_names()?;

        let (data, column_names) = if let Some(columns) = picked {
            let names = data.column_names()?;
            let column_names: Vec<Box<str>> = columns.iter().map(|&i| names[i].clone()).collect();
            (data.read_columns_dmatrix(columns)?, column_names)
        } else if let Some(columns) = columns {
            let columns = parse_index_spec_within(&columns, data.num_columns().unwrap_or(0))?;

            if columns.is_empty() {
                return Err(anyhow::anyhow!("invalid indexes"));
            }

            let _names = data.column_names()?;
            let column_names: Vec<Box<str>> = columns.iter().map(|&i| _names[i].clone()).collect();

            (data.read_columns_dmatrix(columns)?, column_names)
        } else if let Some(column_file) = column_name_file {
            let col_names_to_match = read_col_names(column_file, MAX_COLUMN_NAME_IDX)?;
            let all_names = data.column_names()?;
            let (matched_indices, column_names) =
                match_by_substring(&all_names, &col_names_to_match, "column")?;
            (data.read_columns_dmatrix(matched_indices)?, column_names)
        } else if let Some(col_names_to_match) = column_names_arg {
            let all_names = data.column_names()?;
            let (matched_indices, column_names) =
                match_by_substring(&all_names, &col_names_to_match, "column")?;
            (data.read_columns_dmatrix(matched_indices)?, column_names)
        } else {
            return Err(anyhow::anyhow!(
                "either `column-indices`, `name-file` or `column-names` must be provided"
            ));
        };

        if let Ok(ext) = file_ext(&output) {
            if ext.as_ref() == "parquet" {
                data.to_parquet_with_names(
                    &output,
                    (Some(&row_names), Some("row_name")),
                    Some(&column_names),
                )?;
                return Ok(());
            }
        }

        write_named_tsv(
            &output,
            &data,
            (Some(&row_names), Some("row_name")),
            Some(&column_names),
        )?;

        Ok(())
    })
}

pub fn take_rows(args: &TakeRowsArgs) -> anyhow::Result<()> {
    // The picker's session first: hosted, the rest runs in the session.
    // Nothing selected: picked in the explorer.
    let selected =
        args.row_indices.is_some() || args.name_file.is_some() || args.row_names.is_some();
    let options = "--row-indices, --name-file or --row-names";
    let session = picks("rows", args.batch_process, selected, options)?
        .then(|| open_session("rows", &args.session))
        .transpose()?;
    let mut tui = session.as_ref().map(Tui::new);
    let rows = args.row_indices.clone();
    let row_name_file = args.name_file.clone();
    let row_names_arg = args.row_names.clone();

    let (backend, data_file) = resolve_backend_file(&args.data_file, None)?;

    let output = args.output.clone();

    let picked = match tui.as_mut() {
        Some(tui) => match pick_entries(tui, &args.data_file, Side::Rows, "take")? {
            Some(picked) => Some(picked),
            None => return Ok(()),
        },
        None => None,
    };

    // Written off the screen: a `kill` meanwhile lets it finish.
    off_screen(tui.as_mut(), || {
        let data_backend = open_sparse_matrix(&data_file, &backend)?;

        let (data, row_names) = if let Some(rows) = picked {
            let names = data_backend.row_names()?;
            let row_names: Vec<Box<str>> = rows.iter().map(|&i| names[i].clone()).collect();
            (data_backend.read_rows_dmatrix(rows)?, row_names)
        } else if let Some(rows) = rows {
            let rows = parse_index_spec_within(&rows, data_backend.num_rows().unwrap_or(0))?;

            if rows.is_empty() {
                return Err(anyhow::anyhow!("invalid indexes"));
            }

            let _names = data_backend.row_names()?;
            let row_names: Vec<Box<str>> = rows.iter().map(|&i| _names[i].clone()).collect();

            (data_backend.read_rows_dmatrix(rows)?, row_names)
        } else if let Some(row_name_file) = row_name_file {
            let row_names_to_match = read_col_names(row_name_file, MAX_ROW_NAME_IDX)?;
            let all_names = data_backend.row_names()?;
            let (matched_indices, row_names) =
                match_by_substring(&all_names, &row_names_to_match, "row")?;
            (data_backend.read_rows_dmatrix(matched_indices)?, row_names)
        } else if let Some(row_names_to_match) = row_names_arg {
            let all_names = data_backend.row_names()?;
            let (matched_indices, row_names) =
                match_by_substring(&all_names, &row_names_to_match, "row")?;
            (data_backend.read_rows_dmatrix(matched_indices)?, row_names)
        } else {
            return Err(anyhow::anyhow!(
                "either `row-indices`, `name-file` or `row-names` must be provided"
            ));
        };

        let data_t = data.transpose();
        let column_names = data_backend.column_names()?;

        if let Ok(ext) = file_ext(&output) {
            if ext.as_ref() == "parquet" {
                data_t.to_parquet_with_names(
                    &output,
                    (Some(&column_names), Some("column_name")),
                    Some(&row_names),
                )?;
                return Ok(());
            }
        }

        write_named_tsv(
            &output,
            &data_t,
            (Some(&column_names), Some("column_name")),
            Some(&row_names),
        )?;
        Ok(())
    })
}