AI Flow: Reverse Engineering

//! Photo intake for a set nobody sorted: `init --photos <folder or files>` records the photos as they came, and `stitch`
//! (run by `prep` when it has not been run) sorts them into faces, close-up tiles and overviews, picks the top, and
//! builds one mosaic per face at the close-ups' own resolution. The image registration is Python with OpenCV: the script
//! is `tools/stitch.py`, compiled into the binary, written to `reverse-flow/stitch/stitch.py` and run from there, and
//! the call (command line, script hash, exit status) goes into the run ledger. Its results land where the flow already
//! looks: the mosaics become the run's photos (`photos/top.jpg`, `photos/bottom.jpg`), so `prep`, `scale` and every later
//! step run on them unchanged. `stitch.json` and the evidence pictures say what was done and what is missing.
use crate::ledger;
use serde_json::{json, Value};
use std::path::{Path, PathBuf};

/// The stitcher, as shipped in tools/.
pub const SCRIPT: &str = include_str!("../tools/stitch.py");
/// Where the unsorted photos are kept, and where the stitcher works, inside the flow folder.
pub const RAW_DIR: &str = "photos-raw";
pub const WORK_DIR: &str = "stitch";

fn is_photo(p: &Path) -> bool {
    matches!(p.extension().and_then(|e| e.to_str()).map(|e| e.to_ascii_lowercase()).as_deref(), Some("jpg") | Some("jpeg") | Some("png"))
}

/// A file name split into text and number runs, so that Board2.10 sorts after Board2.9.
fn natural_key(name: &str) -> Vec<(String, u64)> {
    let mut out = vec![]; let (mut text, mut num, mut in_num) = (String::new(), String::new(), false);
    for c in name.chars() {
        if c.is_ascii_digit() { in_num = true; num.push(c); }
        else { if in_num { out.push((std::mem::take(&mut text), num.parse().unwrap_or(0))); num.clear(); in_num = false; } text.push(c.to_ascii_lowercase()); }
    }
    out.push((text, num.parse().unwrap_or(0)));
    out
}

/// The photos behind `--photos`: each argument is a photo or a folder of photos (jpg, jpeg, png; not searched deeper).
/// Sorted by name with numbers in order, no file twice.
pub fn collect(args: &[String]) -> Result<Vec<PathBuf>, String> {
    let mut out: Vec<PathBuf> = vec![];
    for a in args {
        let p = Path::new(a);
        if p.is_dir() {
            let mut found: Vec<PathBuf> = std::fs::read_dir(p).map_err(|e| format!("{a}: {e}"))?.flatten().map(|e| e.path()).filter(|q| q.is_file() && is_photo(q)).collect();
            if found.is_empty() { return Err(format!("{a}: no jpg, jpeg or png photos in this folder")); }
            out.append(&mut found);
        } else if p.is_file() {
            if !is_photo(p) { return Err(format!("{a}: not a jpg, jpeg or png")); }
            out.push(p.to_path_buf());
        } else { return Err(format!("{a}: not a file or folder")); }
    }
    out.sort_by_key(|p| natural_key(&p.file_name().map(|n| n.to_string_lossy().to_string()).unwrap_or_default()));
    out.dedup();
    let mut names = std::collections::HashSet::new();
    for p in &out { if !names.insert(p.file_name().map(|n| n.to_os_string())) { return Err(format!("two photos are both named {}: rename one", p.file_name().unwrap().to_string_lossy())); } }
    Ok(out)
}

/// `init --photos`: copy the photos as they came and return the state entry that says they are not sorted yet.
pub fn record(flow_dir: &Path, files: &[PathBuf]) -> Result<Value, String> {
    let raw = flow_dir.join(RAW_DIR);
    std::fs::create_dir_all(&raw).map_err(|e| e.to_string())?;
    let mut recs = vec![];
    for src in files {
        let name = src.file_name().unwrap().to_string_lossy().to_string();
        let dst = raw.join(&name);
        std::fs::copy(src, &dst).map_err(|e| format!("{}: {e}", src.display()))?;
        recs.push(json!({"file": name, "source": src.display().to_string(), "sha256": crate::sha256(&dst)}));
    }
    Ok(json!({"unsorted": recs, "stitched": false}))
}

/// The run was started with unsorted photos and they have not been stitched yet.
pub fn pending(state: &Value) -> bool {
    state["intake"]["unsorted"].as_array().map(|a| !a.is_empty()).unwrap_or(false) && state["intake"]["stitched"] != json!(true)
}

/// What `next` adds while the photos wait to be sorted.
pub fn hint(flow_dir: &Path) -> Vec<String> {
    let state: Value = std::fs::read_to_string(flow_dir.join("state.json")).ok().and_then(|t| serde_json::from_str(&t).ok()).unwrap_or(json!({}));
    if pending(&state) {
        vec![format!("{} photos are not sorted into faces yet: `prep` does it first (or run `stitch` alone), then read stitch.json and stitch/*-sources.jpg, *-seams.jpg", state["intake"]["unsorted"].as_array().map(|a| a.len()).unwrap_or(0))]
    } else { vec![] }
}

/// The run's photo list after stitching: one mosaic per face that was built.
pub fn photos_from_report(report: &Value) -> Vec<Value> {
    let mut out = vec![];
    for side in ["top", "bottom"] {
        let f = &report["faces"][side];
        if !f.is_object() { continue; }
        let from: Vec<Value> = f["tiles"].as_array().into_iter().flatten().map(|t| t["file"].clone()).collect();
        out.push(json!({"file": format!("{side}.jpg"), "side": side, "source": format!("{WORK_DIR}/{side}.jpg"), "mosaicOf": from, "overview": f["overview"]}));
    }
    out
}

/// One line per face for the hints, and what the report says is missing.
pub fn summary(report: &Value) -> Vec<String> {
    let mut h = vec![];
    for side in ["top", "bottom"] {
        let f = &report["faces"][side];
        if !f.is_object() { continue; }
        let n = f["tiles"].as_array().map(|a| a.len()).unwrap_or(0);
        let px = |v: &Value| v.as_f64().map(|x| format!("{x} px")).unwrap_or("not measured".into());
        h.push(format!("{side}: {n} close-up tile(s){}, {} x {} px; close-ups cover {}% of the board, overview fill (low resolution) {}%, not seen {}%; tiles disagree in their overlaps by {} at the median, {} at the 95th percentile ({} cells)",
            f["overview"].as_str().map(|o| format!(" and overview {o}")).unwrap_or_default(), f["size"][0], f["size"][1],
            f["coverage"]["closeUpsPct"], f["coverage"]["overviewLowResolutionPct"], f["coverage"]["notSeenPct"], px(&f["overlap"]["medPx"]), px(&f["overlap"]["p95Px"]), f["overlap"]["cells"]));
        for n in f["notes"].as_array().into_iter().flatten() { if let Some(s) = n.as_str() { h.push(format!("{side}: {s}")); } }
    }
    if let Some(d) = report["topDecision"].as_object() {
        h.push(format!("top chosen by: {}{}", d["rule"].as_str().unwrap_or(""), d.get("confidence").and_then(|c| c.as_str()).map(|c| format!(" ({c})")).unwrap_or_default()));
    }
    for u in report["unplaced"].as_array().into_iter().flatten() {
        h.push(format!("left out: {} ({})", u["photos"].as_array().into_iter().flatten().filter_map(|x| x.as_str()).collect::<Vec<_>>().join(", "), u["why"].as_str().unwrap_or("")));
    }
    for m in report["missing"].as_array().into_iter().flatten() { if let Some(s) = m.as_str() { h.push(format!("missing: {s}")); } }
    if let Some(s) = report["facesAssumed"].as_str() { h.push(format!("assumed: {s}")); }
    h
}

/// The board's outline in a mosaic, as `stitch` recorded it (stitch.json next to photos/, `boardBoxPx` of that face):
/// the four corners TL, TR, BR, BL in the mosaic's pixels. `prep` falls back on it when its own colour test finds no
/// green board, which is every red or blue one. It is a box round the outline cut from the overview, so it can sit a
/// millimetre or two outside a board that lies turned in the mosaic.
pub fn board_box(photo: &Path, side: &str) -> Option<[crate::geom::P; 4]> {
    let report = crate::items::read_json(&photo.parent()?.parent()?.join("stitch.json")).ok()?;
    let face = &report["faces"][side];
    if face["file"].as_str() != photo.file_name()?.to_str() { return None; }
    let b: Vec<f64> = face["boardBoxPx"].as_array()?.iter().filter_map(|v| v.as_f64()).collect();
    if b.len() != 4 || b[2] - b[0] < 16.0 || b[3] - b[1] < 16.0 { return None; }
    Some([(b[0], b[1]), (b[2], b[1]), (b[2], b[3]), (b[0], b[3])])
}

fn python() -> String { std::env::var("ADOM_REVERSE_PYTHON").unwrap_or_else(|_| "python3".into()) }

/// Sort and stitch the run's unsorted photos. Writes stitch/ (mosaics, stitch.json, evidence), copies the mosaics into
/// photos/, sets the run's photo list, and logs the call. Returns the hints for the caller's answer.
pub fn run(p: &mut ledger::Project, thread: &str) -> Result<Vec<String>, String> {
    let f = p.flow_dir.clone();
    let unsorted: Vec<String> = p.state["intake"]["unsorted"].as_array().into_iter().flatten().filter_map(|r| r["file"].as_str().map(str::to_string)).collect();
    if unsorted.is_empty() { return Err("this run was not started with `init --photos`: there is nothing to sort".into()); }
    let work = f.join(WORK_DIR);
    std::fs::create_dir_all(&work).map_err(|e| e.to_string())?;
    let script = work.join("stitch.py");
    std::fs::write(&script, SCRIPT).map_err(|e| format!("{}: {e}", script.display()))?;
    let mut cmd = std::process::Command::new(python());
    cmd.arg(&script).arg("--out").arg(&work);
    let top_photo = p.state["intake"]["topPhoto"].as_str().map(str::to_string);
    if let Some(t) = &top_photo { cmd.arg("--top-photo").arg(t); }
    for name in &unsorted { cmd.arg(f.join(RAW_DIR).join(name)); }
    cmd.env("STITCH_THREADS", "4");
    let line = format!("{} {}/stitch.py --out {}{} {}", python(), WORK_DIR, WORK_DIR, top_photo.as_ref().map(|t| format!(" --top-photo {t}")).unwrap_or_default(), unsorted.iter().map(|n| format!("{RAW_DIR}/{n}")).collect::<Vec<_>>().join(" "));
    let started = ledger::now();
    let out = cmd.output().map_err(|e| format!("could not start {}: {e}. The stitcher needs python3 with numpy and OpenCV (cv2 with SIFT); set ADOM_REVERSE_PYTHON to choose the interpreter", python()))?;
    std::fs::write(work.join("stitch.log"), [out.stdout.as_slice(), b"\n--- stderr ---\n", out.stderr.as_slice()].concat()).ok();
    let call = json!({"command": line, "cwd": "reverse-flow", "script": format!("{WORK_DIR}/stitch.py"), "scriptSha256": crate::sha256(&script), "started": started, "ended": ledger::now(),
                      "exit": out.status.code(), "log": format!("{WORK_DIR}/stitch.log")});
    if !out.status.success() {
        p.log(thread, "stitch", json!({"call": call, "ok": false}));
        let tail: Vec<String> = String::from_utf8_lossy(&out.stderr).lines().rev().take(4).map(str::to_string).collect::<Vec<_>>().into_iter().rev().collect();
        return Err(format!("the stitcher failed (exit {:?}): {}. Full output in {}/{WORK_DIR}/stitch.log", out.status.code(), tail.join(" | "), f.display()));
    }
    let report = crate::items::read_json(&work.join("stitch.json"))?;
    let mut photos = photos_from_report(&report);
    if photos.is_empty() { return Err("the stitcher built no face: see stitch/stitch.json".into()); }
    std::fs::create_dir_all(f.join("photos")).ok();
    for ph in photos.iter_mut() {
        let name = ph["file"].as_str().unwrap().to_string();
        let dst = f.join("photos").join(&name);
        std::fs::copy(work.join(&name), &dst).map_err(|e| format!("{name}: {e}"))?;
        ph["sha256"] = json!(crate::sha256(&dst));
    }
    std::fs::copy(work.join("stitch.json"), f.join("stitch.json")).map_err(|e| e.to_string())?;
    p.state["photos"] = json!(photos);
    p.state["intake"]["stitched"] = json!(true);
    p.state["intake"]["report"] = json!("stitch.json");
    p.save().map_err(|e| e.to_string())?;
    p.log(thread, "stitch", json!({"call": call, "ok": true, "photos": p.state["photos"], "topDecision": report["topDecision"], "unplaced": report["unplaced"]}));
    let mut h = summary(&report);
    h.push(format!("stitched by `{line}` (Python with OpenCV; the script is the one shipped in tools/stitch.py, its hash is in the ledger)"));
    h.push("Look before going on: stitch/top-sources.jpg and bottom-sources.jpg (which photo each part of the mosaic came from), stitch/*-seams.jpg (the joins at full size), stitch/*-overlap.jpg (where the tiles disagree).".into());
    if photos.len() == 1 { h.push("Only one face was built: bottom copper, via rings and layer evidence will be unknown. Ask for photos of the other face.".into()); }
    Ok(h)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn tmp(name: &str) -> PathBuf { let d = std::env::temp_dir().join(format!("intake-{name}-{}", std::process::id())); std::fs::remove_dir_all(&d).ok(); std::fs::create_dir_all(&d).unwrap(); d }

    #[test]
    fn photos_come_from_folders_and_lists_in_number_order() {
        let d = tmp("collect");
        for n in ["Board2.10.jpeg", "Board2.9.JPG", "Board2.1.JPG", "notes.txt", "Board2.11.png"] { std::fs::write(d.join(n), "x").unwrap(); }
        let got: Vec<String> = collect(&[d.display().to_string()]).unwrap().iter().map(|p| p.file_name().unwrap().to_string_lossy().to_string()).collect();
        assert_eq!(got, ["Board2.1.JPG", "Board2.9.JPG", "Board2.10.jpeg", "Board2.11.png"]);
        // a list of files gives the same, and a file named twice counts once
        let list: Vec<String> = ["Board2.10.jpeg", "Board2.1.JPG", "Board2.1.JPG"].iter().map(|n| d.join(n).display().to_string()).collect();
        assert_eq!(collect(&list).unwrap().len(), 2);
        assert!(collect(&[d.join("notes.txt").display().to_string()]).is_err());
        assert!(collect(&[d.join("missing.jpg").display().to_string()]).is_err());
        std::fs::remove_dir_all(&d).ok();
    }

    #[test]
    fn an_empty_folder_is_an_error_not_an_empty_run() {
        let d = tmp("empty");
        assert!(collect(&[d.display().to_string()]).unwrap_err().contains("no jpg"));
        std::fs::remove_dir_all(&d).ok();
    }

    #[test]
    fn unsorted_photos_wait_until_they_are_stitched() {
        assert!(!pending(&json!({"photos": [{"file": "a.jpg", "side": "top"}]})));        // a run started with --top
        let mut st = json!({"photos": [], "intake": {"unsorted": [{"file": "a.jpg"}], "stitched": false}});
        assert!(pending(&st));
        st["intake"]["stitched"] = json!(true);
        assert!(!pending(&st));
    }

    #[test]
    fn next_says_so_while_the_photos_wait() {
        let d = tmp("hint");
        std::fs::write(d.join("state.json"), json!({"intake": {"unsorted": [{"file": "a.jpg"}, {"file": "b.jpg"}], "stitched": false}}).to_string()).unwrap();
        let h = hint(&d);
        assert_eq!(h.len(), 1);
        assert!(h[0].starts_with("2 photos are not sorted"));
        std::fs::write(d.join("state.json"), json!({"photos": [{"file": "a.jpg", "side": "top"}]}).to_string()).unwrap();
        assert!(hint(&d).is_empty());
        std::fs::remove_dir_all(&d).ok();
    }

    #[test]
    fn the_mosaics_become_the_runs_photos_with_their_faces() {
        let rep = json!({"faces": {"bottom": {"tiles": [{"file": "b1.jpg"}], "overview": null}, "top": {"tiles": [{"file": "t1.jpg"}, {"file": "t2.jpg"}], "overview": "o.jpg"}}});
        let ph = photos_from_report(&rep);
        assert_eq!(ph.len(), 2);
        assert_eq!(ph[0]["file"], "top.jpg"); assert_eq!(ph[0]["side"], "top"); assert_eq!(ph[0]["mosaicOf"], json!(["t1.jpg", "t2.jpg"])); assert_eq!(ph[0]["overview"], "o.jpg");
        assert_eq!(ph[1]["file"], "bottom.jpg"); assert_eq!(ph[1]["side"], "bottom");
        // one face only: one photo, and it is the top
        let one = photos_from_report(&json!({"faces": {"top": {"tiles": [{"file": "t.jpg"}]}}}));
        assert_eq!(one.len(), 1); assert_eq!(one[0]["side"], "top");
    }

    #[test]
    fn the_summary_says_what_is_missing_and_never_hides_low_resolution() {
        let rep = json!({"faces": {"top": {"tiles": [{"file": "t1.jpg"}], "overview": null, "size": [4032, 3024],
            "coverage": {"closeUpsPct": 100.0, "overviewLowResolutionPct": 0.0, "notSeenPct": 0.0}, "overlap": {"cells": 0, "medPx": null, "p95Px": null},
            "notes": ["the board runs out of this photo on the left side(s)"]}},
            "topDecision": {"rule": "only one face was photographed: it is written as the top", "confidence": "none"}, "missing": ["no photo of the other face"],
            "unplaced": [{"photos": ["x.jpg"], "why": "matches no other photo"}]});
        let h = summary(&rep).join("\n");
        assert!(h.contains("not measured"));
        assert!(h.contains("runs out of this photo"));
        assert!(h.contains("missing: no photo of the other face"));
        assert!(h.contains("left out: x.jpg"));
    }

    #[test]
    fn stitch_is_the_first_thing_next_and_auto_run_on_unsorted_photos() {
        let d = tmp("auto");
        std::fs::write(d.join("state.json"), json!({"photos": [], "intake": {"unsorted": [{"file": "a.jpg"}], "stitched": false}}).to_string()).unwrap();
        let fl = crate::flow::load();
        let now = crate::auto::now_lines(&d, &fl);
        assert!(now[0].starts_with("now: `stitch`"), "{now:?}");
        // once stitched, the photos are measured
        std::fs::write(d.join("state.json"), json!({"photos": [{"file": "top.jpg", "side": "top"}], "intake": {"unsorted": [{"file": "a.jpg"}], "stitched": true}}).to_string()).unwrap();
        assert!(crate::auto::now_lines(&d, &fl)[0].starts_with("now: `prep`"));
        std::fs::remove_dir_all(&d).ok();
    }

    #[test]
    fn a_mosaic_hands_its_board_box_to_prep() {
        let d = tmp("box");
        std::fs::create_dir_all(d.join("photos")).unwrap();
        std::fs::write(d.join("stitch.json"), json!({"faces": {"top": {"file": "top.jpg", "boardBoxPx": [300, 200, 7000, 5000]}, "bottom": {"file": "bottom.jpg"}}}).to_string()).unwrap();
        let c = board_box(&d.join("photos/top.jpg"), "top").unwrap();
        assert_eq!(c, [(300.0, 200.0), (7000.0, 200.0), (7000.0, 5000.0), (300.0, 5000.0)]);
        assert!(board_box(&d.join("photos/bottom.jpg"), "bottom").is_none());       // no box recorded for that face
        assert!(board_box(&d.join("photos/other.jpg"), "top").is_none());           // not the mosaic
        std::fs::remove_dir_all(&d).ok();
    }

    #[test]
    fn the_shipped_script_is_the_one_compiled_in() {
        assert!(SCRIPT.starts_with("#!/usr/bin/env python3"));
        assert!(SCRIPT.contains("def main():"));
        // the measurement must copy each cell: OpenCV 5.0.0's phaseCorrelate writes into its inputs
        assert!(SCRIPT.contains(".copy(), gb[y:y + S, x:x + S].copy()"));
    }

    /// Whether this machine can run the stitcher at all.
    fn have_cv() -> bool {
        std::process::Command::new(python()).args(["-c", "import cv2, numpy; cv2.SIFT_create()"]).output().map(|o| o.status.success()).unwrap_or(false)
    }

    /// End to end on a made-up board: four overlapping tiles and a smaller overview of each face, no labels. The stitcher
    /// must put the tiles of each face together, call the face with the dark bodies the top, and agree in the overlaps.
    #[test]
    fn a_made_up_board_is_sorted_and_stitched_with_no_labels() {
        if !have_cv() { eprintln!("skipped: no python3 with OpenCV here"); return; }
        let d = tmp("e2e");
        let gen = d.join("gen.py");
        std::fs::write(&gen, include_str!("../tools/stitch_fixture.py")).unwrap();
        let o = std::process::Command::new(python()).arg(&gen).arg(d.join("in")).output().unwrap();
        assert!(o.status.success(), "{}", String::from_utf8_lossy(&o.stderr));
        let files = collect(&[d.join("in").display().to_string()]).unwrap();
        assert_eq!(files.len(), 10);
        let flow_dir = d.join("reverse-flow");
        std::fs::create_dir_all(&flow_dir).unwrap();
        let intake = record(&flow_dir, &files).unwrap();
        let mut p = ledger::Project::create(&d, "reverse-flow", json!({"photos": [], "intake": intake})).unwrap();
        assert!(pending(&p.state));
        let hints = run(&mut p, "test").unwrap_or_else(|e| panic!("{e}"));
        assert!(!pending(&p.state));
        assert!(hints.iter().any(|h| h.starts_with("top: 4 close-up tile(s)")), "{hints:?}");
        let rep = crate::items::read_json(&flow_dir.join("stitch.json")).unwrap();
        // the fixture's names say which face each tile is: the stitcher never saw that
        for (side, want) in [("top", "a"), ("bottom", "b")] {
            let f = &rep["faces"][side];
            let tiles: Vec<&str> = f["tiles"].as_array().unwrap().iter().map(|t| t["file"].as_str().unwrap()).collect();
            assert_eq!(tiles.len(), 4, "{side}: {tiles:?}");
            assert!(tiles.iter().all(|t| t.starts_with(want)), "{side}: {tiles:?}");
            assert!(f["overview"].as_str().unwrap().starts_with(want));
            assert!(f["overlap"]["cells"].as_u64().unwrap() > 20, "{side}: {}", f["overlap"]);
            assert!(f["overlap"]["medPx"].as_f64().unwrap() < 1.0, "{side}: {}", f["overlap"]);
            assert!(f["coverage"]["closeUpsPct"].as_f64().unwrap() > 90.0, "{side}: {}", f["coverage"]);
            assert!(flow_dir.join("photos").join(format!("{side}.jpg")).is_file());
        }
        assert_eq!(p.state["photos"].as_array().unwrap().len(), 2);
        let ledger_text = std::fs::read_to_string(flow_dir.join("run.jsonl")).unwrap();
        assert!(ledger_text.contains("\"ev\":\"stitch\"") && ledger_text.contains("stitch/stitch.py --out stitch"));
        std::fs::remove_dir_all(&d).ok();
    }
}