//! Independent exact checker for SDP/diagonal-dominance lower-bound certificates
//! for the minimum monochromatic 3-AP density.

use std::{env, fs, process};

fn fail(msg: impl AsRef<str>) -> ! {
    eprintln!("FAIL: {}", msg.as_ref());
    process::exit(1)
}

struct Tokens<'a> { xs: Vec<&'a str>, at: usize }
impl<'a> Tokens<'a> {
    fn new(s: &'a str) -> Self { Self { xs: s.split_whitespace().collect(), at: 0 } }
    fn next(&mut self) -> &'a str {
        if self.at >= self.xs.len() { fail("unexpected end of certificate") }
        let x = self.xs[self.at]; self.at += 1; x
    }
    fn expect(&mut self, want: &str) {
        let got = self.next();
        if got != want { fail(format!("expected token {want:?}, got {got:?}")) }
    }
    fn usize(&mut self, what: &str) -> usize {
        self.next().parse().unwrap_or_else(|_| fail(format!("invalid {what}")))
    }
    fn i128(&mut self, what: &str) -> i128 {
        self.next().parse().unwrap_or_else(|_| fail(format!("invalid {what}")))
    }
}

fn add(a: i128, b: i128, what: &str) -> i128 {
    a.checked_add(b).unwrap_or_else(|| fail(format!("integer overflow in {what}")))
}
fn sub(a: i128, b: i128, what: &str) -> i128 {
    a.checked_sub(b).unwrap_or_else(|| fail(format!("integer overflow in {what}")))
}
fn mul(a: i128, b: i128, what: &str) -> i128 {
    a.checked_mul(b).unwrap_or_else(|| fail(format!("integer overflow in {what}")))
}
fn abs(a: i128) -> i128 {
    a.checked_abs().unwrap_or_else(|| fail("integer overflow in abs"))
}
fn gcd(mut a: i128, mut b: i128) -> i128 {
    a = abs(a); b = abs(b);
    while b != 0 { let r = a % b; a = b; b = r; }
    a
}

/// The PRS block matrix, using zero-based indices for the paper's 1-based rule.
fn a_entry(i0: usize, j0: usize, l: usize) -> i128 {
    let i = i0 + 1;
    let j = j0 + 1;
    let bij = (j + 1 <= 2 * i && 2 * i <= j + l) as i128;
    let bji = (i + 1 <= 2 * j && 2 * j <= i + l) as i128;
    bij + bji
}

fn main() {
    let path = env::args().nth(1).unwrap_or_else(|| fail("usage: check-sdp-lower CERTIFICATE"));
    if env::args().nth(2).is_some() { fail("too many arguments") }
    let text = fs::read_to_string(&path).unwrap_or_else(|e| fail(format!("cannot read {path}: {e}")));
    if text.len() > 20_000_000 { fail("certificate too large") }
    let mut t = Tokens::new(&text);
    t.expect("AP_SDP_LOWER_V1");
    t.expect("L"); let l = t.usize("L");
    t.expect("denom"); let d = t.i128("denom");
    t.expect("factor_scale"); let r = t.i128("factor_scale");
    if l == 0 || l > 1024 || l % 2 != 0 { fail("L must be positive, even, and at most 1024") }
    if !(1..=1_000_000_000).contains(&d) || !(1..=1_000_000_000).contains(&r) {
        fail("denom and factor_scale must lie in 1..=10^9")
    }
    t.expect("diag");
    let mut diag = Vec::with_capacity(l);
    for _ in 0..l {
        let x = t.i128("diagonal entry");
        if x < 0 || x > 1_000_000_000_000 { fail("diagonal entry out of range") }
        diag.push(x);
    }
    t.expect("factor");
    let mut p = vec![0i128; l * l];
    for i in 0..l {
        for j in 0..=i {
            let x = t.i128("factor entry");
            if abs(x) > 10_000_000_000 { fail("factor entry out of range") }
            p[i * l + j] = x;
        }
    }
    t.expect("claimed_diag_sum"); let claimed_sum = t.i128("claimed_diag_sum");
    t.expect("claimed_bound");
    let claimed_num = t.i128("claimed bound numerator");
    let claimed_den = t.i128("claimed bound denominator");
    if t.at != t.xs.len() { fail("trailing tokens") }

    let actual_sum = diag.iter().fold(0i128, |s, &x| add(s, x, "diagonal sum"));
    if actual_sum != claimed_sum { fail(format!("claimed diagonal sum {claimed_sum}, actual {actual_sum}")) }
    let b_support: usize = (0..l).map(|i0| {
        let i = i0 + 1;
        (1..=l).filter(|&j| j + 1 <= 2*i && 2*i <= j + l).count()
    }).sum();
    if b_support != l*l/2 { fail(format!("block relation has unexpected support {b_support}")) }

    // Q = R^2 (D A + diag(diag)) - D P P^T.  If Q is symmetric
    // diagonally dominant with nonnegative diagonal, then Q is PSD.  Therefore
    // D A + diag(diag) = (D/R^2) P P^T + Q/R^2 is PSD exactly.
    let rr = mul(r, r, "R^2");
    let mut qmat = vec![0i128; l * l];
    for i in 0..l {
        for j in 0..l {
            let mut dot = 0i128;
            for k in 0..=i.min(j) {
                dot = add(dot, mul(p[i*l+k], p[j*l+k], "factor dot product"), "factor dot product");
            }
            let c = add(mul(d, a_entry(i,j,l), "scaled A"), if i == j { diag[i] } else { 0 }, "certificate matrix");
            qmat[i*l+j] = sub(mul(rr, c, "R^2 C"), mul(d, dot, "D P P^T"), "residual");
        }
    }
    let mut min_margin: Option<i128> = None;
    for i in 0..l {
        if qmat[i*l+i] < 0 { fail(format!("residual diagonal {i} is negative")) }
        let mut radius = 0i128;
        for j in 0..l {
            if qmat[i*l+j] != qmat[j*l+i] { fail("residual is not symmetric") }
            if i != j { radius = add(radius, abs(qmat[i*l+j]), "Gershgorin radius") }
        }
        let margin = sub(qmat[i*l+i], radius, "diagonal-dominance margin");
        if margin < 0 { fail(format!("residual row {i} is not diagonally dominant (margin {margin})")) }
        min_margin = Some(min_margin.map_or(margin, |m| m.min(margin)));
    }

    // Claimed lower bound is exactly 1/16 - 1/(8L) - sum(diag)/(8 D L^2).
    let li = l as i128;
    let l2 = mul(li, li, "L^2");
    let raw_num = sub(sub(mul(d, l2, "bound numerator"), mul(2*d, li, "bound numerator"), "bound numerator"), 2*actual_sum, "bound numerator");
    let raw_den = mul(16*d, l2, "bound denominator");
    if raw_num <= 0 { fail("certified bound is not positive") }
    let g = gcd(raw_num, raw_den);
    let num = raw_num/g; let den = raw_den/g;
    if claimed_num != num || claimed_den != den {
        fail(format!("claimed bound {claimed_num}/{claimed_den}, actual {num}/{den}"))
    }
    // Require a strict improvement on the published PRS08 lower bound.
    if mul(num, 32768, "comparison") <= mul(1675, den, "comparison") {
        fail("bound does not strictly improve 1675/32768")
    }
    println!("PASS: exact PSD certificate (L={l}, sum={actual_sum}, min residual margin={})", min_margin.unwrap());
    println!("PASS: delta_3 >= {num}/{den} = {:.12}", num as f64 / den as f64);
    println!("PASS: strictly improves PRS08 1675/32768 = {:.12}", 1675f64/32768f64);
}
