major-pickems-sim 0.3.0

use std::{
    collections::{BinaryHeap, HashSet},
    iter::Sum,
    ops::Add,
};

use crate::{
    datatypes::Index,
    reporting::{AssessReport, BasicReport, Report},
    simulation::{Simulation, SwissSystem},
};

/// Candidate team for one pick category.
///
/// Equality and hashing intentionally use only the team index so a candidate can
/// move between probability-ranked pools without becoming a distinct set entry.
#[derive(Debug, Clone, Copy)]
struct Candidate {
    index: Index,
    probability: f32,
}

impl std::hash::Hash for Candidate {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.index.hash(state);
    }
}

impl Eq for Candidate {}

impl PartialEq for Candidate {
    fn eq(&self, other: &Self) -> bool {
        self.index == other.index
    }
}

impl Ord for Candidate {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.probability.total_cmp(&other.probability)
    }
}

impl PartialOrd for Candidate {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

/// Report for selecting pick recommendations from basic outcome probabilities.
#[derive(Debug, Clone, Copy, Default)]
pub struct PicksReport {
    pub basic: BasicReport,
}

impl Add for PicksReport {
    type Output = Self;

    fn add(mut self, rhs: Self) -> Self::Output {
        self.basic = self.basic + rhs.basic;
        self
    }
}

impl Sum for PicksReport {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Self::default(), |acc, report| acc + report)
    }
}

impl Report for PicksReport {
    fn update(&mut self, ss: &SwissSystem) {
        self.basic.update(ss);
    }

    fn format(&self, sim: &Simulation) -> String {
        let [three_zero, advancing, zero_three] = self.basic.calculate_probabilities(sim);

        // Build max-heaps for each pick category so the highest probability
        // candidate is always popped first.
        let candidates = |probabilities: [f32; 16]| -> BinaryHeap<Candidate> {
            probabilities
                .iter()
                .enumerate()
                .map(|(i, p)| Candidate {
                    index: unsafe { Index::from_usize(i) },
                    probability: *p,
                })
                .collect::<BinaryHeap<_>>()
        };

        let mut three_zero_candidates = candidates(three_zero);
        let mut advancing_candidates = candidates(advancing);
        let mut zero_three_candidates = candidates(zero_three);

        // Start with the best advancement picks.
        let mut advancing_picks = HashSet::new();

        for _ in 0..6 {
            if let Some(candidate) = advancing_candidates.pop() {
                advancing_picks.insert(candidate);
            }
        }

        // Select 0-3 picks independently because they do not overlap with the
        // advancement categories in valid terminal outcomes.
        let mut zero_three_picks = HashSet::new();

        for _ in 0..2 {
            if let Some(candidate) = zero_three_candidates.pop() {
                zero_three_picks.insert(candidate);
            }
        }

        // Optimise 3-0 picks by swapping previous advancement picks to maximise win probability.
        let mut three_zero_picks = HashSet::new();

        while three_zero_picks.len() < 2 {
            let mut swap_candidates = BinaryHeap::new();

            while let Some(candidate) = three_zero_candidates.peek()
                && advancing_picks.contains(candidate)
            {
                swap_candidates.push(three_zero_candidates.pop().unwrap());
            }

            match (
                three_zero_candidates.pop(),
                swap_candidates.pop(),
                advancing_candidates.pop(),
            ) {
                // There are teams in all relevant candidate pools.
                (Some(next_three_zero), Some(next_swap), Some(next_advancing))
                    if advancing_picks.contains(&next_swap) =>
                {
                    let swap_advancing = advancing_picks.get(&next_swap).unwrap();

                    // Compare the lost advancement probability against the
                    // gained 3-0 probability from making the swap.
                    let cost = next_advancing.probability - swap_advancing.probability;
                    let reward = next_swap.probability - next_three_zero.probability;

                    // Repopulate candidate pools with candidates that remain
                    // unselected so the next loop considers them again.
                    if reward > cost {
                        three_zero_picks.insert(next_swap);
                        advancing_picks.remove(&next_swap);
                        advancing_picks.insert(next_advancing);
                        three_zero_candidates.push(next_three_zero);
                    } else {
                        three_zero_picks.insert(next_three_zero);
                        advancing_candidates.push(next_advancing);
                    }
                }
                // There are only teams left in the 3-0 candidate pool, and
                // unviable candidates in the advancement pool.
                (Some(next_three_zero), None, Some(next_advancing)) => {
                    three_zero_picks.insert(next_three_zero);
                    advancing_candidates.push(next_advancing);
                }
                // There are only teams left in the 3-0 candidate pool.
                (Some(next_three_zero), None, None) => {
                    three_zero_picks.insert(next_three_zero);
                }
                // The current state no longer makes any sense, either the 3-0 pool is empty
                // or the advancement picks don't contain the next swap candidate.
                state => {
                    unreachable!("invalid state for picking 3-0 teams:\n\n{state:#?}");
                }
            }
        }

        // The picks at this point are potentially still suboptimal, in future I want to further optimise picks
        // using A* or similar to explore more combinations.

        // Assess picks through a second simulation pass using the chosen teams.
        let assessment = sim.run(AssessReport::new(
            three_zero_picks.iter().map(|c| c.index),
            advancing_picks.iter().map(|c| c.index),
            zero_three_picks.iter().map(|c| c.index),
        ));

        // Format results into a string.
        let format_picks = |out: &mut Vec<String>, picks: &HashSet<Candidate>| {
            let mut picks = picks
                .iter()
                .map(|i| (&sim.teams.names[i.index.to_usize()], i.probability * 100.0))
                .collect::<Vec<_>>();

            picks.sort_by(|(_, a), (_, b)| b.total_cmp(a));

            for (i, (name, p)) in picks.into_iter().enumerate() {
                out.push(format!(
                    "{num:<4}{name:<20}{p:>6.1}%",
                    num = format!("{}.", i + 1),
                ));
            }
        };

        let mut out = vec![String::from("\n3-0 picks:")];
        format_picks(&mut out, &three_zero_picks);
        out.push(String::from("\n3-1 or 3-2 picks:"));
        format_picks(&mut out, &advancing_picks);
        out.push(String::from("\n0-3 picks:"));
        format_picks(&mut out, &zero_three_picks);
        out.push(assessment.format(sim));
        out.join("\n")
    }
}