minimax/src/agents/util/partials.rs

use itertools::Itertools;
use std::{mem::swap, num::NonZeroU8};

use crate::{board::Board, util::Symb};

/// Returns an iterator of (idx, coords, char_idx, f32) for each empty slot in the listed partials.
/// - idx is the index of this slot in the board string.
/// - coords are the coordinate of this slot's partial
/// - char_idx is the index of this slot in its partial
/// - f32 is the influence of this slot
pub fn free_slots_by_influence(board: &Board) -> Option<Vec<(usize, f32)>> {
	// Fill all empty slots with fives and compute starting value
	let filled = {
		// This should always result in an evaluatable expression,
		// since parenthesis do not exist.
		// (the only way to divide by zero is by doing something like /(5-2+3) )
		let f = Board::from_board(board.get_board().map(|x| match x {
			None => Symb::from_char('5'),
			_ => x,
		}));

		f
	};

	let base = filled.evaluate()?;

	// Test each slot:
	// Increase its value by 1, and record its effect on the
	// expression's total value.
	// This isn't a perfect metric, but it's pretty good.
	let mut slots: Vec<(usize, f32)> = board
		.get_board()
		.iter()
		.enumerate()
		.filter_map(|(i, s)| if s.is_some() { None } else { Some(i) })
		.map(|i| {
			let mut new_tree = filled.clone();
			new_tree.get_board_mut()[i] = Some(Symb::from_char('6').unwrap());

			// This shouldn't ever be None
			(i, new_tree.evaluate().unwrap() - base)
		})
		.collect();

	// Sort by most to least influence
	slots.sort_by(|a, b| {
		b.1.abs()
			.partial_cmp(&a.1.abs())
			.unwrap_or(std::cmp::Ordering::Equal)
	});
	Some(slots)
}

/// Find the minimum or maximum possible value of the given board,
/// without adding any operations. Returns None if we couldn't find
/// a best board.
pub fn best_board_noop(board: &Board, minimize: bool) -> Option<Board> {
	let n_free = board.get_board().iter().filter(|x| x.is_none()).count();

	// TODO: fix zero division
	let available_numbers = (1..=9)
		.map(|x| match x {
			0 => Symb::Zero,
			x => Symb::Number(NonZeroU8::new(x).unwrap()),
		})
		.filter(|x| !board.contains(*x))
		.collect::<Vec<_>>();

	if n_free > available_numbers.len() {
		return None;
	}

	let slots = free_slots_by_influence(board)?;

	let all_symbols = {
		let mut a = {
			// Number of slots we want to minimize
			let mut neg_count = slots.iter().filter(|(_, x)| *x <= 0.0).count();
			// Number of slots we want to maximize
			let mut pos_count = slots.iter().filter(|(_, x)| *x > 0.0).count();

			if minimize {
				swap(&mut neg_count, &mut pos_count);
			}

			available_numbers
				.iter()
				.take(neg_count)
				.chain(available_numbers.iter().rev().take(pos_count).rev())
				.collect_vec()
		};

		if !minimize {
			a.reverse();
		}
		let mut a_iter = a.into_iter();

		slots
			// Group slots with equal weights
			// and count the number of elements in each group
			.iter()
			.group_by(|x| x.1)
			.into_iter()
			.map(|(_, x)| x.count())
			// Generate the digits we should try for each group of
			// equal-weight slots
			.map(|s| {
				(0..s)
					.map(|_| *a_iter.next().unwrap())
					.permutations(s)
					.unique()
					.collect_vec()
			})
			// Now, covert this to an array of all cartesian products
			// of this set of sets
			.multi_cartesian_product()
			.map(|x| x.iter().flatten().cloned().collect_vec())
			// Finally, attach the coordinate of each slot to each symbol
			.map(|v| slots.iter().map(|x| x.0).zip(v).collect_vec())
			.collect_vec()
	};

	let mut best_board = None;
	let mut best_value = None;
	for i_iter in all_symbols {
		let mut tmp_board = board.clone();
		for (i, s) in i_iter {
			tmp_board.get_board_mut()[i] = Some(s);
		}

		let val = tmp_board.evaluate();

		if let Some(val) = val {
			if minimize {
				if best_value.is_none() || val < best_value.unwrap() {
					best_value = Some(val);
					best_board = Some(tmp_board)
				}
			} else {
				if best_value.is_none() || val > best_value.unwrap() {
					best_value = Some(val);
					best_board = Some(tmp_board)
				}
			}
		}
	}

	best_board
}