taktwerk 0.6.0

//! A campaign file and the runs it plans: the product of its axes (cases, grid, samples,
//! replicates), each run one effective scenario with its own seed. Pure: no runs, no files
//! written.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};
use taktwerk_core::rng::derive_index;
use taktwerk_core::scenario::{EventSpec, ScenarioFile, StopSpec};

use super::sample::{Dist, Method, probabilities};

/// Run seeds keep 63 bits: TOML integers are signed 64-bit.
const SEED_MASK: u64 = i64::MAX as u64;

/// A campaign file.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct CampaignFile {
    /// The project every run runs, relative to the campaign file.
    pub project: PathBuf,
    /// Output directory, relative to the campaign file; default: the file's stem beside it.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub out: Option<PathBuf>,
    /// Runs at a time.
    #[serde(default = "one")]
    pub jobs: usize,
    /// `--speed` of every run: `max` (the default) or a factor.
    #[serde(default = "max")]
    pub speed: String,
    /// The campaign's seed; each run's seed derives from it.
    #[serde(default)]
    pub seed: u64,
    /// A scenario file every run starts from, relative to the campaign file.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub scenario: Option<PathBuf>,
    /// Overrides of every run.
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub set: BTreeMap<String, toml::Value>,
    /// Stops of every run.
    #[serde(default)]
    pub stop: StopSpec,
    /// Events of every run.
    #[serde(default, rename = "event", skip_serializing_if = "Vec::is_empty")]
    pub events: Vec<EventSpec>,
    /// The list axis.
    #[serde(default, rename = "case", skip_serializing_if = "Vec::is_empty")]
    pub cases: Vec<Case>,
    /// The grid axis: key → values, every combination.
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub grid: BTreeMap<String, Vec<toml::Value>>,
    /// The random axis.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub sample: Option<Sample>,
    /// Repeats of every run with another seed.
    #[serde(default = "one")]
    pub replicates: usize,
}

fn one() -> usize {
    1
}

fn max() -> String {
    "max".to_owned()
}

/// One `[[case]]`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Case {
    /// Unique; ASCII letters, digits, `-` and `_`.
    pub name: String,
    /// A scenario file, relative to the campaign file, merged over the common one.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub scenario: Option<PathBuf>,
    /// Overrides of this case.
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub set: BTreeMap<String, toml::Value>,
}

/// `[sample]`: `n` draws of every variable.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Sample {
    /// Draws.
    pub n: usize,
    /// How.
    #[serde(default)]
    pub method: Method,
    /// Key → distribution.
    #[serde(flatten)]
    pub vars: BTreeMap<String, Dist>,
}

/// One planned run.
#[derive(Debug, Clone, PartialEq)]
pub struct RunSpec {
    /// Its name and directory.
    pub name: String,
    /// Its case, if the campaign has cases.
    pub case: Option<String>,
    /// Its grid point, sample and replicate, if the campaign has those axes.
    pub grid: Option<usize>,
    /// Its sample.
    pub sample: Option<usize>,
    /// Its replicate.
    pub replicate: Option<usize>,
    /// The values its axes set, in axis order (case, grid, sample).
    pub values: Vec<(String, toml::Value)>,
    /// Its seed.
    pub seed: u64,
    /// The scenario it runs: common, case, axis values and seed merged.
    pub scenario: ScenarioFile,
}

/// Read scenario files relative to a directory.
pub trait Files {
    /// The scenario file at `path`.
    ///
    /// # Errors
    /// It cannot be read or parsed.
    fn scenario(&self, path: &Path) -> Result<ScenarioFile, String>;
}

/// Plan the runs of `file`, whose relative paths resolve against `dir`.
///
/// Runs are the product of the axes present, in the order case, grid point, sample,
/// replicate. A run's seed is `derive_index(seed, sample · replicates + replicate)`, 63 bits
/// of it: runs that
/// differ only in case or grid point share it (common random numbers), so a comparison between
/// cases or gains sees the same noise.
///
/// # Errors
/// A bad axis, name or distribution, a scenario file that cannot be read, or runs that would
/// not end.
pub fn plan(file: &CampaignFile, dir: &Path, files: &dyn Files) -> Result<Vec<RunSpec>, String> {
    if file.jobs == 0 || file.replicates == 0 {
        return Err("jobs and replicates must be at least 1".into());
    }
    let mut common = match &file.scenario {
        Some(p) => absolute_models(files.scenario(&dir.join(p))?, &dir.join(p)),
        None => ScenarioFile::default(),
    };
    merge(
        &mut common,
        ScenarioFile {
            set: file.set.clone(),
            stop: file.stop.clone(),
            events: file.events.clone(),
            ..ScenarioFile::default()
        },
    );
    let mut cases: Vec<(Option<&Case>, ScenarioFile)> = Vec::new();
    if file.cases.is_empty() {
        cases.push((None, common.clone()));
    }
    for c in &file.cases {
        if c.name.is_empty()
            || !c
                .name
                .bytes()
                .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')
        {
            return Err(format!(
                "case `{}`: a name is ASCII letters, digits, `-` and `_`",
                c.name
            ));
        }
        if file.cases.iter().filter(|o| o.name == c.name).count() > 1 {
            return Err(format!("case `{}` is named twice", c.name));
        }
        let mut s = common.clone();
        if let Some(p) = &c.scenario {
            let path = dir.join(p);
            merge(&mut s, absolute_models(files.scenario(&path)?, &path));
        }
        merge(
            &mut s,
            ScenarioFile {
                set: c.set.clone(),
                ..ScenarioFile::default()
            },
        );
        cases.push((Some(c), s));
    }
    for (key, values) in &file.grid {
        if values.is_empty() {
            return Err(format!("grid `{key}` has no values"));
        }
    }
    let grid = grid_points(&file.grid);
    let samples: Vec<Vec<(String, toml::Value)>> = match &file.sample {
        None => vec![Vec::new()],
        Some(s) => {
            if s.n == 0 || s.vars.is_empty() {
                return Err("[sample] needs n >= 1 and at least one variable".into());
            }
            let mut draws = vec![Vec::new(); s.n];
            for (key, dist) in &s.vars {
                dist.check().map_err(|e| format!("sample `{key}`: {e}"))?;
                for (i, p) in probabilities(s.method, s.n, file.seed, key)
                    .into_iter()
                    .enumerate()
                {
                    draws[i].push((key.clone(), dist.at(p)));
                }
            }
            draws
        }
    };
    let width = |n: usize| (n.max(2) - 1).to_string().len();
    let (gw, sw, rw) = (
        width(grid.len()),
        width(samples.len()),
        width(file.replicates),
    );
    let mut runs = Vec::new();
    for (case, base) in &cases {
        for (g, point) in grid.iter().enumerate() {
            for (s, drawn) in samples.iter().enumerate() {
                for r in 0..file.replicates {
                    let mut parts: Vec<String> = Vec::new();
                    if let Some(c) = case {
                        parts.push(c.name.clone());
                    }
                    if !file.grid.is_empty() {
                        parts.push(format!("g{g:0gw$}"));
                    }
                    if file.sample.is_some() {
                        parts.push(format!("s{s:0sw$}"));
                    }
                    if file.replicates > 1 {
                        parts.push(format!("r{r:0rw$}"));
                    }
                    let name = if parts.is_empty() {
                        "run".to_owned()
                    } else {
                        parts.join("-")
                    };
                    let mut values: Vec<(String, toml::Value)> = case
                        .map(|c| c.set.clone().into_iter().collect())
                        .unwrap_or_default();
                    values.extend(point.iter().cloned());
                    values.extend(drawn.iter().cloned());
                    let mut scenario = base.clone();
                    // 63 bits, so the seed is a TOML integer in the run's scenario.
                    let mut seed =
                        derive_index(file.seed, (s * file.replicates + r) as u64) & SEED_MASK;
                    for (k, v) in point.iter().chain(drawn) {
                        if k == "seed" {
                            seed = v
                                .as_integer()
                                .and_then(|i| u64::try_from(i).ok())
                                .ok_or_else(|| format!("`seed` = {v} is not a seed"))?;
                        } else {
                            scenario.set.insert(k.clone(), v.clone());
                        }
                    }
                    scenario.seed = Some(seed);
                    if scenario.stop.cycles.is_none() && scenario.stop.time_s.is_none() {
                        return Err(format!(
                            "run `{name}` has no stop limit; give [stop] cycles or time_s, so \
                             every run ends"
                        ));
                    }
                    runs.push(RunSpec {
                        name,
                        case: case.map(|c| c.name.clone()),
                        grid: (!file.grid.is_empty()).then_some(g),
                        sample: file.sample.is_some().then_some(s),
                        replicate: (file.replicates > 1).then_some(r),
                        values,
                        seed,
                        scenario,
                    });
                }
            }
        }
    }
    Ok(runs)
}

/// Every combination of the grid's values, keys in order, the last key varying fastest.
fn grid_points(grid: &BTreeMap<String, Vec<toml::Value>>) -> Vec<Vec<(String, toml::Value)>> {
    let mut points = vec![Vec::new()];
    for (key, values) in grid {
        points = points
            .into_iter()
            .flat_map(|p| {
                values.iter().map(move |v| {
                    let mut q = p.clone();
                    q.push((key.clone(), v.clone()));
                    q
                })
            })
            .collect();
    }
    points
}

/// Merge `over` into `base`: overrides and models replace by key, a stop limit given replaces,
/// conditions and events add up, a seed given replaces.
pub fn merge(base: &mut ScenarioFile, over: ScenarioFile) {
    base.set.extend(over.set);
    base.models.extend(over.models);
    base.stop.cycles = over.stop.cycles.or(base.stop.cycles);
    base.stop.time_s = over.stop.time_s.or(base.stop.time_s);
    base.stop.when.extend(over.stop.when);
    base.events.extend(over.events);
    base.seed = over.seed.or(base.seed);
}

/// `file` read from `path`, its model paths made relative to nothing: joined to its directory.
fn absolute_models(mut file: ScenarioFile, path: &Path) -> ScenarioFile {
    let dir = path.parent().unwrap_or_else(|| Path::new(""));
    for p in file.models.values_mut() {
        *p = dir.join(&*p);
    }
    file
}

#[cfg(test)]
#[allow(clippy::unwrap_used, reason = "tests")]
mod tests {
    use super::*;

    struct NoFiles(BTreeMap<PathBuf, &'static str>);

    impl Files for NoFiles {
        fn scenario(&self, path: &Path) -> Result<ScenarioFile, String> {
            let text = self.0.get(path).ok_or(format!("no {}", path.display()))?;
            toml::from_str(text).map_err(|e| e.to_string())
        }
    }

    fn planned(text: &str) -> Result<Vec<RunSpec>, String> {
        let file: CampaignFile = toml::from_str(text).map_err(|e| e.to_string())?;
        let files = NoFiles(BTreeMap::from([(
            PathBuf::from("/c/slow.toml"),
            "[set]\n\"plant.tau\" = 2.0\n[models]\nplant = \"pkg\"\n[[event]]\ncycle = 5\nset = \"x\"\nvalue = 1\n",
        )]));
        plan(&file, Path::new("/c"), &files)
    }

    #[test]
    fn runs_are_the_product_of_the_axes() {
        let runs = planned(
            r#"
            project = "p.toml"
            seed = 1
            replicates = 2
            [stop]
            time_s = 60.0
            [set]
            "ctrl.sp" = 1.0
            [[case]]
            name = "nominal"
            [[case]]
            name = "slow"
            scenario = "slow.toml"
            set = { "ctrl.ki" = 0.1 }
            [grid]
            "ctrl.kp" = [0.5, 1.0, 2.0]
            [sample]
            n = 4
            "plant.gain" = { uniform = [0.8, 1.2] }
            "#,
        )
        .unwrap();
        assert_eq!(runs.len(), 2 * 3 * 4 * 2);
        assert!(runs.iter().all(|r| r.seed <= i64::MAX as u64));
        assert_eq!(runs[0].name, "nominal-g0-s0-r0");
        assert_eq!(runs.last().unwrap().name, "slow-g2-s3-r1");
        let slow = runs.iter().find(|r| r.name == "slow-g1-s2-r0").unwrap();
        let s = &slow.scenario;
        assert_eq!(s.set["ctrl.sp"], 1.0.into());
        assert_eq!(s.set["plant.tau"], 2.0.into());
        assert_eq!(s.set["ctrl.ki"], 0.1.into());
        assert_eq!(s.set["ctrl.kp"], 1.0.into());
        assert!(s.set.contains_key("plant.gain"));
        assert_eq!(s.models["plant"], PathBuf::from("/c/pkg"));
        assert_eq!(s.events.len(), 1);
        assert_eq!(s.stop.time_s, Some(60.0));
        assert_eq!(s.seed, Some(slow.seed));
        let keys: Vec<&str> = slow.values.iter().map(|(k, _)| k.as_str()).collect();
        assert_eq!(keys, ["ctrl.ki", "ctrl.kp", "plant.gain"]);
        // Common random numbers: case and grid point share the seed, samples and replicates
        // do not.
        let seed = |n: &str| runs.iter().find(|r| r.name == n).unwrap().seed;
        assert_eq!(seed("nominal-g0-s2-r1"), seed("slow-g2-s2-r1"));
        assert_ne!(seed("nominal-g0-s2-r1"), seed("nominal-g0-s2-r0"));
        assert_ne!(seed("nominal-g0-s2-r1"), seed("nominal-g0-s3-r1"));
        // The same file plans the same runs.
        assert_eq!(runs, planned_again());
    }

    fn planned_again() -> Vec<RunSpec> {
        planned(
            r#"
            project = "p.toml"
            seed = 1
            replicates = 2
            [stop]
            time_s = 60.0
            [set]
            "ctrl.sp" = 1.0
            [[case]]
            name = "nominal"
            [[case]]
            name = "slow"
            scenario = "slow.toml"
            set = { "ctrl.ki" = 0.1 }
            [grid]
            "ctrl.kp" = [0.5, 1.0, 2.0]
            [sample]
            n = 4
            "plant.gain" = { uniform = [0.8, 1.2] }
            "#,
        )
        .unwrap()
    }

    #[test]
    fn a_grid_alone_names_runs_by_point_and_a_seed_axis_sets_the_seed() {
        let runs = planned(
            "project = \"p.toml\"\n[stop]\ncycles = 10\n[grid]\nseed = [7, 8]\n\"a.k\" = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]\n",
        )
        .unwrap();
        assert_eq!(runs.len(), 22);
        assert_eq!(runs[0].name, "g00");
        assert_eq!(runs[21].name, "g21");
        assert_eq!(runs[0].seed, 7);
        assert_eq!(runs[21].seed, 8);
        assert!(!runs[0].scenario.set.contains_key("seed"));
        let one = planned("project = \"p.toml\"\n[stop]\ncycles = 10\n").unwrap();
        assert_eq!(one.len(), 1);
        assert_eq!(one[0].name, "run");
    }

    #[test]
    fn a_bad_campaign_says_why() {
        for (text, says) in [
            ("project = \"p\"\n[grid]\nk = [1]\n", "no stop limit"),
            (
                "project = \"p\"\njobs = 0\n[stop]\ncycles = 1\n",
                "at least 1",
            ),
            (
                "project = \"p\"\n[stop]\ncycles = 1\n[[case]]\nname = \"a b\"\n",
                "a name is",
            ),
            (
                "project = \"p\"\n[stop]\ncycles = 1\n[[case]]\nname = \"a\"\n[[case]]\nname = \"a\"\n",
                "named twice",
            ),
            (
                "project = \"p\"\n[stop]\ncycles = 1\n[grid]\nk = []\n",
                "has no values",
            ),
            (
                "project = \"p\"\n[stop]\ncycles = 1\n[sample]\nn = 3\nk = { log_uniform = [0.0, 1.0] }\n",
                "sample `k`",
            ),
            (
                "project = \"p\"\n[stop]\ncycles = 1\n[[case]]\nname = \"a\"\nscenario = \"x.toml\"\n",
                "no /c/x.toml",
            ),
            (
                "project = \"p\"\n[stop]\ncycles = 1\n[grid]\nseed = [-1]\n",
                "is not a seed",
            ),
        ] {
            let e = planned(text).unwrap_err();
            assert!(e.contains(says), "{text}: {e}");
        }
    }
}