#![deny(missing_docs)]
//! Evaluate the default DADOES checkpoint and benchmark classification speed.

use std::error::Error;
use std::fmt::{Display, Formatter};
use std::fs;
use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::time::{Duration, Instant};

use dadoes::datasets::{
    DatasetParseError, parse_dadoes_jsonl, parse_empathetic_dialogues_csv,
    parse_fig_loneliness_jsonl, parse_loneliness_causes_csv, parse_um1neko_text_emotion_jsonl,
};
use dadoes::goemotions::{GoEmotionsError, parse_split};
use dadoes::{DadoesClassifier, EmotionClassifier, Mood, OwnedTrainingExample};

const DEFAULT_TEST_SPLIT: &str = "data/raw/goemotions/test.tsv";
const DEFAULT_EXTERNAL_DATA_DIR: &str = "data/raw/external";
const EXTERNAL_DATA_DIR_ENV: &str = "DADOES_EXTERNAL_DATA_DIR";
const INCLUDE_NON_COMMERCIAL_ENV: &str = "DADOES_INCLUDE_NON_COMMERCIAL";
const THRESHOLD: f32 = 0.35;
const BENCHMARK_REPEATS: usize = 50;

/// Runs default checkpoint evaluation and benchmark reporting.
fn main() -> ExitCode {
    match run() {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("{error}");
            ExitCode::from(1)
        }
    }
}

fn run() -> Result<(), EvaluateCliError> {
    let args = EvaluateArgs::from_env()?;
    let mut examples = load_split(&args.test_split)?;
    let external = load_external_test_mix(
        &args.external_data_dir,
        args.include_non_commercial,
        &mut examples,
    )?;
    let load_started = Instant::now();
    let classifier = DadoesClassifier::from_default_model()?;
    let load_duration = load_started.elapsed();
    let label_metrics = per_label_metrics(&classifier, &examples);
    let benchmark = benchmark_classifier(&classifier, &examples, BENCHMARK_REPEATS);

    println!("## Per-mood test metrics");
    println!();
    println!(
        "Loaded mixed test examples: {} (external files: {})",
        examples.len(),
        external.len()
    );
    for summary in &external {
        println!(
            "- {} examples={} path={}",
            summary.name,
            summary.examples,
            summary.path.display()
        );
    }
    println!();
    println!(
        "| Mood | Supervised Examples | Positives | Accuracy | Precision | Recall | F1 | Coverage |"
    );
    println!("|---|---:|---:|---:|---:|---:|---:|---|");
    for metrics in &label_metrics {
        if metrics.has_supervision() {
            println!("{}", metrics.markdown_row());
        }
    }
    let missing_labels = labels_without_supervision(&label_metrics);
    if !missing_labels.is_empty() {
        println!();
        println!(
            "Labels without supervised examples in the loaded mixed test: {}",
            missing_labels.join(", ")
        );
    }

    println!();
    println!("## Recognition benchmark");
    println!();
    println!("| Benchmark | Value |");
    println!("|---|---:|");
    println!("| Model load time | {:.3} ms |", duration_ms(load_duration));
    println!("| Test examples | {} |", examples.len());
    println!("| Repeats | {} |", benchmark.repeats);
    println!("| Total classifications | {} |", benchmark.classifications);
    println!(
        "| Total classification time | {:.3} ms |",
        duration_ms(benchmark.duration)
    );
    println!(
        "| Throughput | {:.1} texts/s |",
        benchmark.texts_per_second()
    );
    println!(
        "| Mean classification latency | {:.3} us/text |",
        benchmark.micros_per_text()
    );
    println!("| Score checksum | {:.6} |", benchmark.score_checksum);

    Ok(())
}

fn load_external_test_mix(
    root: &Path,
    include_non_commercial: bool,
    test: &mut Vec<OwnedTrainingExample>,
) -> Result<Vec<ExternalDatasetSummary>, EvaluateCliError> {
    let mut summaries = Vec::new();
    for file in EXTERNAL_TEST_DATASET_FILES {
        if !file.license.allowed_by(include_non_commercial) {
            continue;
        }

        let path = root.join(file.relative_path);
        if !path.exists() {
            continue;
        }

        let input = fs::read_to_string(&path).map_err(|source| EvaluateCliError::ReadFile {
            path: path.clone(),
            source,
        })?;
        let examples = (file.parser)(&input).map_err(|source| EvaluateCliError::ParseExternal {
            path: path.clone(),
            source,
        })?;
        let examples_count = examples.len();
        test.extend(examples);
        summaries.push(ExternalDatasetSummary {
            name: file.name,
            path,
            examples: examples_count,
        });
    }

    Ok(summaries)
}

fn load_split(path: &Path) -> Result<Vec<OwnedTrainingExample>, EvaluateCliError> {
    let input = fs::read_to_string(path).map_err(|source| EvaluateCliError::ReadFile {
        path: path.to_path_buf(),
        source,
    })?;
    parse_split(&input).map_err(EvaluateCliError::ParseSplit)
}

fn per_label_metrics(
    classifier: &DadoesClassifier,
    examples: &[OwnedTrainingExample],
) -> Vec<LabelMetrics> {
    let mut counters = Mood::ALL
        .iter()
        .copied()
        .map(LabelCounter::new)
        .collect::<Vec<_>>();

    for example in examples {
        let analysis = classifier.classify(example.text());
        for score in analysis.scores() {
            if !example.supervised_labels().contains(&score.mood) {
                continue;
            }
            if let Some(counter) = counters
                .iter_mut()
                .find(|counter| counter.mood == score.mood)
            {
                let target = example
                    .labels()
                    .iter()
                    .copied()
                    .any(|label| label == score.mood);
                let predicted = score.score >= THRESHOLD;
                counter.observe(predicted, target);
            }
        }
    }

    counters.into_iter().map(LabelCounter::metrics).collect()
}

fn benchmark_classifier(
    classifier: &DadoesClassifier,
    examples: &[OwnedTrainingExample],
    repeats: usize,
) -> BenchmarkMetrics {
    let started = Instant::now();
    let mut classifications = 0_usize;
    let mut score_checksum = 0.0_f32;

    for _ in 0..repeats {
        for example in examples {
            let analysis = classifier.classify(example.text());
            classifications += 1;
            if let Some(primary) = analysis.primary_mood() {
                score_checksum += primary.score;
            }
        }
    }

    BenchmarkMetrics {
        repeats,
        classifications,
        duration: started.elapsed(),
        score_checksum,
    }
}

#[derive(Debug, Copy, Clone)]
struct LabelCounter {
    mood: Mood,
    supervised: usize,
    positives: usize,
    true_positive: usize,
    false_positive: usize,
    true_negative: usize,
    false_negative: usize,
}

impl LabelCounter {
    fn new(mood: Mood) -> Self {
        Self {
            mood,
            supervised: 0,
            positives: 0,
            true_positive: 0,
            false_positive: 0,
            true_negative: 0,
            false_negative: 0,
        }
    }

    fn observe(&mut self, predicted: bool, target: bool) {
        self.supervised += 1;
        if target {
            self.positives += 1;
        }
        match (predicted, target) {
            (true, true) => self.true_positive += 1,
            (true, false) => self.false_positive += 1,
            (false, true) => self.false_negative += 1,
            (false, false) => self.true_negative += 1,
        }
    }

    fn metrics(self) -> LabelMetrics {
        let accuracy = ratio(self.true_positive + self.true_negative, self.supervised);
        let precision = ratio(self.true_positive, self.true_positive + self.false_positive);
        let recall = ratio(self.true_positive, self.true_positive + self.false_negative);
        let f1 = harmonic_mean(precision, recall);

        LabelMetrics {
            mood: self.mood,
            supervised: self.supervised,
            positives: self.positives,
            accuracy,
            precision,
            recall,
            f1,
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct LabelMetrics {
    mood: Mood,
    supervised: usize,
    positives: usize,
    accuracy: Option<f32>,
    precision: Option<f32>,
    recall: Option<f32>,
    f1: Option<f32>,
}

impl LabelMetrics {
    fn has_supervision(&self) -> bool {
        self.supervised > 0
    }

    fn markdown_row(&self) -> String {
        format!(
            "| {} | {} | {} | {} | {} | {} | {} | {} |",
            self.mood.as_str(),
            self.supervised,
            self.positives,
            format_metric(self.accuracy),
            format_metric(self.precision),
            format_metric(self.recall),
            format_metric(self.f1),
            "loaded mixed test"
        )
    }
}

#[derive(Debug, Copy, Clone)]
struct BenchmarkMetrics {
    repeats: usize,
    classifications: usize,
    duration: Duration,
    score_checksum: f32,
}

impl BenchmarkMetrics {
    fn texts_per_second(&self) -> f64 {
        usize_to_f64(self.classifications) / self.duration.as_secs_f64()
    }

    fn micros_per_text(&self) -> f64 {
        duration_micros(self.duration) / usize_to_f64(self.classifications)
    }
}

#[derive(Debug, Clone, Eq, PartialEq)]
struct EvaluateArgs {
    test_split: PathBuf,
    external_data_dir: PathBuf,
    include_non_commercial: bool,
}

impl EvaluateArgs {
    fn from_env() -> Result<Self, EvaluateCliError> {
        let args = std::env::args().skip(1).collect::<Vec<_>>();
        let external_data_dir = std::env::var(EXTERNAL_DATA_DIR_ENV)
            .map(PathBuf::from)
            .unwrap_or_else(|_| PathBuf::from(DEFAULT_EXTERNAL_DATA_DIR));
        let include_non_commercial = env_flag(INCLUDE_NON_COMMERCIAL_ENV);
        match args.as_slice() {
            [] => Ok(Self {
                test_split: PathBuf::from(DEFAULT_TEST_SPLIT),
                external_data_dir,
                include_non_commercial,
            }),
            [test_split] => Ok(Self {
                test_split: PathBuf::from(test_split),
                external_data_dir,
                include_non_commercial,
            }),
            _ => Err(EvaluateCliError::Usage),
        }
    }
}

#[derive(Debug)]
enum EvaluateCliError {
    Usage,
    ReadFile {
        path: PathBuf,
        source: std::io::Error,
    },
    ParseSplit(GoEmotionsError),
    ParseExternal {
        path: PathBuf,
        source: DatasetParseError,
    },
    Model(dadoes::ModelIoError),
}

impl Display for EvaluateCliError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            EvaluateCliError::Usage => write!(
                formatter,
                "usage: cargo run --release --bin evaluate -- [test_split]"
            ),
            EvaluateCliError::ReadFile { path, source } => {
                write!(formatter, "failed to read {}: {source}", path.display())
            }
            EvaluateCliError::ParseSplit(error) => {
                write!(formatter, "failed to parse split: {error}")
            }
            EvaluateCliError::ParseExternal { path, source } => {
                write!(
                    formatter,
                    "failed to parse external dataset {}: {source}",
                    path.display()
                )
            }
            EvaluateCliError::Model(error) => write!(formatter, "failed to load model: {error}"),
        }
    }
}

impl Error for EvaluateCliError {}

impl From<dadoes::ModelIoError> for EvaluateCliError {
    fn from(error: dadoes::ModelIoError) -> Self {
        EvaluateCliError::Model(error)
    }
}

fn ratio(numerator: usize, denominator: usize) -> Option<f32> {
    if denominator == 0 {
        return None;
    }
    Some(usize_to_f32(numerator) / usize_to_f32(denominator))
}

fn harmonic_mean(left: Option<f32>, right: Option<f32>) -> Option<f32> {
    let left = left?;
    let right = right?;
    let denominator = left + right;
    if denominator <= f32::EPSILON {
        return Some(0.0);
    }
    Some(2.0 * left * right / denominator)
}

fn format_metric(metric: Option<f32>) -> String {
    match metric {
        Some(value) => format!("{value:.4}"),
        None => "n/a".to_owned(),
    }
}

fn labels_without_supervision(metrics: &[LabelMetrics]) -> Vec<&'static str> {
    metrics
        .iter()
        .filter(|metrics| !metrics.has_supervision())
        .map(|metrics| metrics.mood.as_str())
        .collect()
}

fn env_flag(name: &str) -> bool {
    match std::env::var(name) {
        Ok(value) => matches!(
            value.trim().to_ascii_lowercase().as_str(),
            "1" | "true" | "yes" | "on"
        ),
        Err(_) => false,
    }
}

fn duration_ms(duration: Duration) -> f64 {
    duration.as_secs_f64() * 1_000.0
}

fn duration_micros(duration: Duration) -> f64 {
    duration.as_secs_f64() * 1_000_000.0
}

fn usize_to_f32(value: usize) -> f32 {
    value.to_string().parse::<f32>().unwrap_or(f32::INFINITY)
}

fn usize_to_f64(value: usize) -> f64 {
    value.to_string().parse::<f64>().unwrap_or(f64::INFINITY)
}

type ExternalParser = fn(&str) -> Result<Vec<OwnedTrainingExample>, DatasetParseError>;

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum LicenseClass {
    Public,
    NonCommercial,
}

impl LicenseClass {
    fn allowed_by(self, include_non_commercial: bool) -> bool {
        match self {
            LicenseClass::Public => true,
            LicenseClass::NonCommercial => include_non_commercial,
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct ExternalDatasetFile {
    name: &'static str,
    license: LicenseClass,
    relative_path: &'static str,
    parser: ExternalParser,
}

#[derive(Debug, Clone)]
struct ExternalDatasetSummary {
    name: &'static str,
    path: PathBuf,
    examples: usize,
}

const EXTERNAL_TEST_DATASET_FILES: &[ExternalDatasetFile] = &[
    ExternalDatasetFile {
        name: "dadoes-domain",
        license: LicenseClass::Public,
        relative_path: "dadoes-domain/test.jsonl",
        parser: parse_dadoes_jsonl,
    },
    ExternalDatasetFile {
        name: "facebook/empathetic_dialogues",
        license: LicenseClass::NonCommercial,
        relative_path: "empathetic_dialogues/test.csv",
        parser: parse_empathetic_dialogues_csv,
    },
    ExternalDatasetFile {
        name: "FIG-Loneliness/FIG-Loneliness",
        license: LicenseClass::NonCommercial,
        relative_path: "fig_loneliness/test.jsonl",
        parser: parse_fig_loneliness_jsonl,
    },
    ExternalDatasetFile {
        name: "yael-katsman/Loneliness-Causes-and-Intensity",
        license: LicenseClass::Public,
        relative_path: "loneliness_causes/test_data.csv",
        parser: parse_loneliness_causes_csv,
    },
    ExternalDatasetFile {
        name: "Um1neko/text_emotion",
        license: LicenseClass::Public,
        relative_path: "um1neko/test.json",
        parser: parse_um1neko_text_emotion_jsonl,
    },
];