use std::fs;
use z3::{Config, SatResult, Solver, ast::Int, with_z3_config};
const COLORS: u64 = 3;
fn main() {
    let numbers: Option<usize> = if let Some(numbers) = std::env::args().skip(1).next()
        && let Ok(numbers) = numbers.parse()
    {
        Some(numbers)
    } else {
        None
    };
    let mut config = Config::new();
    config.set_model_generation(true);
    with_z3_config(&config, || {
        for i in 1.. {
            let i = if let Some(numbers) = numbers {
                numbers
            } else {
                i
            };
            let solver = Solver::new();
            let nums: Vec<_> = (1..=i).map(|n| Int::new_const(n.to_string())).collect();
            nums.iter().for_each(|n| {
                solver.assert(&n.ge(0));
                solver.assert(&n.le(COLORS - 1));
            });
            for diff in (1..=nums.len().isqrt()).map(|n| n * n) {
                for j in 0..nums.len() - diff {
                    solver.assert(&nums[j].ne(&nums[j + diff]));
                }
            }
            match solver.check() {
                SatResult::Sat => {
                    println!("Works for {COLORS} colors and {i} numbers.");
                    if let Some(model) = solver.get_model() {
                        let model_str = model.to_string();
                        let mut model: Vec<_> = model_str.lines().collect();
                        model.sort_by_key(|line| {
                            line.split(" -> ").next().unwrap().parse::<u32>().unwrap()
                        });
                        fs::write(
                            format!("models/model-c{COLORS}-n{i}.txt"),
                            format!("{}", model.join("\n")),
                        )
                        .unwrap();
                    }
                }
                SatResult::Unsat => {
                    println!("Does not work for {COLORS} colors and {i} numbers.");
                    break;
                }
                SatResult::Unknown => {
                    println!("Unknown for {COLORS} colors and {i} numbers.");
                    break;
                }
            }
            if numbers.is_some() {
                break;
            }
        }
    });
}

[package]
name = "colors"
version = "0.1.0"
edition = "2024"
[dependencies]
z3 = "0.19.15"