use clap::Args;
use kuva::plot::DotPlot;
use kuva::render::layout::Layout;
use kuva::render::plots::Plot;
use kuva::render::render::render_multiple;
use crate::data::{ColSpec, DataTable, InputArgs};
use crate::layout_args::{apply_axis_args, apply_base_args, AxisArgs, BaseArgs};
use crate::output::write_output;
#[derive(Args, Debug)]
pub struct DotArgs {
#[arg(long)]
pub x_col: Option<ColSpec>,
#[arg(long)]
pub y_col: Option<ColSpec>,
#[arg(long)]
pub size_col: Option<ColSpec>,
#[arg(long)]
pub color_col: Option<ColSpec>,
#[arg(long, default_value = "viridis")]
pub colormap: String,
#[arg(long)]
pub max_radius: Option<f64>,
#[arg(long)]
pub size_legend: Option<String>,
#[arg(long)]
pub colorbar: Option<String>,
#[command(flatten)]
pub input: InputArgs,
#[command(flatten)]
pub base: BaseArgs,
#[command(flatten)]
pub axis: AxisArgs,
}
use crate::data::parse_colormap;
pub fn run(args: DotArgs) -> Result<(), String> {
let proj: Vec<ColSpec> = vec![
args.x_col.clone().unwrap_or(ColSpec::Index(0)),
args.y_col.clone().unwrap_or(ColSpec::Index(1)),
args.size_col.clone().unwrap_or(ColSpec::Index(2)),
args.color_col.clone().unwrap_or(ColSpec::Index(3)),
];
let table = DataTable::parse(
args.input.input.as_deref(),
args.input.no_header,
args.input.delimiter,
&proj,
)?;
let x_col = args.x_col.unwrap_or(ColSpec::Index(0));
let y_col = args.y_col.unwrap_or(ColSpec::Index(1));
let size_col = args.size_col.unwrap_or(ColSpec::Index(2));
let color_col = args.color_col.unwrap_or(ColSpec::Index(3));
let xs = table.col_str(&x_col)?;
let ys = table.col_str(&y_col)?;
let sizes = table.col_f64(&size_col)?;
let colors = table.col_f64(&color_col)?;
let pts: Vec<(String, String, f64, f64)> = xs
.into_iter()
.zip(ys)
.zip(sizes)
.zip(colors)
.map(|(((x, y), s), c)| (x, y, s, c))
.collect();
let mut plot = DotPlot::new()
.with_data(pts)
.with_color_map(parse_colormap(&args.colormap));
if let Some(r) = args.max_radius {
plot = plot.with_max_radius(r);
}
if let Some(ref label) = args.size_legend {
plot = plot.with_size_legend(label.clone());
}
if let Some(ref label) = args.colorbar {
plot = plot.with_colorbar(label.clone());
}
let plots = vec![Plot::DotPlot(plot)];
let layout = Layout::auto_from_plots(&plots);
let layout = apply_base_args(layout, &args.base);
let layout = apply_axis_args(layout, &args.axis);
let layout = layout.with_x_tick_rotate(-45.0);
let scene = render_multiple(plots, layout);
write_output(scene, &args.base)
}