// 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::HashMap, time::Instant};

use battlesnake_game_types::{
    compact_representation::standard::CellBoard4Snakes11x11,
    types::{
        build_snake_id_map, LengthGettableGame, Move, RandomReasonableMovesGame,
        ReasonableMovesGame, SimulableGame, SimulatorInstruments, SnakeIDMap, SnakeId,
        VictorDeterminableGame, YouDeterminableGame,
    },
    wire_representation::Game,
};
use log::{error, info};
use ordered_float::OrderedFloat;
use rand::{prelude::*, thread_rng};
use rocket::time::{ext::NumericalDuration, Duration};
use serde_json::{json, Value};

// 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)]
pub struct GameState {
    calculation_time: Duration,
    snake_id_map: SnakeIDMap,
}

// start is called when your Battlesnake begins a game
#[must_use]
pub fn start(game: &Game) -> GameState {
    info!("GAME START");
    let snake_id_map = build_snake_id_map(game);
    let calculation_time = (game.game.timeout / 2).milliseconds();

    GameState {
        calculation_time,
        snake_id_map,
    }
}

// end is called when your Battlesnake finishes a game
pub fn end(game: &Game, state: GameState) {
    std::mem::drop(state);
    info!("GAME OVER after {} turns", game.turn);
}

// 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, state: &mut GameState, start: &Instant) -> Move {
    let calc_start = Instant::now();
    if calc_start - *start > 10.milliseconds() {
        error!(
            "The calculation was started long after the request ({}ms)",
            (calc_start - *start).as_millis()
        );
    }

    let deadline = *start + state.calculation_time;

    let name = game.you.name.clone();
    let turn = game.turn;
    let solo = game.game.ruleset.name == "solo";
    let Ok(board) = CellBoard4Snakes11x11::convert_from_game(game, &state.snake_id_map) else {
        error!("Unable to fit board");
        return Move::Down;
    };

    let mut tree = Node {
        statistic: Statistics {
            played: 0,
            won: HashMap::new(),
        },
        child_statistics: HashMap::new(),
        childs: HashMap::new(),
    };

    while Instant::now() < deadline {
        if solo {
            tree.monte_carlo_solo_step(&board);
        } else {
            tree.monte_carlo_step(&board);
        }
    }

    let actions = tree.child_statistics.entry(*board.you_id()).or_default();

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

    #[allow(clippy::cast_precision_loss)]
    let chosen = actions
        .iter()
        .max_by_key(|(_, stat)| stat.played)
        .map(|(direction, _)| *direction)
        .or_else(|| {
            board
                .random_reasonable_move_for_each_snake(&mut thread_rng())
                .find(|(snake_id, _)| snake_id == board.you_id())
                .map(|(_, direction)| direction)
        })
        .unwrap_or(Move::Down);

    info!(
        "DIRECTION {turn}: {chosen:?} after {}ms ({name})",
        start.elapsed().as_millis(),
    );
    chosen
}

#[derive(Debug, Clone, Copy)]
struct Instruments;

impl SimulatorInstruments for Instruments {
    fn observe_simulation(&self, _duration: std::time::Duration) {}
}

#[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: HashMap<SnakeId, usize>,
}

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

#[derive(Debug, PartialEq, Eq, Clone, Default)]
struct Node {
    statistic: Statistics,
    child_statistics: HashMap<SnakeId, HashMap<Move, ActionStatistic>>,
    childs: HashMap<[Option<(Move, u16)>; 4], Node>,
}

impl Node {
    /// Performs one monte carlo simulation step
    ///
    /// Returns the snake that has won the simulation
    fn monte_carlo_step(&mut self, board: &CellBoard4Snakes11x11) -> Option<SnakeId> {
        let stop_condition = CellBoard4Snakes11x11::is_over;
        let winner = if stop_condition(board) {
            board.get_winner()
        } else if self.statistic.played == 0 {
            // We didn't simulate a game for this node yet. Do that
            let mut board = *board;
            while !stop_condition(&board) {
                let rng = &mut thread_rng();
                let moves = board.random_reasonable_move_for_each_snake(rng);
                let (_, new_board) = board
                    .simulate_with_moves(
                        &Instruments,
                        moves.map(|(snake_id, direction)| (snake_id, [direction])),
                    )
                    .next()
                    .unwrap();
                board = new_board;
            }
            board.get_winner()
        } else {
            // select a node to simulate
            let possible_actions = board.reasonable_moves_for_each_snake();

            let actions = possible_actions
                .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::<Vec<_>>();

            let (_, board) = board
                .simulate_with_moves(&Instruments, actions.iter().copied())
                .next()
                .unwrap();
            let mut map_actions = [None; 4];
            for (i, action) in map_actions.iter_mut().enumerate() {
                *action = actions
                    .iter()
                    .find(|(snake_id, _)| snake_id.as_usize() == i)
                    .and_then(|(snake_id, moves)| {
                        Some((*moves.first()?, board.get_length(snake_id)))
                    });
            }
            let winner = self
                .childs
                .entry(map_actions)
                .or_default()
                .monte_carlo_step(&board);

            // update child statistics
            for (token, action) in &actions {
                let entry = self
                    .child_statistics
                    .entry(*token)
                    .or_default()
                    .entry(action[0])
                    .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);
        }
        winner
    }

    /// Performs one monte carlo simulation step for a solo game
    ///
    /// Returns the lengths before death
    fn monte_carlo_solo_step(&mut self, board: &CellBoard4Snakes11x11) -> u16 {
        let stop_condition = |board: &CellBoard4Snakes11x11| board.alive_snake_count() == 0;
        let winner = if self.statistic.played == 0 {
            // We didn't simulate a game for this node yet. Do that
            let mut board = *board;
            while !stop_condition(&board) {
                let moves =
                    board
                        .reasonable_moves_for_each_snake()
                        .filter_map(|(snake_id, moves)| {
                            Some((snake_id, [*moves.choose(&mut thread_rng())?]))
                        });
                let Some((_, new_board)) = board.simulate_with_moves(&Instruments, moves).next()
                else {
                    break;
                };
                if stop_condition(&new_board) {
                    break;
                }
                board = new_board;
            }
            let winner = board.get_length(board.you_id());
            winner
        } else {
            // select a node to simulate
            let possible_actions = board.reasonable_moves_for_each_snake();

            let actions = possible_actions
                .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,
                            )
                            * 11.0
                            * 11.0;
                        OrderedFloat(exploitation + exploration)
                    })?;
                    Some((token, [selected]))
                })
                .collect::<Vec<_>>();

            let (_, new_board) = board
                .simulate_with_moves(&Instruments, actions.iter().copied())
                .next()
                .unwrap();
            let mut map_actions = [None; 4];
            for (i, action) in map_actions.iter_mut().enumerate() {
                *action = actions
                    .iter()
                    .find(|(snake_id, _)| snake_id.as_usize() == i)
                    .and_then(|(snake_id, moves)| {
                        Some((*moves.first()?, new_board.get_length(snake_id)))
                    });
            }
            let winner = if stop_condition(&new_board) {
                board.get_length(board.you_id())
            } else {
                self.childs
                    .entry(map_actions)
                    .or_default()
                    .monte_carlo_solo_step(&new_board)
            };

            // update child statistics
            let entry = self
                .child_statistics
                .entry(*new_board.you_id())
                .or_default()
                .entry(
                    actions
                        .iter()
                        .find(|(snake_id, _)| snake_id == new_board.you_id())
                        .map(|(_, action)| action[0])
                        .unwrap(),
                )
                .or_default();
            entry.played += 1;
            entry.won += usize::from(winner);

            winner
        };
        self.statistic.played += 1;
        self.statistic
            .won
            .entry(*board.you_id())
            .and_modify(|won| *won += usize::from(winner))
            .or_insert_with(|| usize::from(winner));
        winner
    }
}
