use std::env;
use std::fs;
use std::process;

const MAX_WORK: u128 = 200_000_000;

#[derive(Debug)]
struct Candidate {
    k: usize,
    q: usize,
    colors: Vec<u8>,
    claimed_pairs: u128,
    claimed_num: u128,
    claimed_den: u128,
}

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

fn parse_u128(name: &str, text: &str) -> u128 {
    text.parse::<u128>()
        .unwrap_or_else(|_| fail(format!("{name} is not a base-10 nonnegative integer")))
}

fn parse_usize(name: &str, text: &str) -> usize {
    let value = parse_u128(name, text);
    usize::try_from(value).unwrap_or_else(|_| fail(format!("{name} does not fit in usize")))
}

fn parse_fraction(text: &str) -> (u128, u128) {
    let mut pieces = text.split('/');
    let num_text = pieces.next().unwrap_or("");
    let den_text = pieces
        .next()
        .unwrap_or_else(|| fail("claimed_bound must have the form NUM/DEN"));
    if pieces.next().is_some() || num_text.is_empty() || den_text.is_empty() {
        fail("claimed_bound must have exactly one '/' and the form NUM/DEN");
    }
    let num = parse_u128("claimed_bound numerator", num_text);
    let den = parse_u128("claimed_bound denominator", den_text);
    if den == 0 {
        fail("claimed_bound denominator is zero");
    }
    (num, den)
}

fn parse_candidate(input: &str) -> Candidate {
    let lines: Vec<&str> = input.lines().collect();
    if lines.len() != 6 {
        fail(format!(
            "expected exactly 6 lines, including the AP_PERIODIC_V1 header; found {}",
            lines.len()
        ));
    }
    if lines[0] != "AP_PERIODIC_V1" {
        fail("line 1 must be exactly AP_PERIODIC_V1");
    }

    fn field<'a>(line: &'a str, expected: &str, line_no: usize) -> &'a str {
        let mut words = line.split_ascii_whitespace();
        let key = words.next().unwrap_or("");
        let value = words.next().unwrap_or("");
        if key != expected || value.is_empty() || words.next().is_some() {
            fail(format!(
                "line {line_no} must be exactly '{expected} VALUE' with one value"
            ));
        }
        value
    }

    let k = parse_usize("k", field(lines[1], "k", 2));
    let q = parse_usize("period", field(lines[2], "period", 3));
    let colors_text = field(lines[3], "colors", 4);
    let claimed_pairs = parse_u128(
        "claimed_mono_pairs",
        field(lines[4], "claimed_mono_pairs", 5),
    );
    let (claimed_num, claimed_den) =
        parse_fraction(field(lines[5], "claimed_bound", 6));

    if k < 3 {
        fail("k must be at least 3");
    }
    if q == 0 {
        fail("period must be positive");
    }
    if colors_text.len() != q {
        fail(format!(
            "colors has byte-length {}, but period is {q}",
            colors_text.len()
        ));
    }
    let colors: Vec<u8> = colors_text
        .bytes()
        .enumerate()
        .map(|(i, b)| match b {
            b'0' | b'1' => b - b'0',
            _ => fail(format!(
                "colors byte {} is not '0' or '1'",
                i + 1
            )),
        })
        .collect();

    let work = (q as u128)
        .checked_mul(q as u128)
        .and_then(|x| x.checked_mul(k as u128))
        .unwrap_or_else(|| fail("period^2 * k overflows u128"));
    if work > MAX_WORK {
        fail(format!(
            "certificate requests {work} elementary checks, exceeding safety limit {MAX_WORK}"
        ));
    }

    Candidate {
        k,
        q,
        colors,
        claimed_pairs,
        claimed_num,
        claimed_den,
    }
}

fn gcd(mut a: u128, mut b: u128) -> u128 {
    while b != 0 {
        let r = a % b;
        a = b;
        b = r;
    }
    a
}

fn reduced(num: u128, den: u128) -> (u128, u128) {
    let g = gcd(num, den);
    (num / g, den / g)
}

fn count_mono_pairs(candidate: &Candidate) -> u128 {
    let mut count = 0_u128;
    for a in 0..candidate.q {
        for d in 0..candidate.q {
            let color = candidate.colors[a];
            let mut position = a;
            let mut mono = true;
            for _ in 1..candidate.k {
                // position and d are both below q, so this addition cannot overflow
                // after the parser's work bound has forced q to be modest.
                position = (position + d) % candidate.q;
                if candidate.colors[position] != color {
                    mono = false;
                    break;
                }
            }
            if mono {
                count += 1;
            }
        }
    }
    count
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() != 2 {
        fail(format!("usage: {} CANDIDATE.txt", args[0]));
    }
    let input = fs::read_to_string(&args[1])
        .unwrap_or_else(|e| fail(format!("cannot read '{}': {e}", args[1])));
    let candidate = parse_candidate(&input);

    let actual_pairs = count_mono_pairs(&candidate);
    if candidate.claimed_pairs != actual_pairs {
        fail(format!(
            "claimed_mono_pairs is {}, but exact enumeration gives {actual_pairs}",
            candidate.claimed_pairs
        ));
    }

    let q = candidate.q as u128;
    let expected_den = 2_u128
        .checked_mul((candidate.k - 1) as u128)
        .and_then(|x| x.checked_mul(q))
        .and_then(|x| x.checked_mul(q))
        .unwrap_or_else(|| fail("bound denominator overflows u128"));
    let expected = reduced(actual_pairs, expected_den);
    let claimed = reduced(candidate.claimed_num, candidate.claimed_den);
    if claimed != expected {
        fail(format!(
            "claimed_bound reduces to {}/{}, but the periodic certificate gives {}/{}",
            claimed.0, claimed.1, expected.0, expected.1
        ));
    }

    println!("PASS");
    println!("k = {}", candidate.k);
    println!(
        "monochromatic residue pairs = {} / {}",
        actual_pairs,
        q * q
    );
    println!(
        "certified asymptotic upper bound: delta_{} <= {}/{}",
        candidate.k, expected.0, expected.1
    );
}
