use std::collections::HashMap;

// Define types for readability
type Action = usize; // Action indices: 0 = fold/pass, 1 = bet/raise
type StrategyMap = HashMap<String, Vec<f64>>; // Stores cumulative average strategies
type RegretMap = HashMap<String, Vec<f64>>; // Stores accumulated regrets

#[derive(Debug, Clone)]
struct GameState {
    history: String, // Keeps track of the sequence of moves
    current_player: usize, // Indicates which player's turn it is (0 or 1)
}

impl GameState {
    // Checks if the game has reached a terminal state
    fn is_terminal(&self) -> bool {
        // Termination occurs if history is sufficiently long or ends with certain action
        self.history.len() > 4 || self.history.ends_with("1")
    }

    // Calculates utility for each player at terminal states
    fn utility(&self, player: usize) -> f64 {
        // Utility is positive for player 0 winning and negative for player 1
        if self.history.ends_with("1") {
            if player == 0 { 1.0 } else { -1.0 }
        } else {
            0.0
        }
    }

    // Returns the list of legal actions from the current state
    fn get_legal_actions(&self) -> Vec<Action> {
        vec![0, 1] // Two actions possible: fold/pass (0) and bet/raise (1)
    }
}

// Manages training through CFR by maintaining regrets and strategies.
struct CFRTrainer {
    regret_map: RegretMap, // Holds regret values for each state-action pair
    strategy_sum_map: StrategyMap, // Holds cumulative strategies for each state
}

impl CFRTrainer {
    // Initializes CFRTrainer with empty maps
    fn new() -> Self {
        CFRTrainer {
            regret_map: HashMap::new(),
            strategy_sum_map: HashMap::new(),
        }
    }

    // Conducts training over a number of iterations
    fn train(&mut self, iterations: usize) {
        for _ in 0..iterations {
            let mut game_state = GameState {
                history: String::new(),
                current_player: 0,
            };
            self.cfr(&mut game_state, 1.0, 1.0); 
        }
        self.compute_average_strategy();
    }

    // Performs the recursive CFR to compute utilities and update regrets
    fn cfr(&mut self, state: &mut GameState, p0: f64, p1: f64) -> f64 {
        if state.is_terminal() {
            // Terminal utility directly impacts regret updates, simplicity aids learning
            return state.utility(state.current_player);
        }

        let player = state.current_player;
        let current_history = state.history.clone();
        let legal_actions = state.get_legal_actions();
        let strategy = self.calculate_strategy(&current_history);

        let mut util = vec![0.0; legal_actions.len()];
        let mut node_util = 0.0;

        for (i, &_action) in legal_actions.iter().enumerate() {
            let mut next_state = state.clone();
            next_state.history.push_str(&i.to_string());
            next_state.current_player = 1 - player;

            // Recursive utility calculation and backward regret propagation
            util[i] = match player {
                0 => -self.cfr(&mut next_state, p0 * strategy[i], p1),
                1 => -self.cfr(&mut next_state, p0, p1 * strategy[i]),
                _ => 0.0,
            };
            node_util += strategy[i] * util[i];
        }

        // Update regrets for current strategy based on utility differences
        let regret_entry = self.regret_map.entry(current_history.clone()).or_insert(vec![0.0; legal_actions.len()]);
        for i in 0..legal_actions.len() {
            let regret = util[i] - node_util;
            if player == 0 {
                regret_entry[i] += p1 * regret;
            } else {
                regret_entry[i] += p0 * regret;
            }
        }

        node_util
    }

    // Calculates strategy given current regrets, norms, and updates strategy sums
    fn calculate_strategy(&mut self, history: &String) -> Vec<f64> {
        let legal_actions = 2;
        let mut strategy = vec![0.0; legal_actions];
        let mut normalizing_sum = 0.0;

        let regret_entry = self.regret_map.entry(history.clone()).or_insert(vec![0.0; legal_actions]);
        for i in 0..legal_actions {
            strategy[i] = (regret_entry[i]).max(0.0);
            normalizing_sum += strategy[i];
        }

        // Normalize the strategy; if all regrets are non-positive, assign equal probabilities
        if normalizing_sum > 0.0 {
            for s in &mut strategy {
                *s /= normalizing_sum;
            }
        } else {
            for s in &mut strategy {
                *s = 1.0 / legal_actions as f64;
            }
        }

        // Update strategy sum for averaging
        let strategy_sum = self.strategy_sum_map.entry(history.clone()).or_insert(vec![0.0; legal_actions]);
        for i in 0..legal_actions {
            strategy_sum[i] += strategy[i];
        }

        strategy
    }

    // Computes and prints average strategies across all simulation iterations
    fn compute_average_strategy(&mut self) {
        for (history, strategy_sum) in &self.strategy_sum_map {
            let sum = strategy_sum.iter().sum::<f64>();
            if sum > 0.0 {
                let avg_strategy = strategy_sum.iter().map(|v| v / sum).collect::<Vec<f64>>();
                println!("History: {} => Average Strategy: {:?}", history, avg_strategy);
            }
        }
    }
}

fn main() {
    // Initializes and runs CFR training for a specified number of iterations
    let mut trainer = CFRTrainer::new();
    trainer.train(10000);
}

edit the code, so it will work for NLTH poker