AI Flow: Reverse Engineering

//! `parts snap`: a part's measured position and rotation from its own footprint. The pads of the
//! library footprint are laid on the photo's metal (exposed pads, solder, leads) at every offset within
//! reach and every quarter turn; the pose where the pads cover the most metal and the ring just outside
//! them the least wins. A turn that scores the same as the AI's rotation (a symmetric pad pattern)
//! keeps the AI's rotation: pin 1 comes from the mark, not from the pads.
use crate::board::{to_px, RES};
use crate::raster::{self, Plane};
use crate::sexp;
use image::RgbImage;

#[derive(Clone, Debug)]
pub struct LPad { pub x: f64, pub y: f64, pub w: f64, pub h: f64, pub rot: f64, pub round: bool, pub drill: f64 }

/// Pads of a footprint in its own frame (mm).
pub fn pads(text: &str) -> Vec<LPad> {
    let Ok(n) = sexp::parse(text) else { return vec![] };
    n.children("pad").filter_map(|p| {
        let at = p.child("at")?; let size = p.child("size")?;
        let shape = p.arg(2).unwrap_or_default();
        Some(LPad { x: at.num(0)?, y: at.num(1)?, rot: at.num(2).unwrap_or(0.0), w: size.num(0)?, h: size.num(1)?, round: shape == "circle", drill: p.child("drill").and_then(|d| d.num(0)).unwrap_or(0.0) })
    }).filter(|p| p.w > 0.05 && p.h > 0.05).collect()
}

/// Metal likelihood per pixel (0..1): exposed pads, solder, gold, leads; lightly blurred.
pub fn metal_plane(img: &RgbImage) -> Plane {
    let tone = crate::prep::mask_tone(img);
    let m = Plane { w: img.width() as usize, h: img.height() as usize, d: img.pixels().map(|p| match crate::prep::material(p, &tone) {
        crate::prep::Mat::Metal | crate::prep::Mat::Gold => 1.0, crate::prep::Mat::Silk => 0.6, _ => 0.0 }).collect() };
    raster::blur(&m, 1, 2)
}

/// Sample points inside each pad (skipping its drill) and on a ring 0.25 mm outside it, in footprint mm.
fn probes(pads: &[LPad]) -> (Vec<(f64, f64)>, Vec<(f64, f64)>) {
    let (mut inside, mut ring) = (vec![], vec![]);
    for p in pads {
        let (c, s) = (p.rot.to_radians().cos(), p.rot.to_radians().sin());
        let put = |lx: f64, ly: f64| (p.x + lx * c + ly * s, p.y - lx * s + ly * c);
        let (hw, hh) = (p.w / 2.0, p.h / 2.0);
        let step = (p.w.min(p.h) / 4.0).clamp(0.05, 0.25);
        let mut u = -hw + step / 2.0;
        while u < hw { let mut v = -hh + step / 2.0;
            while v < hh {
                let inpad = if p.round { (u / hw).powi(2) + (v / hh).powi(2) <= 1.0 } else { true };
                if inpad && (u * u + v * v).sqrt() > p.drill / 2.0 + 0.05 { inside.push(put(u, v)); }
                v += step; }
            u += step; }
        let g = 0.25; let n = 16;
        for k in 0..n { let t = k as f64 / n as f64 * std::f64::consts::TAU; ring.push(put((hw + g) * t.cos(), (hh + g) * t.sin())); }
    }
    (inside, ring)
}

/// The best pose near (x, y) for these pads. `turns` are the rotations to try (degrees, KiCad).
/// Returns (x, y, rot, score) with score = metal inside - metal outside (about 0.5+ is a clear match).
pub fn best_pose(metal: &Plane, pads: &[LPad], x: f64, y: f64, turns: &[f64], reach: f64) -> (f64, f64, f64, f64) { best_pose_at(metal, pads, x, y, turns, reach, None) }

/// The pads' bounding-box centre in the footprint frame: where a part's middle is, whatever its origin.
pub fn pad_centre(pads: &[LPad]) -> (f64, f64) {
    let (x0, y0, x1, y1) = pads.iter().fold((f64::MAX, f64::MAX, f64::MIN, f64::MIN), |a, q| (a.0.min(q.x - q.w / 2.0), a.1.min(q.y - q.h / 2.0), a.2.max(q.x + q.w / 2.0), a.3.max(q.y + q.h / 2.0)));
    ((x0 + x1) / 2.0, (y0 + y1) / 2.0)
}

/// Chip sizes in the KiCad libraries, small to large.
pub const CHIP_SIZES: [&str; 5] = ["0201_0603Metric", "0402_1005Metric", "0603_1608Metric", "0805_2012Metric", "1206_3216Metric"];
/// The pads of a neighbouring size count as fitting the metal better when their score (pads on metal minus ring on metal) is this
/// much higher, up or down. Down used to need no margin at all; on a sharp photo the smaller pads then won whenever the solder
/// round a part was wider than they are (board 10: 11 of 34 chips given at their true size were shrunk, 2026-11-10).
pub const SIZE_MARGIN: f64 = 0.08;
/// With no body reading, pads alone change a size only when the neighbouring size fits this much better. Read on the left half of
/// board 10 (2026-11-10): chips given at their true size gained at most 0.125 from a wrong neighbour, chips given one size too
/// large gained 0.12 to 0.3 from the right one, so 0.15 kept all 35 true sizes and put 34 of 35 oversized ones right.
pub const PADS_ALONE: f64 = 0.15;
/// A body read on the photo belongs to a part when its centre lies within this of the part's own, mm: under half the pitch at
/// which 0402 parts stand side by side, the closest neighbours seen so far.
pub const BODY_WITHIN: f64 = 0.35;

/// The package index (into `CHIP_SIZES`) the part's own body reads as on the photo, from the body-first reader's findings on that
/// face (`passives::find_with`), or None when no body reads at that spot.
pub fn body_read(found: &[crate::passives::Passive], x: f64, y: f64) -> Option<usize> {
    found.iter().map(|p| { let c = p.centre(); (crate::geom::dist(crate::board::to_mm(c.0, c.1), (x, y)), p.pk) })
        .filter(|(d, _)| *d <= BODY_WITHIN).min_by(|a, b| a.0.partial_cmp(&b.0).unwrap()).map(|(_, k)| k)
}

/// The size a chip part takes, and why. `k` is the size it was given, `own` its pads' fit, `fits` the fit of each neighbouring size
/// tried, `body` what `body_read` says. The body comes first: the lands of a part are not the part, and a smaller pair of pads
/// fits inside any larger patch of solder. Where the body reads, the size goes to what it measures, when the pads of that size fit
/// no worse than the given ones by `SIZE_MARGIN`, and stays otherwise. Where no body reads, the size changes only on a pad fit
/// `PADS_ALONE` better, and the reason says what the pads preferred, for whoever reads parts.json.
pub fn settle_size(k: usize, own: f64, fits: &[(usize, f64)], body: Option<usize>) -> (usize, String) {
    let name = |j: usize| &CHIP_SIZES[j][..4];
    let pads = fits.iter().filter(|(_, sc)| *sc >= own + SIZE_MARGIN).max_by(|a, b| a.1.partial_cmp(&b.1).unwrap()).map(|(j, _)| *j);
    match body {
        Some(b) if b == k => (k, match pads { Some(j) => format!("the body reads {}: kept, though {} pads fit the metal better", name(k), name(j)), None => format!("the body reads {}", name(k)) }),
        Some(b) => match fits.iter().find(|(j, _)| *j == b) {
            Some((_, sc)) if *sc >= own - SIZE_MARGIN => (b, format!("the body reads {} and its pads fit ({sc:.2} against {own:.2})", name(b))),
            Some((_, sc)) => (k, format!("the body reads {} but its pads fit worse ({sc:.2} against {own:.2}): kept, check the crop", name(b))),
            None => (k, format!("the body reads {}, more than one size away: kept, check the crop", name(b))),
        },
        None => {
            match fits.iter().filter(|(_, sc)| *sc >= own + PADS_ALONE).max_by(|a, b| a.1.partial_cmp(&b.1).unwrap()) {
                Some((j, sc)) => (*j, format!("no body read here; {} pads fit the metal far better ({sc:.2} against {own:.2})", name(*j))),
                None => (k, match pads { Some(j) => format!("no body read here; {} pads fit the metal a little better, which does not change a size", name(j)), None => "no body read here; the given size's pads fit as well as its neighbours'".to_string() }),
            }
        }
    }
}
/// A two-pad chip footprint of a standard size: (name up to the size, index into `CHIP_SIZES`).
pub fn chip_size(id: &str) -> Option<(String, usize)> {
    if !(id.starts_with("Capacitor_SMD:C_") || id.starts_with("Resistor_SMD:R_") || id.starts_with("Inductor_SMD:L_")) { return None; }
    CHIP_SIZES.iter().position(|z| id.ends_with(z)).map(|k| (id[..id.len() - CHIP_SIZES[k].len()].to_string(), k))
}

/// How far a part may be moved to find its pads: under half its own short side, so a small part in a
/// row can never be carried onto its neighbour's pads; never more than 1 mm.
pub fn reach(pads: &[LPad]) -> f64 {
    let (x0, y0, x1, y1) = pads.iter().fold((f64::MAX, f64::MAX, f64::MIN, f64::MIN), |a, q| (a.0.min(q.x - q.w / 2.0), a.1.min(q.y - q.h / 2.0), a.2.max(q.x + q.w / 2.0), a.3.max(q.y + q.h / 2.0)));
    if pads.is_empty() { return 1.0; }
    (0.45 * (x1 - x0).min(y1 - y0)).clamp(0.12, 1.0)
}

/// As `best_pose`; with `centre` (the pads' centre in the footprint frame), (x, y) is where the part's
/// middle was read on the photo, and each turn searches around the origin that puts the middle there.
pub fn best_pose_at(metal: &Plane, pads: &[LPad], x: f64, y: f64, turns: &[f64], reach: f64, centre: Option<(f64, f64)>) -> (f64, f64, f64, f64) {
    let (inside, ring) = probes(pads);
    if inside.is_empty() { return (x, y, turns[0], 0.0); }
    let mut best = (x, y, turns[0], f64::MIN);
    let step = 1.0 / RES * 2.0;
    for &rot in turns {
        let (c, s) = (rot.to_radians().cos(), rot.to_radians().sin());
        let (x, y) = match centre { Some((cx, cy)) => (x - (cx * c + cy * s), y - (-cx * s + cy * c)), None => (x, y) };
        let tin: Vec<(f64, f64)> = inside.iter().map(|(lx, ly)| (lx * c + ly * s, -lx * s + ly * c)).collect();
        let tout: Vec<(f64, f64)> = ring.iter().map(|(lx, ly)| (lx * c + ly * s, -lx * s + ly * c)).collect();
        let mean = |pts: &[(f64, f64)], ox: f64, oy: f64| -> f64 {
            let mut sum = 0.0; let mut n = 0.0;
            for (dx, dy) in pts { let (u, v) = to_px(ox + dx, oy + dy); if u >= 0.0 && v >= 0.0 && u < metal.w as f64 - 1.0 && v < metal.h as f64 - 1.0 { sum += metal.sample(u as f32, v as f32) as f64; n += 1.0; } }
            if n > 0.0 { sum / n } else { 0.0 }
        };
        let mut oy = -reach;
        while oy <= reach + 1e-9 { let mut ox = -reach;
            while ox <= reach + 1e-9 {
                let sc = mean(&tin, x + ox, y + oy) - mean(&tout, x + ox, y + oy);
                // prefer the AI's own rotation on a tie, and the nearer offset
                let sc_adj = sc - 0.002 * (ox * ox + oy * oy).sqrt() + if rot == turns[0] { 0.01 } else { 0.0 };
                if sc_adj > best.3 { best = (x + ox, y + oy, rot, sc_adj); }
                ox += step; }
            oy += step; }
    }
    best
}

/// Quarter turns that map the pad pattern onto itself (a resistor at 180, a QFN at 90): such a turn is
/// the same part, so the search never tries it and pin 1 stays the AI's.
pub fn symmetric_turns(pads: &[LPad]) -> Vec<f64> {
    let (cx, cy) = pad_centre(pads);
    let mut out = vec![];
    for t in [90.0f64, 180.0, 270.0] {
        let (c, s) = (t.to_radians().cos(), t.to_radians().sin());
        let same = pads.iter().all(|p| {
            let (dx, dy) = (p.x - cx, p.y - cy);
            let (rx, ry) = (cx + dx * c + dy * s, cy - dx * s + dy * c);
            let quarter = (t / 90.0) as i64 % 2 == 1;
            pads.iter().any(|q| (q.x - rx).abs() < 0.05 && (q.y - ry).abs() < 0.05 && {
                let (pw, ph) = if quarter ^ ((p.rot - q.rot).rem_euclid(180.0).abs() > 45.0) { (p.h, p.w) } else { (p.w, p.h) };
                (pw - q.w).abs() < 0.05 && (ph - q.h).abs() < 0.05 })
        });
        if same { out.push(t); }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    fn two(w: f64, h: f64, pitch: f64) -> Vec<LPad> { [-1.0, 1.0].iter().map(|s| LPad { x: s * pitch / 2.0, y: 0.0, w, h, rot: 0.0, round: false, drill: 0.0 }).collect() }
    #[test]
    fn a_chip_size_is_read_from_its_name() {
        assert_eq!(chip_size("Capacitor_SMD:C_0402_1005Metric"), Some(("Capacitor_SMD:C_".to_string(), 1)));
        assert_eq!(chip_size("Resistor_SMD:R_1206_3216Metric"), Some(("Resistor_SMD:R_".to_string(), 4)));
        assert_eq!(chip_size("Package_SO:SOIC-8_3.9x4.9mm_P1.27mm"), None);
        assert_eq!(chip_size("Capacitor_SMD:C_0402_1005Metric_Pad0.74x0.62mm_HandSolder"), None);
    }
    #[test]
    fn pads_of_the_right_size_fit_better_than_the_next_size_up() {
        // metal of an 0402 on its lands: two blocks 0.5 wide, from 0.2 to 0.75 mm either side of the middle
        let mut m = Plane { w: 400, h: 400, d: vec![0.0; 160000] };
        let (bx, by) = (crate::board::ORIGIN + 5.0, crate::board::ORIGIN + 5.0);
        for s in [-1.0, 1.0] { let (u0, v0) = to_px(bx + s * 0.2, by - 0.25); let (u1, v1) = to_px(bx + s * 0.75, by + 0.25);
            for y in v0 as usize..=v1 as usize { for x in (u0.min(u1) as usize)..=(u0.max(u1) as usize) { m.d[y * 400 + x] = 1.0f32; } } }
        let small = two(0.56, 0.62, 0.96); let big = two(0.875, 0.95, 1.55);
        let a = best_pose_at(&m, &small, bx, by, &[0.0], reach(&small), None).3;
        let b = best_pose_at(&m, &big, bx, by, &[0.0], reach(&big), None).3;
        assert!(a >= b + SIZE_MARGIN, "0402 pads {a:.2} against 0603 pads {b:.2}");
    }
    #[test]
    fn a_size_follows_the_body_before_the_pads() {
        // board 10, C24: an 0402 whose 0201 pads fit better (0.523 against 0.398). No body reading and under `PADS_ALONE`: it stays 0402
        let (k, why) = settle_size(1, 0.398, &[(0, 0.523), (2, 0.105)], None);
        assert_eq!(k, 1); assert!(why.contains("0201 pads fit the metal a little better"), "{why}");
        // board 10, R2 given as 0603: the 0402 pads fit far better (0.272 against 0.095), so it becomes 0402 without a body reading
        assert_eq!(settle_size(2, 0.095, &[(1, 0.272), (3, 0.05)], None).0, 1);
        // the body reads 0402: it stays, whatever the pads say
        assert_eq!(settle_size(1, 0.398, &[(0, 0.523), (2, 0.105)], Some(1)).0, 1);
        // given as 0805, the body reads 0603 and the 0603 pads fit: it becomes 0603
        assert_eq!(settle_size(3, 0.24, &[(2, 0.56), (4, 0.1)], Some(2)).0, 2);
        // the body reads a size whose pads fit clearly worse, or one two sizes away: no change
        assert_eq!(settle_size(1, 0.50, &[(0, 0.30), (2, 0.2)], Some(0)).0, 1);
        assert_eq!(settle_size(1, 0.50, &[(0, 0.30), (2, 0.2)], Some(3)).0, 1);
        // a body within reach is the nearest one
        let p = |x: f64, y: f64, pk: usize| crate::passives::Passive { kind: "tan", x, y, rot: 0.0, pk, q: 1.0, body_mm: (0.5, 0.5), cost: 0.1, ends: 0.9, across: 0.0, axis: (20.0, 20.0, 0.0), own_levels: true };
        let found = [p(200.5, 200.5, 1), p(260.5, 200.5, 2)];
        assert_eq!(body_read(&found, crate::board::ORIGIN + 5.1, crate::board::ORIGIN + 5.0), Some(1));
        assert_eq!(body_read(&found, crate::board::ORIGIN + 6.4, crate::board::ORIGIN + 5.2), Some(2));
        assert_eq!(body_read(&found, crate::board::ORIGIN + 5.75, crate::board::ORIGIN + 5.0), None);
    }
    #[test]
    fn reach_follows_the_part() {
        // an 0201 (pads 0.3 x 0.3, 0.6 apart) may move well under the 0.6 mm to the next part of a row
        let r = reach(&two(0.3, 0.3, 0.6)); assert!(r < 0.2, "{r}");
        // a wide part keeps the full millimetre
        let r = reach(&two(1.5, 6.0, 5.0)); assert!((r - 1.0).abs() < 1e-9, "{r}");
    }
    #[test]
    fn a_small_part_is_not_carried_onto_its_neighbour() {
        // two 0201 land patterns 0.7 mm apart; the reading sits 0.1 mm off the left one
        let mut m = Plane { w: 400, h: 400, d: vec![0.0; 160000] };
        let mut pad = |cx: f64, cy: f64| { let (u, v) = to_px(cx, cy); for y in (v as i64 - 6)..=(v as i64 + 6) { for x in (u as i64 - 6)..=(u as i64 + 6) { m.d[y as usize * 400 + x as usize] = 1.0f32; } } };
        let (bx, by) = (crate::board::ORIGIN + 3.0, crate::board::ORIGIN + 3.0);
        // left part: only half its metal shows; right part: bright and whole
        pad(bx - 0.3, by); pad(bx + 0.4, by + 0.7); pad(bx + 1.0, by + 0.7);
        let pads = two(0.3, 0.3, 0.6);
        let far = best_pose_at(&m, &pads, bx + 0.1, by + 0.1, &[0.0], 1.0, None);
        let near = best_pose_at(&m, &pads, bx + 0.1, by + 0.1, &[0.0], reach(&pads), None);
        assert!((far.1 - by).abs() > 0.5, "with a 1 mm reach the search leaves for the neighbour: {far:?}");
        assert!((near.1 - by).abs() < 0.2 && (near.0 - bx).abs() < 0.2, "with its own reach it stays: {near:?}");
    }
}