use std::{
    collections::HashMap,
    path::{Path, PathBuf},
};

use clap::Parser;
use expr::LinExpr;
use grb::*;
use itertools::Itertools;
use ModelSense::Minimize;

fn read_input(filename: impl AsRef<Path>) -> Vec<Vec<u8>> {
    let text = std::fs::read_to_string(filename).unwrap();

    let mut r = Vec::new();

    for line in text.lines().skip(1) {
        r.push(line.as_bytes().to_vec());
    }

    r
}

#[derive(Debug, Parser)]
struct Args {
    filename: PathBuf,
    outfilename: PathBuf,
}

fn solve(input: &[Vec<u8>]) -> String {
    let mut columns = Vec::new();

    for i in 0..input[0].len() {
        let mut nextfree: u8 = 0;
        let mut c = Vec::new();

        for l in 0..input.len() {
            if let Some(n) = input[0..l].iter().position(|r| r[i] == input[l][i]) {
                c.push(c[n]);
            } else {
                c.push(nextfree);
                nextfree += 1;
            }
        }

        columns.push(c);
    }

    let counts = columns.iter().counts();

    let mut model = Model::new("model1").unwrap();

    let mut variables = HashMap::new();

    for (&k, &c) in &counts {
        let mut v = Vec::new();
        let mut sum = Expr::Linear(LinExpr::new());
        for i in 0..=(k.iter().max().copied().unwrap()) {
            let var = add_intvar!(model, name: &format!("{:?}:{}", k, i), bounds: 0..).unwrap();

            sum = sum + var;
            v.push(var);
        }

        model.add_constr(&format!("{:?}", k), c!(sum == c)).unwrap();
        variables.insert(k, v);
    }

    let d = add_intvar!(model, name: "D", bounds: 0..).unwrap();

    for i in 0..input.len() {
        let mut sum = Expr::Linear(LinExpr::new());

        for &k in counts.keys() {
            for j in 0..=*k.iter().max().unwrap() {
                if k[i] != j {
                    sum = sum + variables[k][j as usize];
                }
            }
        }

        model.add_constr(&format!("{:?}", i), c!(sum <= d)).unwrap();
    }

    model.set_objective(d, Minimize).unwrap();

    model.update().unwrap();

    // model.write("test.lp").unwrap();

    model.optimize().unwrap();

    // for vars in &variables {
    //     for v in vars {
    //         dbg!(model.get_obj_attr(attr::X, v).unwrap());
    //     }
    // }

    let mut result = String::new();

    for (i, col) in columns.iter().enumerate() {
        let t = model
            .get_obj_attr_batch(attr::X, variables[col].clone())
            .unwrap();

        let mut p = columns[0..i].iter().filter(|&c| c == col).count();

        let mut j = 0;
        while t[j] as usize <= p {
            p -= t[j] as usize;
            j += 1;
        }

        let l = col.iter().position(|&c| c as usize == j).unwrap();

        // dbg!(i, col, t, j, l, char::from_u32(input[l][i] as u32).unwrap());

        result.push(char::from_u32(input[l][i] as u32).unwrap());
    }

    result
}

fn main() {
    let args = Args::parse();

    let input = read_input(args.filename);

    let r = solve(&input);

    std::fs::write(args.outfilename, r).unwrap();
}
