AI Flow: Reverse Engineering

//! Part candidates in board millimetres. Bodies found on the board images get a package match by
//! measured size (with and without their terminations), a provisional reference, an orientation in
//! KiCad degrees and a crop for reading the marking. Rows of plated holes become connector candidates
//! (pitch and count give the footprint), and large square pads with a hole become lug candidates.
use crate::board::{to_mm, RES};
use image::RgbImage;
use serde_json::{json, Value};

/// (package, KiCad footprint, length, width with terminations, mm). Sizes from the KiCad library's own
/// footprint names and JEDEC/EIA nominal bodies; a match within 12% on both axes counts.
pub const PACKAGES: &[(&str, &str, f64, f64)] = &[
    ("0402", "Resistor_SMD:R_0402_1005Metric", 1.0, 0.5),
    ("0603", "Resistor_SMD:R_0603_1608Metric", 1.6, 0.8),
    ("0805", "Resistor_SMD:R_0805_2012Metric", 2.0, 1.25),
    ("1206", "Resistor_SMD:R_1206_3216Metric", 3.2, 1.6),
    ("1210", "Resistor_SMD:R_1210_3225Metric", 3.2, 2.5),
    ("2010", "Resistor_SMD:R_2010_5025Metric", 5.0, 2.5),
    ("2512", "Resistor_SMD:R_2512_6332Metric", 6.3, 3.2),
    ("SOD-123", "Diode_SMD:D_SOD-123", 3.7, 1.6),
    ("SMA", "Diode_SMD:D_SMA", 5.2, 2.6),
    ("SMB", "Diode_SMD:D_SMB", 5.4, 3.6),
    ("SOT-23", "Package_TO_SOT_SMD:SOT-23", 2.9, 2.4),
    ("SOT-23-8", "Package_TO_SOT_SMD:SOT-23-8", 2.9, 2.8),
    ("SOT-223", "Package_TO_SOT_SMD:SOT-223-3_TabPin2", 7.0, 6.5),
    ("SOIC-8", "Package_SO:SOIC-8_3.9x4.9mm_P1.27mm", 6.0, 4.9),
    ("PDFN-8 5x6", "Package_TO_SOT_SMD:TDSON-8-1", 6.15, 5.15),
    ("QFN 3x3", "Package_DFN_QFN:QFN-16-1EP_3x3mm_P0.5mm_EP1.7x1.7mm", 3.0, 3.0),
    ("QFN 4x4", "Package_DFN_QFN:QFN-24-1EP_4x4mm_P0.5mm_EP2.6x2.6mm", 4.0, 4.0),
    ("QFN 5x5", "Package_DFN_QFN:QFN-32-1EP_5x5mm_P0.5mm_EP3.45x3.45mm", 5.0, 5.0),
    ("QFN 6x6", "Package_DFN_QFN:QFN-40-1EP_6x6mm_P0.5mm_EP4.6x4.6mm", 6.0, 6.0),
    ("QFN 7x7", "Package_DFN_QFN:QFN-48-1EP_7x7mm_P0.5mm_EP5.6x5.6mm", 7.0, 7.0),
    ("Inductor 6x6", "Inductor_SMD:L_Bourns-SRN6045TA", 6.0, 6.0),
    ("Inductor 8x8", "Inductor_SMD:L_Bourns-SRN8040_8x8.15mm", 8.15, 8.0),
];

/// Ceramic capacitors show only their tan body: the length less both terminations (EIA nominal:
/// 0402 0.25 mm, 0603 0.35, 0805 0.5, 1206 0.5, 1210 0.5 mm per end).
pub const CERAMIC_BODIES: &[(&str, &str, f64, f64)] = &[
    ("0402", "Capacitor_SMD:C_0402_1005Metric", 0.5, 0.5),
    ("0603", "Capacitor_SMD:C_0603_1608Metric", 0.9, 0.8),
    ("0805", "Capacitor_SMD:C_0805_2012Metric", 1.25, 1.0),
    ("1206", "Capacitor_SMD:C_1206_3216Metric", 2.2, 1.6),
    ("1210", "Capacitor_SMD:C_1210_3225Metric", 2.5, 2.2),
];

fn f(v: &Value, k: &str) -> f64 { v[k].as_f64().unwrap_or(0.0) }

fn match_in(table: &[(&str, &str, f64, f64)], l: f64, w: f64) -> Vec<(String, String, f64)> {
    let (l, w) = if l >= w { (l, w) } else { (w, l) };
    let mut m: Vec<(String, String, f64)> = table.iter().map(|(n, fp, pl, pw)| { let (pl, pw) = if pl >= pw { (*pl, *pw) } else { (*pw, *pl) };
        (n.to_string(), fp.to_string(), ((l - pl) / pl).abs().max(((w - pw) / pw).abs())) }).collect();
    m.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap()); m.truncate(3); m
}

/// Best package matches for measured (length, width) in mm, with the worse-axis relative error.
pub fn match_package(l: f64, w: f64) -> Vec<(String, String, f64)> {
    let mut m: Vec<(String, String, f64)> = PACKAGES.iter().map(|(n, fp, pl, pw)| {
        let e = ((l - pl) / pl).abs().max(((w - pw) / pw).abs()); (n.to_string(), fp.to_string(), e)
    }).collect();
    m.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap());
    m.truncate(3);
    m
}

fn prefix(colour: &str, pkg: &str) -> &'static str {
    match (colour, pkg) {
        ("tan", _) => "C",
        (_, p) if p.starts_with("SOD") || p.starts_with("SM") => "D",
        (_, p) if p.starts_with("PDFN") || p == "SOT-23" => "Q",
        (_, p) if p.starts_with("Inductor") => "L",
        (_, p) if p.starts_with("QFN") || p.starts_with("SOIC") || p.starts_with("SOT-223") || p.starts_with("SOT-23-") => "U",
        ("black", p) if ["0402", "0603", "0805", "1206", "1210", "2010", "2512"].contains(&p) => "R",
        ("yellow", _) | ("bright-green", _) | ("white", _) => "J",
        _ => "X",
    }
}

/// Bodies (board-image pixels) to candidates in board mm. `side` is top or bottom.
pub fn candidates(bodies: &[Value], side: &str, start: usize) -> Vec<Value> {
    let mut out = vec![];
    for (k, b) in bodies.iter().enumerate() {
        let (l, w) = (f(b, "long") / RES, f(b, "short") / RES);
        let (bl, bw) = (f(&b["body"], "long") / RES, f(&b["body"], "short") / RES);
        if l < 0.7 && w < 0.4 { continue; }
        let colour = b["colour"].as_str().unwrap_or("");
        // terminations count for passives and leaded parts; the moulding alone for square packages
        // ceramics by their tan body; everything else by its full extent or its moulding, whichever fits
        let m = if colour == "tan" { match_in(CERAMIC_BODIES, bl, bw) } else { let a = match_package(l, w); let b = match_package(bl, bw); if b[0].2 < a[0].2 { b } else { a } };
        let fits = m[0].2 < 0.15;
        let pos = to_mm(f(b, "x"), f(b, "y"));
        let id = format!("P{:03}", start + out.len() + 1);
        // KiCad rotation is counter-clockwise on screen; the image angle runs clockwise
        let rot = (-(f(b, "angle"))).rem_euclid(180.0);
        let rot = [0.0, 45.0, 90.0, 135.0, 180.0].iter().copied().find(|s| (rot - s).abs() < 4.0).map(|s| s % 180.0).unwrap_or((rot * 10.0).round() / 10.0);
        out.push(json!({"id": id, "side": side, "pos_mm": [(pos.0 * 1000.0).round() / 1000.0, (pos.1 * 1000.0).round() / 1000.0], "rot": rot,
                        "size_mm": [(l * 100.0).round() / 100.0, (w * 100.0).round() / 100.0], "body_mm": [(bl * 100.0).round() / 100.0, (bw * 100.0).round() / 100.0],
                        "colour": colour, "fits": fits, "prefix": if fits { prefix(colour, &m[0].0) } else { "X" },
                        "package": if fits { json!(m[0].0) } else { Value::Null }, "footprint": if fits { json!(m[0].1) } else { Value::Null },
                        "matches": m.iter().map(|(n, fp, e)| json!({"package": n, "footprint": fp, "error": (e * 1000.0).round() / 1000.0})).collect::<Vec<_>>(),
                        "source": format!("body {k}")}));
    }
    out
}

/// Connector candidates from hole rows (board mm holes) and lug candidates from large square pads.
pub fn hole_parts(holes: &[Value], rows: &[Value], start: usize) -> Vec<Value> {
    let mut out = vec![];
    for r in rows {
        let ids: Vec<usize> = r["holes"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect();
        let pitch = r["pitch_mm"].as_f64().unwrap_or(0.0);
        let std = crate::board::PITCHES.iter().copied().find(|p| ((pitch - p) / p).abs() < 0.03);
        let Some(p) = std else { continue };
        let n = ids.len();
        let first = &holes[ids[0]]; let last = &holes[ids[n - 1]];
        let (x0, y0, x1, y1) = (f(first, "x"), f(first, "y"), f(last, "x"), f(last, "y"));
        let vertical = (y1 - y0).abs() > (x1 - x0).abs();
        let fp = match p { 2.0 => format!("Connector_PinHeader_2.00mm:PinHeader_1x{n:02}_P2.00mm_Vertical"), 2.54 => format!("Connector_PinHeader_2.54mm:PinHeader_1x{n:02}_P2.54mm_Vertical"),
                           3.5 => format!("TerminalBlock_Phoenix:TerminalBlock_Phoenix_PT-1,5-{n}-3.5-H_1x{n:02}_P3.50mm_Horizontal"), 5.0 | 5.08 => format!("TerminalBlock_Phoenix:TerminalBlock_Phoenix_MKDS-1,5-{n}-5.08_1x{n:02}_P5.08mm_Horizontal"),
                           _ => format!("Connector_PinHeader_2.54mm:PinHeader_1x{n:02}_P2.54mm_Vertical") };
        // pin 1 of a library header sits at the footprint origin, the row running +y at rotation 0
        let rot: f64 = if vertical { if y1 > y0 { 0.0 } else { 180.0 } } else if x1 > x0 { 90.0 } else { 270.0 };
        out.push(json!({"id": format!("P{:03}", start + out.len() + 1), "side": "top", "pos_mm": [x0, y0], "rot": rot, "pitch_mm": p, "pins": n,
                        "holes": ids, "fits": true, "prefix": "J", "package": format!("1x{n} {p} mm"), "footprint": fp, "origin": "pin 1 (first hole of the row: check which end is pin 1 from the square pad)",
                        "source": "hole row"}));
    }
    for (i, h) in holes.iter().enumerate() {
        if f(h, "outer_mm") > 3.5 && f(h, "hole_mm") >= 1.5 {
            out.push(json!({"id": format!("P{:03}", start + out.len() + 1), "side": "top", "pos_mm": [f(h, "x"), f(h, "y")], "rot": 0.0, "fits": true, "prefix": "TP",
                            "package": format!("lug {:.1} mm, hole {:.2} mm", f(h, "outer_mm"), f(h, "hole_mm")), "footprint": "generated:lug", "lug": {"pad_mm": f(h, "outer_mm"), "drill_mm": f(h, "hole_mm")}, "holes": [i], "source": "large square pad"}));
        }
    }
    out
}

/// A crop around a candidate, enlarged for reading markings.
pub fn crop(img: &RgbImage, pos_mm: (f64, f64), size_mm: (f64, f64), scale: u32) -> RgbImage {
    let (cx, cy) = crate::board::to_px(pos_mm.0, pos_mm.1);
    let half = ((size_mm.0.max(size_mm.1) / 2.0 + 1.0) * RES) as i64;
    let (x0, y0) = ((cx as i64 - half).max(0) as u32, (cy as i64 - half).max(0) as u32);
    let w = ((2 * half) as u32).min(img.width() - x0); let h = ((2 * half) as u32).min(img.height() - y0);
    let c = image::imageops::crop_imm(img, x0, y0, w, h).to_image();
    image::imageops::resize(&c, w * scale, h * scale, image::imageops::FilterType::CatmullRom)
}

/// A contact sheet: crops in a grid, each labelled with its id.
pub fn contact_sheet(crops: &[(String, RgbImage)], cell: u32, cols: u32) -> RgbImage {
    let rows = (crops.len() as u32).div_ceil(cols).max(1);
    let mut sheet = RgbImage::from_pixel(cols * cell, rows * cell, image::Rgb([20, 20, 24]));
    for (k, (id, c)) in crops.iter().enumerate() {
        let (gx, gy) = ((k as u32 % cols) * cell, (k as u32 / cols) * cell);
        let s = image::imageops::resize(c, cell - 8, cell - 22, image::imageops::FilterType::Triangle);
        image::imageops::overlay(&mut sheet, &s, (gx + 4) as i64, (gy + 18) as i64);
        crate::draw::label(&mut sheet, (gx as f64 + 4.0, gy as f64 + 3.0), id.trim_start_matches('P'), [255, 255, 0], 3);
    }
    sheet
}

/// Resistor and inductor marking codes: 3-digit (472), 4-digit (4701), R-notation (R002, 4R7) and
/// EIA-96 (30B). Returns ohms (or the code's base unit) and how it was read.
pub fn decode(code: &str) -> Option<(f64, String)> {
    let c = code.trim().to_uppercase();
    const EIA: [f64; 96] = [100., 102., 105., 107., 110., 113., 115., 118., 121., 124., 127., 130., 133., 137., 140., 143., 147., 150., 154., 158., 162., 165., 169., 174., 178., 182., 187., 191., 196., 200., 205., 210., 215., 221., 226., 232., 237., 243., 249., 255., 261., 267., 274., 280., 287., 294., 301., 309., 316., 324., 332., 340., 348., 357., 365., 374., 383., 392., 402., 412., 422., 432., 442., 453., 464., 475., 487., 499., 511., 523., 536., 549., 562., 576., 590., 604., 619., 634., 649., 665., 681., 698., 715., 732., 750., 768., 787., 806., 825., 845., 866., 887., 909., 931., 953., 976.];
    if c.contains('R') && c.chars().all(|ch| ch.is_ascii_digit() || ch == 'R') && c.matches('R').count() == 1 {
        return c.replace('R', ".").parse::<f64>().ok().map(|v| (v, format!("R-notation {c}")));
    }
    if c.len() == 3 && c.chars().take(2).all(|ch| ch.is_ascii_digit()) && c.chars().nth(2).unwrap().is_ascii_alphabetic() {
        let i: usize = c[..2].parse().ok()?;
        if !(1..=96).contains(&i) { return None; }
        let m = match c.chars().nth(2)? { 'Z' => 0.001, 'Y' | 'R' => 0.01, 'X' | 'S' => 0.1, 'A' => 1.0, 'B' | 'H' => 10.0, 'C' => 100.0, 'D' => 1000.0, 'E' => 10000.0, 'F' => 100000.0, _ => return None };
        return Some((EIA[i - 1] * m, format!("EIA-96 {c}: {} x {m}", EIA[i - 1])));
    }
    if (c.len() == 3 || c.len() == 4) && c.chars().all(|ch| ch.is_ascii_digit()) {
        let (sig, exp) = c.split_at(c.len() - 1);
        let v = sig.parse::<f64>().ok()? * 10f64.powi(exp.parse::<i32>().ok()?);
        return Some((v, format!("{}-digit {c}", c.len())));
    }
    None
}

pub fn si(v: f64, unit: &str) -> String {
    let (d, p) = if v >= 1e6 { (1e6, "M") } else if v >= 1e3 { (1e3, "k") } else if v >= 1.0 { (1.0, "") } else { (1e-3, "m") };
    let s = format!("{:.3}", v / d); let s = s.trim_end_matches('0').trim_end_matches('.');
    format!("{s}{p}{unit}")
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn codes_and_packages() {
        assert_eq!(decode("472").unwrap().0, 4700.0);
        assert_eq!(decode("4701").unwrap().0, 4700.0);
        assert_eq!(decode("R002").unwrap().0, 0.002);
        assert_eq!(decode("30B").unwrap().0, 2000.0);
        assert_eq!(decode("59C").unwrap().0, 40200.0);
        assert_eq!(decode("01C").unwrap().0, 10000.0);
        assert_eq!(si(40200.0, "Ω"), "40.2kΩ");
        assert_eq!(match_package(1.62, 0.83)[0].0, "0603");
        assert_eq!(match_package(6.1, 5.95)[0].0, "QFN 6x6");
    }
}