use clap::Args;

use kuva::plot::roc::{RocGroup, RocPlot};
use kuva::render::layout::Layout;
use kuva::render::palette::Palette;
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;

/// ROC curve — Receiver Operating Characteristic plot from score/label data.
#[derive(Args, Debug)]
pub struct RocArgs {
    /// Score column (classifier output; 0-based index or header name; default: 0).
    #[arg(long)]
    pub score_col: Option<ColSpec>,

    /// Label column (0/1 or false/true; 1 = positive; default: 1).
    #[arg(long)]
    pub label_col: Option<ColSpec>,

    /// Group by this column — one ROC curve per unique value.
    #[arg(long)]
    pub color_by: Option<ColSpec>,

    /// Hide the diagonal no-skill reference line.
    #[arg(long)]
    pub no_diagonal: bool,

    /// Show DeLong 95% confidence interval bands.
    #[arg(long)]
    pub ci: bool,

    /// Annotate each curve with its AUC value.
    #[arg(long)]
    pub auc_label: bool,

    /// Legend title.
    #[arg(long)]
    pub legend: Option<String>,

    #[command(flatten)]
    pub input: InputArgs,

    #[command(flatten)]
    pub base: BaseArgs,
    #[command(flatten)]
    pub axis: AxisArgs,
}

pub fn run(args: RocArgs) -> Result<(), String> {
    let table = DataTable::parse(
        args.input.input.as_deref(),
        args.input.no_header,
        args.input.delimiter,
    )?;

    let score_col = args.score_col.unwrap_or(ColSpec::Index(0));
    let label_col = args.label_col.unwrap_or(ColSpec::Index(1));

    let pal = Palette::category10();

    let mut plot = RocPlot::new();

    if args.no_diagonal {
        plot = plot.with_diagonal(false);
    }
    if let Some(legend) = args.legend {
        plot = plot.with_legend(legend);
    }

    if let Some(ref cb) = args.color_by {
        let groups = table.group_by(cb)?;
        for (i, (name, subtable)) in groups.into_iter().enumerate() {
            let scores = subtable.col_f64(&score_col)?;
            let labels = subtable.col_f64(&label_col)?;
            let predictions: Vec<(f64, bool)> = scores
                .into_iter()
                .zip(labels)
                .map(|(s, l)| (s, l > 0.5))
                .collect();
            let mut group = RocGroup::new(&name)
                .with_raw(predictions)
                .with_color(pal[i % pal.len()].to_string());
            if args.ci {
                group = group.with_ci(true);
            }
            if args.auc_label {
                group = group.with_auc_label(true);
            }
            plot = plot.with_group(group);
        }
    } else {
        let scores = table.col_f64(&score_col)?;
        let labels = table.col_f64(&label_col)?;
        let predictions: Vec<(f64, bool)> = scores
            .into_iter()
            .zip(labels)
            .map(|(s, l)| (s, l > 0.5))
            .collect();
        let mut group = RocGroup::new("Model")
            .with_raw(predictions)
            .with_color(pal[0].to_string());
        if args.ci {
            group = group.with_ci(true);
        }
        if args.auc_label {
            group = group.with_auc_label(true);
        }
        plot = plot.with_group(group);
    }

    let plots = vec![Plot::Roc(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 scene = render_multiple(plots, layout);
    write_output(scene, &args.base)
}
