path: root/src/strategy/minimax.rs

use crate::command::{Action, Command};
use crate::game::{GameBoard, SimulationOutcome};

use std::cmp;
use std::collections::HashMap;
use std::ops::*;
use time::{Duration, PreciseTime};

pub fn choose_move(
    state: &GameBoard,
    previous_root: Option<Node>,
    start_time: PreciseTime,
    max_time: Duration,
) -> (Command, Node) {
    let mut root_node = match previous_root {
        None => Node {
            score_sum: ScoreSum::new(),
            player_score_sums: [HashMap::new(), HashMap::new()],
            unexplored: move_combos(state),
            children: HashMap::new(),
        },
        Some(mut node) => node
            .children
            .drain()
            .map(|(_k, n)| n)
            .find(|_n| false) // TODO: Use the last player / opponent move and worm positions to use this cache.
            .unwrap_or_else(|| {
                eprintln!("Previous round did not appear in the cache");
                Node {
                    score_sum: ScoreSum::new(),
                    player_score_sums: [HashMap::new(), HashMap::new()],
                    unexplored: move_combos(state),
                    children: HashMap::new(),
                }
            }),
    };

    while start_time.to(PreciseTime::now()) < max_time {
        let _ = expand_tree(&mut root_node, &state);
    }

    eprintln!("Number of simulations: {}", root_node.score_sum.visit_count);
    for (command, score_sum) in &root_node.player_score_sums[0] {
        eprintln!(
            "{} = {} ({} visits)",
            command,
            score_sum.avg().val,
            score_sum.visit_count
        );
    }

    let chosen_command = best_player_move(&root_node);

    root_node
        .children
        .retain(|[c1, _], _| *c1 == chosen_command);

    (chosen_command, root_node)
}

pub struct Node {
    score_sum: ScoreSum,
    player_score_sums: [HashMap<Command, ScoreSum>; 2],
    unexplored: Vec<[Command; 2]>,
    children: HashMap<[Command; 2], Node>,
}

#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
struct Score {
    val: f32,
}

impl AddAssign for Score {
    fn add_assign(&mut self, other: Self) {
        self.val = self.val + other.val;
    }
}

impl Div<u32> for Score {
    type Output = Self;
    fn div(self, other: u32) -> Self {
        Score {
            val: self.val / other as f32,
        }
    }
}

impl cmp::Eq for Score {}
impl cmp::Ord for Score {
    fn cmp(&self, other: &Score) -> cmp::Ordering {
        self.val
            .partial_cmp(&other.val)
            .unwrap_or(cmp::Ordering::Equal)
    }
}

struct ScoreSum {
    sum: Score,
    visit_count: u32,
}

impl ScoreSum {
    fn new() -> ScoreSum {
        ScoreSum {
            sum: Score { val: 0. },
            visit_count: 0,
        }
    }
    fn with_initial(score: Score) -> ScoreSum {
        ScoreSum {
            sum: score,
            visit_count: 1,
        }
    }
    fn avg(&self) -> Score {
        self.sum / self.visit_count
    }
}

impl AddAssign<Score> for ScoreSum {
    fn add_assign(&mut self, other: Score) {
        self.sum += other;
        self.visit_count = self.visit_count.saturating_add(1);
    }
}

fn expand_tree(node: &mut Node, state: &GameBoard) -> Score {
    if state.outcome != SimulationOutcome::Continue {
        score(state)
    } else if let Some(commands) = node.unexplored.pop() {
        // TODO: Explore preemptively doing the rollout
        let mut new_state = state.clone();
        new_state.simulate(commands);
        let score = score(&new_state);
        let unexplored = if new_state.outcome == SimulationOutcome::Continue {
            move_combos(&new_state)
        } else {
            Vec::new()
        };

        let new_node = Node {
            score_sum: ScoreSum::with_initial(score),
            player_score_sums: [HashMap::new(), HashMap::new()],
            unexplored,
            children: HashMap::new(),
        };
        node.children.insert(commands, new_node);
        update(node, commands, score);

        if false && node.unexplored.is_empty() {
            // TODO: Prune dominated moves
            // Any move that, regardless of opponent move, there's some other move that gives a better score.
            // Prune for both players.
            // TODO: Prune opponent moves too
            // TODO: Propagate score up
            // TODO: prune moves from the explored states
            let mut to_prune: Vec<Command> = Vec::new();
            for player_move in node.player_score_sums[0].keys() {
                let mut better_moves: Vec<Command> = node.player_score_sums[0]
                    .keys()
                    .filter(|m| **m != *player_move)
                    .cloned()
                    .collect();
                for opponent_move in node.player_score_sums[1].keys() {
                    let baseline_score = node
                        .children
                        .get(&[*player_move, *opponent_move])
                        .unwrap()
                        .score_sum
                        .avg();

                    better_moves.retain(|possible_better_move| {
                        let possibly_better_score = node
                            .children
                            .get(&[*possible_better_move, *opponent_move])
                            .unwrap()
                            .score_sum
                            .avg();
                        baseline_score >= possibly_better_score
                    });
                }
                if !better_moves.is_empty() {
                    to_prune.push(player_move.clone());
                }
            }
            for pruning in to_prune {
                node.player_score_sums[0].remove(&pruning);
            }
        }

        score
    } else {
        let commands = choose_existing(node);
        // TODO: Is there anyway I can avoid this clone? Clone before
        // calling update_tree and just pass ownership all the way
        // down.
        let mut new_state = state.clone();
        new_state.simulate(commands);
        let score = expand_tree(
            node.children
                .get_mut(&commands)
                .expect("The existing node hasn't been tried yet"),
            &new_state,
        );
        update(node, commands, score);
        score
    }
}

fn move_combos(state: &GameBoard) -> Vec<[Command; 2]> {
    let player_moves = pruned_moves(state, 0);
    let opponent_moves = pruned_moves(state, 1);
    debug_assert!(!player_moves.is_empty(), "No player moves");
    debug_assert!(!opponent_moves.is_empty(), "No opponent moves");

    let mut result = Vec::with_capacity(player_moves.len() * opponent_moves.len());
    for p in &player_moves {
        for o in &opponent_moves {
            result.push([*p, *o]);
        }
    }

    result
}

fn best_player_move(node: &Node) -> Command {
    node.player_score_sums[0]
        .iter()
        .max_by_key(|(_command, score_sum)| score_sum.avg())
        .map(|(command, _score_sum)| *command)
        .unwrap_or_else(|| Command::new(Action::DoNothing))
}

fn score(state: &GameBoard) -> Score {
    let max_health =
        (state.players[0].max_worm_health() - state.players[1].max_worm_health()) as f32;
    let points = (state.players[0].score() - state.players[1].score()) as f32;

    const MAX_HEALTH_WEIGHT: f32 = 1.;
    const POINTS_WEIGHT: f32 = 0.;
    const VICTORY_WEIGHT: f32 = 3000.;

    // TODO: Try adding new features here, like max worm health, weighted in some way
    // TODO: Distance to dirt heatmap? Probably less relevant these days.
    Score {
        val: match state.outcome {
            SimulationOutcome::PlayerWon(0) => VICTORY_WEIGHT,
            SimulationOutcome::PlayerWon(1) => -VICTORY_WEIGHT,
            _ => max_health * MAX_HEALTH_WEIGHT + points * POINTS_WEIGHT,
        },
    }
}

fn choose_existing(node: &Node) -> [Command; 2] {
    [choose_one_existing(node, 0), choose_one_existing(node, 1)]
}

fn choose_one_existing(node: &Node, player_index: usize) -> Command {
    let ln_n = (node.score_sum.visit_count as f32).ln();
    let c = 100.;
    let multiplier = if player_index == 0 { 1. } else { -1. };
    node.player_score_sums[player_index]
        .iter()
        .max_by_key(|(_command, score_sum)| {
            (multiplier * (score_sum.avg().val + c * (ln_n / score_sum.visit_count as f32).sqrt()))
                as i32
        })
        .map(|(command, _score_sum)| *command)
        .unwrap_or_else(|| Command::new(Action::DoNothing))
}

fn update(node: &mut Node, commands: [Command; 2], score: Score) {
    *node.player_score_sums[0]
        .entry(commands[0])
        .or_insert_with(ScoreSum::new) += score;
    *node.player_score_sums[1]
        .entry(commands[1])
        .or_insert_with(ScoreSum::new) += score;
    node.score_sum += score;
}

fn pruned_moves(state: &GameBoard, player_index: usize) -> Vec<Command> {
    let sim_with_idle_opponent = |cmd| {
        let mut idle_commands = [
            Command::new(Action::DoNothing),
            Command::new(Action::DoNothing),
        ];
        idle_commands[player_index] = cmd;
        let mut state_cpy = state.clone();
        state_cpy.simulate(idle_commands);
        state_cpy
    };

    let mut do_nothing_state = state.clone();
    do_nothing_state.simulate([
        Command::new(Action::DoNothing),
        Command::new(Action::DoNothing),
    ]);

    let opponent_index = GameBoard::opponent(player_index);
    let my_starting_health = do_nothing_state.players[player_index].health();
    let opponent_starting_health = do_nothing_state.players[opponent_index].health();

    state
        .valid_moves(player_index)
        .into_iter()
        .filter(|command| {
            // TODO: Filtering out of snowball moves to only freeze opponents
            // TODO: Filter bombs out to only hurt opponents

            // NB: These rules should pass for doing nothing, otherwise
            // we need some other mechanism for sticking in a do
            // nothing option. Unfortunately, sitting in lava is a situation where this prunes all moves currently :(

            let idle_opponent_state = sim_with_idle_opponent(*command);
            let hurt_self = idle_opponent_state.players[player_index].health() < my_starting_health;
            let hurt_opponent =
                idle_opponent_state.players[opponent_index].health() < opponent_starting_health;
            let is_select = command.worm.is_some();

            !hurt_self && (!is_select || hurt_opponent)
        })
        .collect()
}