cost_adr.rs

use fsrs::{
    CombinedProgressState, CostAdrEvaluationConfig, CostAdrEvaluationResult, CostAdrPolicy,
    CostAdrTrainingConfig, DEFAULT_PARAMETERS, FSRS, FSRSError, SimulationResult, SimulatorConfig,
    simulate_with_cost_adr_policy,
};
use std::env;
use std::error::Error;
use std::fmt::Display;
use std::io::{Error as IoError, ErrorKind, Write};
use std::str::FromStr;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration as StdDuration, Instant};

const EXAMPLE_DEFAULT_BASE_SEED: u64 = 42;
const EXAMPLE_EVALUATION_SEED_OFFSET: u64 = 2_000_000;
const EXAMPLE_FINAL_ROLLOUT_SEED_OFFSET: u64 = 3_000_000;

const USAGE: &str = "\
Usage: cargo run --release --features experimental_cost_adr --example cost_adr -- [OPTIONS]

Options:
  --days <usize>            Simulation learn span in days (default: 1825)
  --deck <usize>            Simulated deck size (default: 10000)
  --learn-limit <usize>     New-card limit per day (default: 10)
  --review-limit <usize>    Review limit per day (default: 9999)
  --cost-limit-minutes <f32>
                            Study time limit per day in minutes (default: 720.0)
  --pop <usize>             CMA-ES population size (default: 16)
  --gen <usize>             CMA-ES generation count (default: 20)
  --seed <u64>              Base seed for optimizer and derived simulation seeds (default: 42)
  --sigma0 <f32>            CMA-ES initial sigma (default: 1.0)
  -h, --help                Print help
";

struct ExampleConfig {
    days: usize,
    deck_size: usize,
    learn_limit: usize,
    review_limit: usize,
    cost_limit_minutes: f32,
    population_size: usize,
    generations: usize,
    seed: Option<u64>,
    sigma0: f32,
}

impl Default for ExampleConfig {
    fn default() -> Self {
        let training_config = CostAdrTrainingConfig::default();
        Self {
            days: 1_825,
            deck_size: 10_000,
            learn_limit: 10,
            review_limit: 9_999,
            cost_limit_minutes: 720.0,
            population_size: training_config.population_size,
            generations: training_config.generations,
            seed: training_config.seed,
            sigma0: training_config.sigma0,
        }
    }
}

fn invalid_arg(message: impl Into<String>) -> Box<dyn Error> {
    Box::new(IoError::new(ErrorKind::InvalidInput, message.into()))
}

fn parse_value<T>(flag: &str, value: &str) -> Result<T, Box<dyn Error>>
where
    T: FromStr,
    T::Err: Display,
{
    value
        .parse()
        .map_err(|err| invalid_arg(format!("invalid value for {flag}: {value} ({err})")))
}

fn next_arg_value<I>(args: &mut I, flag: &str) -> Result<String, Box<dyn Error>>
where
    I: Iterator<Item = String>,
{
    args.next()
        .ok_or_else(|| invalid_arg(format!("missing value for {flag}")))
}

fn arg_value<I>(
    args: &mut I,
    flag: &str,
    inline_value: Option<String>,
) -> Result<String, Box<dyn Error>>
where
    I: Iterator<Item = String>,
{
    match inline_value {
        Some(value) => Ok(value),
        None => next_arg_value(args, flag),
    }
}

fn seed_with_offset(seed: Option<u64>, offset: u64) -> Option<u64> {
    seed.map(|seed| seed.wrapping_add(offset))
}

fn parse_args() -> Result<Option<ExampleConfig>, Box<dyn Error>> {
    let mut config = ExampleConfig::default();
    let mut args = env::args().skip(1);
    while let Some(arg) = args.next() {
        if arg == "-h" || arg == "--help" {
            print!("{USAGE}");
            return Ok(None);
        }

        let (flag, inline_value) = if let Some((flag, value)) = arg.split_once('=') {
            (flag, Some(value.to_string()))
        } else {
            (arg.as_str(), None)
        };

        match flag {
            "--days" => {
                config.days = parse_value(flag, &arg_value(&mut args, flag, inline_value)?)?
            }
            "--deck" => {
                config.deck_size = parse_value(flag, &arg_value(&mut args, flag, inline_value)?)?
            }
            "--learn-limit" => {
                config.learn_limit = parse_value(flag, &arg_value(&mut args, flag, inline_value)?)?
            }
            "--review-limit" => {
                config.review_limit = parse_value(flag, &arg_value(&mut args, flag, inline_value)?)?
            }
            "--cost-limit-minutes" => {
                config.cost_limit_minutes =
                    parse_value(flag, &arg_value(&mut args, flag, inline_value)?)?
            }
            "--pop" => {
                config.population_size =
                    parse_value(flag, &arg_value(&mut args, flag, inline_value)?)?
            }
            "--gen" => {
                config.generations = parse_value(flag, &arg_value(&mut args, flag, inline_value)?)?
            }
            "--seed" => {
                config.seed = Some(parse_value(
                    flag,
                    &arg_value(&mut args, flag, inline_value)?,
                )?)
            }
            "--sigma0" => {
                config.sigma0 = parse_value(flag, &arg_value(&mut args, flag, inline_value)?)?
            }
            _ => return Err(invalid_arg(format!("unknown argument: {flag}"))),
        }
    }
    Ok(Some(config))
}

fn format_optional(value: Option<f32>) -> String {
    value
        .map(|value| format!("{value:.6}"))
        .unwrap_or_else(|| "n/a".to_string())
}

fn spawn_progress_printer(progress: Arc<Mutex<CombinedProgressState>>) -> thread::JoinHandle<()> {
    thread::spawn(move || {
        let mut stderr = std::io::stderr();
        loop {
            let (current, total, finished) = {
                let progress = progress.lock().unwrap();
                (progress.current(), progress.total(), progress.finished())
            };
            eprint!("\r{}", progress_line(current, total, finished));
            let _ = stderr.flush();

            if finished {
                eprintln!();
                break;
            }

            thread::sleep(StdDuration::from_millis(100));
        }
    })
}

fn progress_line(current: usize, total: usize, finished: bool) -> String {
    if total == 0 {
        return "Cost ADR training [initializing]".to_string();
    }

    const WIDTH: usize = 32;
    let current = current.min(total);
    let filled = ((current * WIDTH) / total).min(WIDTH);
    let bar = format!("{}{}", "=".repeat(filled), " ".repeat(WIDTH - filled));
    let percent = current as f32 * 100.0 / total as f32;
    let status = if finished { "done" } else { "training" };
    format!("Cost ADR {status} [{bar}] {current}/{total} ({percent:.1}%)")
}

fn print_default_policy_evaluation(result: &CostAdrEvaluationResult) {
    println!(
        "default_policy_baseline_hypervolume={:.6} default_policy_hypervolume={:.6} default_policy_hypervolume_delta={:.6}",
        result.baseline_hypervolume, result.scheduler_hypervolume, result.hypervolume_delta
    );
    println!(
        "default_policy_span_coverage_percent={:.3} covered_span={:.3} total_span={:.3} covered_targets={}/{} baseline_frontier={} scheduler_frontier={}",
        result.auc_metrics.span_coverage_percent,
        result.auc_metrics.covered_span,
        result.auc_metrics.total_span,
        result.auc_metrics.covered_target_count,
        result.auc_metrics.target_count,
        result.auc_metrics.baseline_frontier_count,
        result.auc_metrics.scheduler_frontier_count
    );
    println!(
        "default_policy_same_target_time_saved_auc={} baseline_time_auc={} relative_same_target_time_saved_auc_percent={}",
        format_optional(result.auc_metrics.same_target_time_saved_auc),
        format_optional(result.auc_metrics.baseline_time_auc),
        format_optional(
            result
                .auc_metrics
                .relative_same_target_time_saved_auc_percent
        )
    );
    println!("default policy cost-weight rollout points:");
    for point in &result.scheduler_metrics {
        println!(
            "  w={:<7.1} avg_dr={:>8} memorized_avg={:>9.3} time_avg_min={:>7.3} mem_per_min={:>9.3} reviews={} lapses={}",
            point.goal_cost_weight,
            format_optional(point.average_desired_retention),
            point.metrics.memorized_average,
            point.metrics.time_average,
            point.metrics.memorized_per_minute,
            point.metrics.total_reviews,
            point.metrics.total_lapses
        );
    }
}

fn print_cost_adr_simulation(label: &str, result: &SimulationResult) {
    let total_cost = result.cost_per_day.iter().sum::<f32>();
    let time_average = if result.cost_per_day.is_empty() {
        0.0
    } else {
        total_cost / result.cost_per_day.len() as f32 / 60.0
    };
    let memorized_average = if result.memorized_cnt_per_day.is_empty() {
        0.0
    } else {
        result.memorized_cnt_per_day.iter().sum::<f32>() / result.memorized_cnt_per_day.len() as f32
    };
    let total_reviews = result.review_cnt_per_day.iter().sum::<usize>()
        + result.learn_cnt_per_day.iter().sum::<usize>();
    let total_lapses = result.cards.iter().map(|card| card.lapses).sum::<u32>();

    println!(
        "{label} avg_dr={} memorized_avg={:.3} time_avg_min={:.3} mem_per_min={:.3} reviews={} lapses={} total_cost_seconds={:.3}",
        format_optional(result.average_desired_retention),
        memorized_average,
        time_average,
        if time_average > 0.0 {
            memorized_average / time_average
        } else {
            0.0
        },
        total_reviews,
        total_lapses,
        total_cost
    );
}

fn main() -> fsrs::Result<()> {
    let Some(example_config) = parse_args().map_err(|err| {
        eprintln!("{err}");
        eprintln!();
        eprint!("{USAGE}");
        FSRSError::InvalidInput
    })?
    else {
        return Ok(());
    };

    let config = SimulatorConfig {
        deck_size: example_config.deck_size,
        learn_span: example_config.days,
        learn_limit: example_config.learn_limit,
        review_limit: example_config.review_limit,
        max_cost_perday: example_config.cost_limit_minutes * 60.0,
        ..Default::default()
    };
    let progress = CombinedProgressState::new_shared();
    let training_config = CostAdrTrainingConfig {
        population_size: example_config.population_size,
        generations: example_config.generations,
        sigma0: example_config.sigma0,
        seed: example_config.seed,
        simulation_seed: None,
        progress: Some(progress.clone()),
        ..Default::default()
    };

    let default_policy = CostAdrPolicy::new(training_config.initial_coefficients.clone())?;
    let evaluation_config = CostAdrEvaluationConfig {
        seed: seed_with_offset(example_config.seed, EXAMPLE_EVALUATION_SEED_OFFSET),
        ..Default::default()
    };
    println!("Cost ADR default policy evaluation");
    let started = Instant::now();
    let default_evaluation =
        default_policy.evaluate(&config, &DEFAULT_PARAMETERS, &evaluation_config)?;
    println!(
        "default_policy_duration_seconds={:.3}",
        started.elapsed().as_secs_f32()
    );
    print_default_policy_evaluation(&default_evaluation);

    // For a production user, replace DEFAULT_PARAMETERS with parameters trained
    // from that user's revlog via compute_parameters.
    let progress_printer = spawn_progress_printer(progress.clone());
    let started = Instant::now();
    let result = CostAdrPolicy::train_single_user(&config, &DEFAULT_PARAMETERS, &training_config);
    let wall_seconds = started.elapsed().as_secs_f32();
    let _ = progress_printer.join();
    let result = result?;

    println!("Cost ADR single-user FSRS training");
    println!(
        "config days={} deck={} learn_limit={} review_limit={} cost_limit_minutes={} pop={} gen={} sigma0={} base_seed={}",
        example_config.days,
        example_config.deck_size,
        example_config.learn_limit,
        example_config.review_limit,
        example_config.cost_limit_minutes,
        example_config.population_size,
        example_config.generations,
        example_config.sigma0,
        example_config.seed.unwrap_or(EXAMPLE_DEFAULT_BASE_SEED)
    );
    println!(
        "duration_seconds={:.3} result_training_seconds={:.3}",
        wall_seconds, result.training_seconds
    );
    println!(
        "baseline_hypervolume={:.6} best_hypervolume={:.6} best_hypervolume_delta={:.6}",
        result.baseline_hypervolume, result.best_hypervolume, result.best_hypervolume_delta
    );
    println!(
        "span_coverage_percent={:.3} covered_span={:.3} total_span={:.3} covered_targets={}/{} baseline_frontier={} scheduler_frontier={}",
        result.best_auc_metrics.span_coverage_percent,
        result.best_auc_metrics.covered_span,
        result.best_auc_metrics.total_span,
        result.best_auc_metrics.covered_target_count,
        result.best_auc_metrics.target_count,
        result.best_auc_metrics.baseline_frontier_count,
        result.best_auc_metrics.scheduler_frontier_count
    );
    println!(
        "same_target_time_saved_auc={} baseline_time_auc={} relative_same_target_time_saved_auc_percent={}",
        format_optional(result.best_auc_metrics.same_target_time_saved_auc),
        format_optional(result.best_auc_metrics.baseline_time_auc),
        format_optional(
            result
                .best_auc_metrics
                .relative_same_target_time_saved_auc_percent
        )
    );
    if let Some(last) = result.history.last() {
        println!(
            "last_generation={} best_delta={:.6} generation_best_delta={:.6} mean_delta={:.6} sigma={:.6}",
            last.generation,
            last.best_hypervolume_delta,
            last.generation_best_hypervolume_delta,
            last.mean_hypervolume_delta,
            last.sigma
        );
    }
    println!("selected cost-weight rollout points:");
    for point in &result.best_cost_weight_metrics {
        println!(
            "  w={:<7.1} avg_dr={:>8} memorized_avg={:>9.3} time_avg_min={:>7.3} mem_per_min={:>9.3} reviews={} lapses={}",
            point.goal_cost_weight,
            format_optional(point.average_desired_retention),
            point.metrics.memorized_average,
            point.metrics.time_average,
            point.metrics.memorized_per_minute,
            point.metrics.total_reviews,
            point.metrics.total_lapses
        );
    }

    // In production, persist the policy with the user's FSRS parameters.
    // CostAdrPolicy derives serde Serialize/Deserialize.
    let mut user_policy = result.policy.clone();
    user_policy.max_interval_days = Some(config.max_ivl);
    println!(
        "persist policy coefficient_count={}",
        user_policy.coefficients.len()
    );
    println!("policy={user_policy:#?}");

    let rollout_cost_weight = training_config.cost_weights[0];
    let rollout_seed = Some(
        example_config
            .seed
            .unwrap_or(EXAMPLE_DEFAULT_BASE_SEED)
            .wrapping_add(EXAMPLE_FINAL_ROLLOUT_SEED_OFFSET),
    );
    let rollout = simulate_with_cost_adr_policy(
        &config,
        &DEFAULT_PARAMETERS,
        &user_policy,
        rollout_cost_weight,
        rollout_seed,
        None,
    )?;
    print_cost_adr_simulation(
        &format!("simulate_with_cost_adr_policy w={rollout_cost_weight:.1}"),
        &rollout,
    );

    let fsrs = FSRS::new(&DEFAULT_PARAMETERS)?;
    let previous_state = fsrs.next_states(None, 0.9, 0)?.good.memory;
    println!("\nruntime schedule by cost weight:");
    for &cost_weight in &training_config.cost_weights {
        let next_states = user_policy.next_states(&fsrs, Some(previous_state), cost_weight, 7)?;
        println!("  w={cost_weight:.1}");
        for (rating, scheduled) in [
            ("Again", &next_states.again),
            ("Hard", &next_states.hard),
            ("Good", &next_states.good),
            ("Easy", &next_states.easy),
        ] {
            let interval_days = scheduled.interval.round().max(1.0) as u32;
            println!(
                "    rating={} stability={:.3} difficulty={:.3} desired_retention={:.6} interval_days={}",
                rating,
                scheduled.memory.stability,
                scheduled.memory.difficulty,
                scheduled.desired_retention,
                interval_days
            );
        }
    }

    println!("\nconsecutive Good interval sequence by cost weight:");
    for &cost_weight in &training_config.cost_weights {
        println!("  w={cost_weight:.1}");
        let mut state = None;
        let mut days_elapsed = 0;
        for i in 1..=10 {
            let next_states = user_policy.next_states(&fsrs, state, cost_weight, days_elapsed)?;
            let good = &next_states.good;
            let interval_days = good.interval.round().max(1.0) as u32;
            println!(
                "    review={} stability={:.3} difficulty={:.3} desired_retention={:.6} interval_days={}",
                i,
                good.memory.stability,
                good.memory.difficulty,
                good.desired_retention,
                interval_days
            );
            state = Some(good.memory);
            days_elapsed = interval_days;
        }
    }

    Ok(())
}