// Welcome to
// __________         __    __  .__                               __
// \______   \_____ _/  |__/  |_|  |   ____   ______ ____ _____  |  | __ ____
//  |    |  _/\__  \\   __\   __\  | _/ __ \ /  ___//    \\__  \ |  |/ // __ \
//  |    |   \ / __ \|  |  |  | |  |_\  ___/ \___ \|   |  \/ __ \|    <\  ___/
//  |________/(______/__|  |__| |____/\_____>______>___|__(______/__|__\\_____>
//
// This file can be a nice home for your Battlesnake logic and helper functions.
//
// To get you started we've included code to prevent your Battlesnake from moving backwards.
// For more info see docs.battlesnake.com

use core::f64;
use std::{
    collections::{BTreeMap, HashMap},
    sync::{Arc, LazyLock, Mutex},
    time::{Duration, Instant},
};

use log::{error, info, warn};
use ordered_float::OrderedFloat;
use rand::{prelude::*, thread_rng};
use serde_json::{json, Value};

use crate::{
    simulation::{self, SnakeToken},
    Action, Battlesnake, Board, Direction, Game,
};

// info is called when you create your Battlesnake on play.battlesnake.com
// and controls your Battlesnake's appearance
// TIP: If you open your Battlesnake URL in a browser you should see this data
#[must_use]
pub fn info() -> Value {
    info!("INFO");

    json!({
        "apiversion": "1",
        "author": "der-informatiker",
        "color": "#00FFEE",
        "head": "smart-caterpillar",
        "tail": "mouse",
    })
}

#[derive(Debug, Clone)]
struct GameInfo {
    calculation_time: Arc<Mutex<Duration>>,
    token_mapping: Arc<BTreeMap<String, SnakeToken>>,
    my_token: SnakeToken,
    tree: Arc<Mutex<Option<Node>>>,
}

static GAME_INFOS: LazyLock<Mutex<HashMap<(String, String), GameInfo>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

// start is called when your Battlesnake begins a game
pub fn start(game: &Game, _turn: i32, board: &Board, you: &Battlesnake) {
    info!("GAME START");
    let token_mapping = Arc::new(SnakeToken::from_board(board));
    let my_token = token_mapping[&you.id];
    let Ok(mut game_infos) = GAME_INFOS.lock() else {
        error!("unable to lock game infos");
        return;
    };
    game_infos.insert(
        (game.id.clone(), you.id.clone()),
        GameInfo {
            calculation_time: Arc::new(Mutex::new(Duration::from_millis(
                u64::from(game.timeout) / 2,
            ))),
            token_mapping,
            my_token,
            tree: Arc::new(Mutex::new(None)),
        },
    );
}

// end is called when your Battlesnake finishes a game
pub fn end(game: &Game, turn: i32, _board: &Board, you: &Battlesnake) {
    info!("GAME OVER after {turn} turns");
    let Ok(mut game_infos) = GAME_INFOS.lock() else {
        error!("unable to lock game infos");
        return;
    };
    game_infos.remove(&(game.id.clone(), you.id.clone()));
}

// move is called on every turn and returns your next move
// Valid moves are "up", "down", "left", or "right"
// See https://docs.battlesnake.com/api/example-move for available data
pub fn get_move(
    game: &Game,
    turn: i32,
    board: &Board,
    you: &Battlesnake,
    start: &Instant,
) -> Option<Action> {
    let calc_start = Instant::now();
    if calc_start - *start > Duration::from_millis(10) {
        error!(
            "The calculation was started long after the request ({}ms)",
            (calc_start - *start).as_millis()
        );
    }
    let game_info = GAME_INFOS
        .lock()
        .ok()
        .and_then(|guard| guard.get(&(game.id.clone(), you.id.clone())).cloned())
        .unwrap_or_else(|| {
            let token_mapping = Arc::new(SnakeToken::from_board(board));
            let my_token = token_mapping[&you.id];
            GameInfo {
                calculation_time: Arc::new(Mutex::new(Duration::from_millis(
                    u64::from(game.timeout) / 2,
                ))),
                token_mapping,
                my_token,
                tree: Arc::new(Mutex::new(None)),
            }
        });

    let board = simulation::Board::from_game_board(
        board,
        &game_info.token_mapping,
        turn,
        game.ruleset.settings.food_spawn_chance,
        game.ruleset.settings.minimum_food,
        game.ruleset.name == "constrictor",
    );
    let possible_actions = board.possible_actions().get(&game_info.my_token).cloned()?;
    if possible_actions.is_empty() {
        info!("No movement options in turn {turn}");
        return None;
    }

    // do some latency compensation
    let deadline = *start
        + game_info.calculation_time.lock().map_or_else(
            |_| Duration::from_millis(u64::from(game.timeout) / 2),
            |mut guard| {
                let target_latency = game.timeout / 2;
                let latency = you.latency.parse().unwrap_or_else(|e| {
                    warn!("Unable to parse latency: {e}");
                    target_latency
                });
                let last_computation_time = u32::try_from(guard.as_millis()).unwrap_or(0);
                let computation_time =
                    (last_computation_time + target_latency).saturating_sub(latency);
                *guard =
                    Duration::from_millis(u64::from(computation_time.clamp(1, target_latency)));
                *guard
            },
        );

    let mut tree_guard = game_info.tree.lock();
    let tree = match tree_guard {
        Err(ref e) => {
            error!("unable to lock tree: {e}");
            None
        }
        Ok(ref mut guard) => guard.as_mut(),
    };
    let mut tree = tree
        .and_then(|node| {
            let snake_length_direction: BTreeMap<_, _> = board
                .snakes()
                .map(|snake| {
                    let length = board.snake_length(snake).unwrap_or_default();
                    let action = board.last_action(snake).unwrap_or(Direction::Up);
                    (snake, (action, length))
                })
                .collect();
            let node_key = node
                .childs
                .keys()
                .find(|child| {
                    child.iter().all(|(snake, (direction, length))| {
                        snake_length_direction
                            .get(snake)
                            .copied()
                            .unwrap_or((*direction, 0))
                            == (*direction, *length)
                    })
                })?
                .clone();
            let node = node.childs.remove(&node_key)?;
            info!(
                "using previous node with {} simulations",
                node.statistic.played
            );
            Some(node)
        })
        .unwrap_or_default();

    while Instant::now() < deadline {
        let mut board = board.clone();
        if game.ruleset.name == "solo" {
            let _ = tree.monte_carlo_step_solo(&mut board, &deadline);
        } else {
            let _ = tree.monte_carlo_step(&mut board, &deadline);
        }
    }

    let actions = tree.child_statistics.entry(game_info.my_token).or_default();

    info!("actions {}: {actions:?}", you.name);

    #[allow(clippy::cast_precision_loss)]
    let chosen = actions
        .iter()
        .max_by_key(|(_, stat)| OrderedFloat(stat.won as f64 / stat.played as f64))
        .map(|(direction, _)| *direction)
        .or_else(|| possible_actions.iter().choose(&mut thread_rng()).copied())?;

    if let Ok(ref mut guard) = tree_guard {
        **guard = Some(tree);
    }
    std::mem::drop(tree_guard);

    info!(
        "DIRECTION {turn}: {chosen:?} after {}ms ({})",
        start.elapsed().as_millis(),
        you.name,
    );
    Some(Action {
        r#move: chosen,
        shout: None,
    })
}

#[derive(Debug, PartialEq, Eq, Clone, Default)]
struct Statistics {
    /// Number of times this node was simulated
    played: usize,
    /// Number of times this node was simulated and the agent has won.
    won: BTreeMap<SnakeToken, usize>,
}

#[derive(Debug, PartialEq, Eq, Clone, Default)]
struct ActionStatistic {
    played: usize,
    won: usize,
}

struct DeadlineError;

#[derive(Debug, PartialEq, Eq, Clone, Default)]
struct Node {
    statistic: Statistics,
    child_statistics: BTreeMap<SnakeToken, BTreeMap<Direction, ActionStatistic>>,
    childs: BTreeMap<BTreeMap<SnakeToken, (Direction, usize)>, Node>,
}

impl Node {
    /// Performs one monte carlo simulation step
    ///
    /// Returns the snake that has won the simulation
    fn monte_carlo_step(
        &mut self,
        board: &mut simulation::Board,
        deadline: &Instant,
    ) -> Result<Option<SnakeToken>, DeadlineError> {
        let stop_condition =
            |board: &simulation::Board| board.alive_snakes() <= 1 || Instant::now() >= *deadline;
        let winner = if stop_condition(board) {
            if Instant::now() >= *deadline {
                return Err(DeadlineError);
            }
            board.snakes().next()
        } else if self.statistic.played == 0 {
            // We didn't simulate a game for this node yet. Do that
            board.simulate_until(&mut thread_rng(), stop_condition);
            if Instant::now() >= *deadline {
                return Err(DeadlineError);
            }
            board.snakes().next()
        } else {
            // select a node to simulate
            let possible_actions = board.possible_actions();

            let actions = possible_actions
                .iter()
                .filter_map(|(token, actions)| {
                    let statistics = self.child_statistics.entry(*token).or_default();
                    let selected = actions.iter().copied().max_by_key(|direction| {
                        let statistics = statistics.entry(*direction).or_default();
                        if statistics.played == 0 {
                            return OrderedFloat(f64::INFINITY);
                        }
                        #[allow(clippy::cast_precision_loss)]
                        let exploitation = statistics.won as f64 / statistics.played as f64;
                        #[allow(clippy::cast_precision_loss)]
                        let exploration = f64::consts::SQRT_2
                            * f64::sqrt(
                                f64::ln(self.statistic.played as f64) / statistics.played as f64,
                            );
                        OrderedFloat(exploitation + exploration)
                    })?;
                    Some((*token, selected))
                })
                .collect();

            board.simulate_actions(&actions, &mut thread_rng());
            let map_actions = actions
                .iter()
                .map(|(&snake, &action)| (snake, (action, board.snake_length(snake).unwrap_or(0))))
                .collect();
            let winner = self
                .childs
                .entry(map_actions)
                .or_default()
                .monte_carlo_step(board, deadline)?;

            // update child statistics
            for (token, action) in &actions {
                let entry = self
                    .child_statistics
                    .entry(*token)
                    .or_default()
                    .entry(*action)
                    .or_default();
                entry.played += 1;
                if Some(*token) == winner {
                    entry.won += 1;
                }
            }

            winner
        };
        self.statistic.played += 1;
        if let Some(token) = winner {
            self.statistic
                .won
                .entry(token)
                .and_modify(|won| *won += 1)
                .or_insert(1);
        }
        Ok(winner)
    }

    /// Performs one monte carlo simulation step for a solo game
    ///
    /// Returns the lengths before death
    fn monte_carlo_step_solo(
        &mut self,
        board: &mut simulation::Board,
        deadline: &Instant,
    ) -> Result<BTreeMap<SnakeToken, usize>, DeadlineError> {
        let lengths = if self.statistic.played == 0 {
            // We didn't simulate a game for this node yet. Do that
            let mut lengths: BTreeMap<_, _> = board
                .snakes()
                .filter_map(|snake| Some((snake, board.snake_length(snake)?)))
                .collect();
            board.simulate_until(&mut thread_rng(), |board| {
                if Instant::now() >= *deadline {
                    return true;
                }
                for snake in board.snakes() {
                    if let Some(length) = board.snake_length(snake) {
                        lengths.insert(snake, length);
                    }
                }
                board.alive_snakes() == 0
            });
            if Instant::now() >= *deadline {
                return Err(DeadlineError);
            }
            lengths
        } else {
            // select a node to simulate
            let possible_actions = board.possible_actions();

            let actions = possible_actions
                .iter()
                .filter_map(|(token, actions)| {
                    let statistics = self.child_statistics.entry(*token).or_default();
                    let selected = actions.iter().copied().max_by_key(|direction| {
                        let statistics = statistics.entry(*direction).or_default();
                        if statistics.played == 0 {
                            return OrderedFloat(f64::INFINITY);
                        }
                        #[allow(clippy::cast_precision_loss)]
                        let exploitation = statistics.won as f64
                            / board.spaces() as f64
                            / statistics.played as f64;
                        #[allow(clippy::cast_precision_loss)]
                        let exploration = f64::consts::SQRT_2
                            * f64::sqrt(
                                f64::ln(self.statistic.played as f64) / statistics.played as f64,
                            );
                        OrderedFloat(exploitation + exploration)
                    })?;
                    Some((*token, selected))
                })
                .collect();

            if Instant::now() >= *deadline {
                return Err(DeadlineError);
            }
            board.simulate_actions(&actions, &mut thread_rng());
            let map_actions = actions
                .iter()
                .map(|(&snake, &action)| (snake, (action, board.snake_length(snake).unwrap_or(0))))
                .collect();
            let lengths = self
                .childs
                .entry(map_actions)
                .or_default()
                .monte_carlo_step_solo(board, deadline)?;

            // update child statistics
            for (token, action) in &actions {
                let entry = self
                    .child_statistics
                    .entry(*token)
                    .or_default()
                    .entry(*action)
                    .or_default();
                entry.played += 1;
                if let Some(length) = lengths.get(token) {
                    entry.won += length;
                }
            }

            lengths
        };
        self.statistic.played += 1;
        for (token, length) in &lengths {
            self.statistic
                .won
                .entry(*token)
                .and_modify(|won| *won += length)
                .or_insert(*length);
        }
        Ok(lengths)
    }
}
