Ozai/src/azul.rs

use std::collections::HashMap;
use std::convert::TryFrom;
use std::ops::{Add, AddAssign};

use serde::{Deserialize, Serialize};

use rand::distributions::WeightedIndex;
use rand::prelude::*;

#[derive(Serialize, Deserialize)]
pub struct GameState {
    n_players: usize,
    current_player: PlayerName,
    starting_player: PlayerName,
    player_names: Vec<PlayerName>,
    game_end: bool,
    rounds: usize,
    days: usize,
    bag: TileSet,
    lid: TileSet,
    factories: Vec<TileSet>,
    market: TileSetWithStart,
    players: HashMap<PlayerName, Player>,
    #[serde(skip)]
    #[serde(default = "make_rng")]
    rng: Box<dyn RngCore + Send>,
}

#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
pub struct PlayerName(String);

impl From<String> for PlayerName {
    fn from(value: String) -> Self {
        Self(value)
    }
}

impl From<&str> for PlayerName {
    fn from(value: &str) -> Self {
        value.to_string().into()
    }
}

fn make_rng() -> Box<dyn RngCore + Send> {
    Box::new(StdRng::from_entropy())
}

impl Clone for GameState {
    fn clone(&self) -> Self {
        GameState {
            n_players: self.n_players,
            current_player: self.current_player.clone(),
            starting_player: self.starting_player.clone(),
            player_names: self.player_names.clone(),
            game_end: self.game_end,
            rounds: self.rounds,
            days: self.days,
            bag: self.bag.clone(),
            lid: self.lid.clone(),
            factories: self.factories.clone(),
            market: self.market.clone(),
            players: self.players.clone(),
            rng: make_rng(),
        }
    }
}

impl GameState {
    pub fn new(mut player_names: Vec<PlayerName>) -> Result<GameState, &'static str> {
        let n_players = player_names.len();
        if n_players < 2 {
            return Err("Can't create game with less than two players");
        }
        if n_players > 4 {
            return Err("Can't create game with more than 4 players");
        }
        if n_players != 2 {
            return Err("Only two players is currently supported");
        }

        let n_factories = 1 + 2 * n_players;

        let mut factories: Vec<TileSet> = Vec::with_capacity(n_factories);

        // let mut factories
        for _ in 0..n_factories {
            factories.push(TileSet::default());
        }

        let mut players = HashMap::<PlayerName, Player>::new();

        for player in &player_names {
            players.insert(player.clone(), Player::default());
        }

        let mut rng = Box::new(StdRng::from_entropy());
        player_names.shuffle(&mut rng);
        let starting_player = player_names[0].to_owned();

        let game = GameState {
            n_players,
            current_player: starting_player.clone(),
            starting_player,
            player_names,
            game_end: false,
            rounds: 0,
            days: 0,
            bag: TileSet {
                blue: 20,
                yellow: 20,
                red: 20,
                black: 20,
                white: 20,
            },
            lid: TileSet::default(),
            factories,
            market: TileSetWithStart {
                start: 1,
                blue: 0,
                yellow: 0,
                red: 0,
                black: 0,
                white: 0,
            },
            players,
            rng,
        };

        Ok(game)
    }

    pub fn set_ready(&mut self, player_name: &PlayerName) -> Result<(), &str> {
        self.players
            .get_mut(player_name)
            .ok_or("That player is not part of this game")?
            .ready = true;
        Ok(())
    }

    pub fn all_ready(&self) -> bool {
        self.players.values().all(|x| x.ready)
    }

    /// Fills the factories from the bag
    /// Will replenish the bag from the lid when empty
    /// Will return with partially filled factories if out of tiles
    pub fn fill(&mut self) -> Result<(), &'static str> {
        if !self.only_start() {
            return Err("Cannot fill, there are still tiles left to be picked");
        }

        for factory in &mut self.factories {
            for _ in 0..4 {
                if self.bag.is_empty() && !self.lid.is_empty() {
                    self.bag = self.lid.clone();
                    self.lid = TileSet::default();
                } else if self.bag.is_empty() {
                    return Ok(());
                }
                let choices = [
                    Color::Blue,
                    Color::Yellow,
                    Color::Red,
                    Color::Black,
                    Color::White,
                ];
                let weights = [
                    self.bag.blue,
                    self.bag.yellow,
                    self.bag.red,
                    self.bag.black,
                    self.bag.white,
                ];

                let dist = WeightedIndex::new(&weights).unwrap();
                let picked = choices[dist.sample(&mut self.rng)];

                self.bag.add_color(picked, -1)?;
                factory.add_color(picked, 1)?;
            }
        }

        Ok(())
    }

    /// Returns whether or not the market and factories are empty
    fn is_empty(&self) -> bool {
        let factories = self.factories.iter().all(|f| f.is_empty());
        let market = self.market.is_empty();

        factories && market
    }

    /// Returns whether or not only the start tile is in play
    fn only_start(&self) -> bool {
        let factories = self.factories.iter().all(|f| f.is_empty());
        let market = Color::Blue
            .into_iter()
            .all(|c| self.market.get_color(c) == 0);

        factories && market
    }

    /// Scores the game and clears the boards into the lid
    /// Doesn't check if game is ready to be scored!
    fn score_unchecked(&mut self) -> Result<(), &'static str> {
        for (player_name, player) in self.players.iter_mut() {
            for row in 0..5 {
                if player.pattern_lines[row].len() == (row + 1) {
                    let color = player.pattern_lines[row]
                        .color
                        .ok_or("patternline filled with None color")?;
                    let index = Player::get_wall_index(row, color)?;
                    player.wall[row][index] = true;
                    player.points += player.connected(row, index);

                    self.lid.add_color(color, row as isize)?;
                    player.pattern_lines[row] = PatternLine::default()
                }
            }

            let negative = match player.floor.len() {
                0 => 0,
                1 => 1,
                2 => 2,
                3 => 4,
                4 => 6,
                5 => 8,
                6 => 11,
                _ => 14,
            };
            player.points -= negative;

            if player.floor.start == 1 {
                self.starting_player = player_name.clone();
                player.floor.start = 0;
            }

            self.lid += player.floor.try_into()?;
            player.floor = TileSetWithStart::default()
        }

        Ok(())
    }

    /// Manages the entire lifecycle of the game
    /// After a move, will check if game should be scored,
    /// After scoring, either ends the game and calculates bonus
    /// or fills the factories again and sets the starting player
    /// Also tracks rounds and days
    pub fn do_iter(&mut self, game_move: GameMove) -> Result<(), MoveErr> {
        if self.game_end {
            return Err(MoveErr::Other("Game is over"));
        }

        self.do_move(game_move).map_err(|e| e)?;

        if self.is_empty() {
            self.score_unchecked().map_err(|e| MoveErr::Other(e))?;
            self.game_end = self
                .players
                .values()
                .any(|p| p.wall.iter().any(|r| r.iter().all(|&v| v)));

            if self.game_end {
                for player in self.players.values_mut() {
                    player.points += player.bonus_score().map_err(|e| MoveErr::Other(e))?;
                }
            } else {
                self.fill().map_err(|e| MoveErr::Other(e))?;
                self.current_player = self.starting_player.clone();
                self.days += 1;
            }
        }

        self.rounds += 1;

        Ok(())
    }

    /// Does a move according to the policy defined in the move
    fn do_move(&mut self, game_move: GameMove) -> Result<(), MoveErr> {
        match game_move.policy {
            Policy::Strict => self.do_move_strict(game_move),
            Policy::Loose => self.do_move_loose(game_move),
        }
    }

    /// Does a move, if the move is placed where it can't fit, places it on the floor
    fn do_move_loose(&mut self, game_move: GameMove) -> Result<(), MoveErr> {
        match self.do_move_strict(game_move.clone()) {
            Ok(()) => Ok(()),
            Err(MoveErr::Dst(_)) => {
                let mut new_game_move = game_move;
                new_game_move.destination = Destination::Floor;
                self.do_move_strict(new_game_move)
            }
            e => e,
        }
    }

    /// Does a move, errors out if the move is weird in any way
    fn do_move_strict(&mut self, game_move: GameMove) -> Result<(), MoveErr> {
        if self.current_player != game_move.player {
            return Err(MoveErr::Player("Not this player's turn"));
        }

        // let player = &mut self.players[current_player];

        let color = game_move.color;
        let src = game_move.source;
        let dst = game_move.destination;

        match (color, src, dst) {
            (Color::Start, _, _) => {
                return Err(MoveErr::Color("You can't take the start tile specifically"))
            }
            (c, s, _) if self.get_color(s, c) == 0 => {
                return Err(MoveErr::Color("Source does not contain that color"))
            }
            (_, Source::Factory(f), _) if f > self.factories.len() => {
                return Err(MoveErr::Src("Not a valid factory"))
            }
            (_, _, Destination::PatternLine(l))
                if self.players[&self.current_player].pattern_lines[l].number > l + 1 =>
            {
                return Err(MoveErr::Dst("That PatternLine is full!"))
            }
            (c, _, Destination::PatternLine(l))
                if self.players[&self.current_player].pattern_lines[l].color != Some(c)
                    && !self.players[&self.current_player].pattern_lines[l]
                        .color
                        .is_none() =>
            {
                return Err(MoveErr::Dst(
                    "That pattern line already contains tiles of a different color!",
                ))
            }
            (c, _, Destination::PatternLine(l))
                if self.players[&self.current_player].wall[l]
                    [Player::get_wall_index(l, c).unwrap()]
                    == true =>
            {
                return Err(MoveErr::Dst(
                    "That pattern line and color is already filled on the wall",
                ))
            }
            (c, s, Destination::PatternLine(p)) => {
                let amount = self.take_tiles(s, c);

                let pattern_line = &mut self
                    .players
                    .get_mut(&self.current_player)
                    .unwrap()
                    .pattern_lines[p];

                let remaining_capacity = p + 1 - pattern_line.len();
                let to_line = usize::min(amount, remaining_capacity);
                let to_floor = amount - to_line;

                pattern_line.color = Some(c);
                pattern_line.number += to_line;
                self.players
                    .get_mut(&self.current_player)
                    .unwrap()
                    .floor
                    .add_color(c, to_floor as isize)
                    .map_err(|e| MoveErr::Other(e))?;
            }
            (c, s, Destination::Floor) => {
                let amount = self.take_tiles(s, c);
                self.players
                    .get_mut(&self.current_player)
                    .unwrap()
                    .floor
                    .add_color(c, amount as isize)
                    .map_err(|e| MoveErr::Other(e))?;
            }
        };

        self.current_player = self
            .player_names
            .iter()
            .cycle()
            .skip_while(|a| **a != self.current_player)
            .skip(1)
            .next()
            .unwrap()
            .to_owned();

        Ok(())
    }

    /// Gets how many tiles of a given color is in a source
    fn get_color(&self, src: Source, color: Color) -> usize {
        match src {
            Source::Market => self.market.get_color(color),
            Source::Factory(f) => self.factories[f].get_color(color),
        }
    }

    /// Takes tiles from a source and causes side-effects
    fn take_tiles(&mut self, src: Source, color: Color) -> usize {
        let amount = self.get_color(src, color);
        let player = &mut self.players.get_mut(&self.current_player).unwrap();

        // Sideffects
        match src {
            Source::Market => {
                if self.market.start == 1 {
                    self.market.start -= 1;
                    player.floor.start += 1;
                }
                let _ = self.market.add_color(color, -(amount as isize));
            }
            Source::Factory(f) => {
                let factory = &mut self.factories[f];
                let _ = factory.add_color(color, -(amount as isize));
                self.market += *factory;
                factory.clear();
            }
        };

        return amount;
    }
}

pub enum MoveErr {
    Player(&'static str),
    Color(&'static str),
    Src(&'static str),
    Dst(&'static str),
    Other(&'static str),
}

#[derive(Debug, Serialize, Deserialize, Clone, Default, Copy)]
struct TileSet {
    blue: usize,
    yellow: usize,
    red: usize,
    black: usize,
    white: usize,
}

impl TileSet {
    fn len(&self) -> usize {
        self.blue + self.yellow + self.red + self.black + self.white
    }
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn add_color(&mut self, color: Color, n: isize) -> Result<(), &'static str> {
        match color {
            Color::Blue => {
                self.blue = usize::checked_add_signed(self.blue, n).ok_or("would overflow!")?
            }
            Color::Yellow => {
                self.yellow = usize::checked_add_signed(self.yellow, n).ok_or("would overflow")?
            }
            Color::Red => {
                self.red = usize::checked_add_signed(self.red, n).ok_or("would overflow")?
            }
            Color::Black => {
                self.black = usize::checked_add_signed(self.black, n).ok_or("would overflow")?
            }
            Color::White => {
                self.white = usize::checked_add_signed(self.white, n).ok_or("would overflow")?
            }
            Color::Start => return Err("tried to add Start tiles to TileSet"),
        }

        Ok(())
    }

    fn get_color(&self, color: Color) -> usize {
        match color {
            Color::Start => 0,
            Color::Blue => self.blue,
            Color::Yellow => self.yellow,
            Color::Red => self.red,
            Color::Black => self.black,
            Color::White => self.white,
        }
    }

    fn clear(&mut self) {
        *self = Self::default();
    }
}

impl Add for TileSet {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        Self {
            blue: self.blue + other.blue,
            yellow: self.yellow + other.yellow,
            red: self.red + other.red,
            black: self.black + other.black,
            white: self.white + other.white,
        }
    }
}

impl AddAssign for TileSet {
    fn add_assign(&mut self, other: Self) {
        *self = *self + other
    }
}

impl TryFrom<TileSetWithStart> for TileSet {
    type Error = &'static str;

    fn try_from(value: TileSetWithStart) -> Result<Self, Self::Error> {
        if value.start == 0 {
            Ok(TileSet {
                blue: value.blue,
                yellow: value.yellow,
                red: value.red,
                black: value.black,
                white: value.white,
            })
        } else {
            Err("Can't convert, tileset had a start tile inside")
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, Default, Copy)]
struct TileSetWithStart {
    start: usize,
    blue: usize,
    yellow: usize,
    red: usize,
    black: usize,
    white: usize,
}

impl TileSetWithStart {
    fn len(&self) -> usize {
        self.blue + self.yellow + self.red + self.black + self.white + self.start
    }
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
    fn get_color(&self, color: Color) -> usize {
        match color {
            Color::Start => self.start,
            Color::Blue => self.blue,
            Color::Yellow => self.yellow,
            Color::Red => self.red,
            Color::Black => self.black,
            Color::White => self.white,
        }
    }

    fn add_color(&mut self, color: Color, n: isize) -> Result<(), &'static str> {
        match color {
            Color::Start if n == 1 && self.start == 0 => {
                self.start = usize::checked_add_signed(self.white, n).ok_or("would overflow")?
            }
            Color::Start => return Err("Tried to add more than 1 start tile to TileSet"),
            Color::Blue => {
                self.blue = usize::checked_add_signed(self.blue, n).ok_or("would overflow!")?
            }
            Color::Yellow => {
                self.yellow = usize::checked_add_signed(self.yellow, n).ok_or("would overflow")?
            }
            Color::Red => {
                self.red = usize::checked_add_signed(self.red, n).ok_or("would overflow")?
            }
            Color::Black => {
                self.black = usize::checked_add_signed(self.black, n).ok_or("would overflow")?
            }
            Color::White => {
                self.white = usize::checked_add_signed(self.white, n).ok_or("would overflow")?
            }
        }

        Ok(())
    }
}

impl Add for TileSetWithStart {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        Self {
            start: self.start + other.start,
            blue: self.blue + other.blue,
            yellow: self.yellow + other.yellow,
            red: self.red + other.red,
            black: self.black + other.black,
            white: self.white + other.white,
        }
    }
}

impl Add<TileSet> for TileSetWithStart {
    type Output = Self;

    fn add(self, other: TileSet) -> Self {
        Self {
            start: self.start,
            blue: self.blue + other.blue,
            yellow: self.yellow + other.yellow,
            red: self.red + other.red,
            black: self.black + other.black,
            white: self.white + other.white,
        }
    }
}

impl AddAssign<TileSet> for TileSetWithStart {
    fn add_assign(&mut self, other: TileSet) {
        *self = *self + other
    }
}

impl From<TileSet> for TileSetWithStart {
    fn from(value: TileSet) -> Self {
        TileSetWithStart {
            start: 0,
            blue: value.blue,
            yellow: value.yellow,
            red: value.red,
            black: value.black,
            white: value.white,
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, Default)]
struct Player {
    ready: bool,
    points: usize,
    pattern_lines: [PatternLine; 5],
    wall: Wall,
    floor: TileSetWithStart,
}

impl Player {
    fn get_wall_index(row: usize, color: Color) -> Result<usize, &'static str> {
        match (row, color) {
            (0, Color::Blue) => Ok(0),
            (0, Color::Yellow) => Ok(1),
            (0, Color::Red) => Ok(2),
            (0, Color::Black) => Ok(3),
            (0, Color::White) => Ok(4),

            (1, Color::Blue) => Ok(1),
            (1, Color::Yellow) => Ok(2),
            (1, Color::Red) => Ok(3),
            (1, Color::Black) => Ok(4),
            (1, Color::White) => Ok(0),

            (2, Color::Blue) => Ok(2),
            (2, Color::Yellow) => Ok(3),
            (2, Color::Red) => Ok(4),
            (2, Color::Black) => Ok(0),
            (2, Color::White) => Ok(1),

            (3, Color::Blue) => Ok(3),
            (3, Color::Yellow) => Ok(4),
            (3, Color::Red) => Ok(0),
            (3, Color::Black) => Ok(1),
            (3, Color::White) => Ok(2),

            (4, Color::Blue) => Ok(4),
            (4, Color::Yellow) => Ok(0),
            (4, Color::Red) => Ok(1),
            (4, Color::Black) => Ok(2),
            (4, Color::White) => Ok(3),

            _ => return Err("Not a valid tile on the wall"),
        }
    }

    /// given a coordinate on the board, returns how many tiles are in its group
    fn connected(&self, row: usize, column: usize) -> usize {
        // This implementation scans for contigious runs of tiles
        // until it finds the group the tile we're checking is part of
        // Does this for both the row and the column
        // then returns the sum

        let wall = self.wall;

        if wall[row][column] == false {
            return 0;
        }

        let mut sum = 0;

        // Count connected tiles in the on the row
        let mut count = 0;
        let mut active = false;
        for i in 0..5 {
            if active == true && wall[row][i] == false {
                break;
            } else if wall[row][i] == false {
                count = 0;
            } else if (row, i) == (row, column) {
                active = true;
                count += 1;
            } else {
                count += 1
            }
        }
        sum += count;

        // Count connected tiles in the column
        let mut count = 0;
        let mut active = false;
        for i in 0..5 {
            if active == true && wall[i][column] == false {
                break;
            } else if wall[i][column] == false {
                count = 0;
            } else if (i, column) == (row, column) {
                active = true;
                count += 1;
            } else {
                count += 1
            }
        }
        sum += count;

        return sum;
    }

    /// Calculates end-of-game bonus
    fn bonus_score(&self) -> Result<usize, &'static str> {
        let mut bonus = 0;

        // Horizontal
        for row in self.wall {
            if row.iter().all(|&x| x == true) {
                bonus += 2;
            };
        }

        // Vertical
        for column in 0..5 {
            let mut tiles = 0;
            for row in 0..5 {
                if self.wall[row][column] == true {
                    tiles += 1;
                }
            }
            if tiles == 5 {
                bonus += 7
            }
        }

        // 5 of a color
        for color in Color::Blue.into_iter() {
            let mut tiles = 0;
            for row in 0..5 {
                let index = Player::get_wall_index(row, color)?;
                if self.wall[row][index] == true {
                    tiles += 1
                }
            }

            if tiles == 5 {
                bonus += 10;
            }
        }

        return Ok(bonus);
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, Default)]
struct PatternLine {
    color: Option<Color>,
    number: usize,
}

impl PatternLine {
    fn len(&self) -> usize {
        self.number
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)]
enum Color {
    #[serde(rename = "start")]
    Start,
    #[serde(rename = "blue")]
    Blue,
    #[serde(rename = "yellow")]
    Yellow,
    #[serde(rename = "red")]
    Red,
    #[serde(rename = "black")]
    Black,
    #[serde(rename = "white")]
    White,
}

impl IntoIterator for Color {
    type Item = Color;
    type IntoIter = ColorIter;

    /// iterates through blue, yellow, red, black, and white
    fn into_iter(self) -> Self::IntoIter {
        ColorIter { current: self }
    }
}

struct ColorIter {
    current: Color,
}

impl Iterator for ColorIter {
    type Item = Color;

    fn next(&mut self) -> Option<Color> {
        match self.current {
            Color::Blue => {
                let next = Color::Yellow;
                self.current = next;
                Some(next)
            }
            Color::Yellow => {
                let next = Color::Red;
                self.current = next;
                Some(next)
            }
            Color::Red => {
                let next = Color::Black;
                self.current = next;
                Some(next)
            }
            Color::Black => {
                let next = Color::White;
                self.current = next;
                Some(next)
            }
            _ => None,
        }
    }
}

type Row = [bool; 5];
type Wall = [Row; 5];

#[derive(Serialize, Deserialize, Clone)]
pub struct GameMove {
    player: PlayerName, // Safeguard to check that the bot knows which player it is
    policy: Policy,     // What policy to run moves under, so bots can do less error-handling
    color: Color,       // What color to select
    source: Source,     // What source to move the tiles from
    destination: Destination, // Where to place the tiles
}

#[derive(Serialize, Deserialize, Clone, Copy)]
enum Policy {
    /// Anything weird will return an error
    #[serde(rename = "strict")]
    Strict,
    #[serde(rename = "loose")]
    /// If you place tiles wrongly, they will fall on the floor, anything else will return an error
    Loose,
    // Anything weird will instead play a random move
    // #[serde(rename = "random")]
    // Random,
}

#[derive(Serialize, Deserialize, Clone, Copy)]
enum Source {
    #[serde(rename = "market")]
    Market,
    #[serde(untagged)]
    Factory(usize),
}

#[derive(Serialize, Deserialize, Clone, Copy)]
enum Destination {
    #[serde(rename = "floor")]
    Floor,
    #[serde(untagged)]
    PatternLine(usize),
}

// Tests

use serde_json::to_string_pretty;

#[test]
fn test_fill() {
    let mut game = GameState::new(vec!["Bot1".into(), "Bot2".into()]).unwrap();

    // println!("{}", to_string_pretty(&game).unwrap());
    assert!(game.only_start());

    game.fill();
    //println!("{}", to_string_pretty(&game).unwrap());

    assert_eq!(game.bag.len(), 100 - (5 * 4));
    assert_eq!(game.factories.iter().map(|f| f.len()).sum::<usize>(), 5 * 4);
}