#ifndef DYNAMIC_PROGRAM_GAME_STATE_H
#define DYNAMIC_PROGRAM_GAME_STATE_H

#include <array>
#include <cstdint>
#include <algorithm>
#include <cstddef>
#include <unordered_map>
#include <bitset>
#include <limits>
#include <optional>
#include <boost/container/static_vector.hpp>
#include <list>
#include <ostream>


namespace Hanabi {

    using rank_t = std::uint8_t;
    using suit_t = std::uint8_t;
    using clue_t = std::uint8_t;
    using player_t = std::uint8_t;
    using hand_index_t = std::uint8_t;

    using state_t = std::uint32_t;

/**
 * We will generally assume that stacks are played from n to 0
 * Playing a 0 will yield a clue
 * Therefore, for the default hanabi, we will play 4,3,2,1,0 in that order
 * on each stack. A stack with no cards played implicitly has value 5 on it
 * This is just easier to implement, since then the remaining number of cards
 * to be played is always the current number of the stack
 */
    constexpr rank_t starting_card_rank = 5;
    constexpr suit_t max_suit_index = 5;
    constexpr size_t max_card_duplicity = 3;
    constexpr clue_t max_num_clues = 8;
    constexpr uint8_t not_in_starting_hand = std::numeric_limits<uint8_t>::max();

    constexpr std::array<std::string, 6> suit_initials{"r", "y", "g", "b", "p", "t"};

    struct Card {
        suit_t suit;
        rank_t rank;
        uint8_t local_index;
        bool in_starting_hand;
        bool initial_trash;

        Card &operator++();

        const Card operator++(int);

        auto operator<=>(const Card &) const = default;
    };
}

namespace std {
    template<>
    struct hash<Hanabi::Card> {
        std::size_t operator()(Hanabi::Card const& card) const noexcept {
            return card.suit * 6 + card.rank;
        }
    };
}

namespace Hanabi {

std::ostream &operator<<(std::ostream &os, const Card &card);

constexpr Card r0 = {0, 0};
constexpr Card r1 = {0, 1};
constexpr Card r2 = {0, 2};
constexpr Card r3 = {0, 3};
constexpr Card r4 = {0, 4};
constexpr Card y0 = {1, 0};
constexpr Card y1 = {1, 1};
constexpr Card y2 = {1, 2};
constexpr Card y3 = {1, 3};
constexpr Card y4 = {1, 4};

/**
 * To store:
 * - Draw pile size
 * - Distribution of cards
 *      - Which cards exist?
 * - Number of clues
 */

template <size_t num_suits>
using Stacks = std::array<rank_t, num_suits>;

template <size_t num_suits>
std::ostream& operator<<(std::ostream &os, const Stacks<num_suits> &stacks);

struct CardMultiplicity {
  Card card;
  std::uint8_t multiplicity;

  auto operator<=>(const CardMultiplicity &) const = default;
};

template<typename T>
struct InnerCardArray {
    template<size_t N>
    using array_t = std::array<T, N>;
};

template<>
struct InnerCardArray<bool> {
    template<size_t N>
    using array_t = std::bitset<N>;
};

template <suit_t num_suits, typename T> struct CardArray {
    using value_type = T;

    CardArray() = default;
    explicit CardArray(value_type default_val);

    void fill(value_type val);

    const value_type &operator[](const Card &card) const;

    value_type &operator[](const Card &card);

    auto operator<=>(const CardArray &) const = default;

private:
    using inner_array_t = typename InnerCardArray<T>::template array_t<starting_card_rank>;
    std::array<inner_array_t , num_suits> _array {};
};

enum class ActionType {
    play = 0,
    discard = 1,
    clue = 2,
    color_clue = 2,
    rank_clue = 3,
    end_game = 4,
    vote_terminate = 10,
};

struct BacktrackAction {
  // The card that was discarded or played
  Card discarded{};
  // Index of card in hand that was discarded or played
  hand_index_t index{};
  // Multiplicity of new draw (needed for probability calculations)
  hand_index_t multiplicity{};
};

class HanabiStateIF {
public:
    virtual double backtrack(size_t depth) = 0;

    virtual void clue() = 0;
    virtual BacktrackAction discard(hand_index_t index) = 0;
    virtual BacktrackAction play(hand_index_t index) = 0;

    [[nodiscard]] virtual hand_index_t find_card_in_hand(const Card& card) const = 0;
    [[nodiscard]] virtual bool is_trash(const Card& card) const = 0;
    [[nodiscard]] virtual bool is_playable(const Card& card) const = 0;
    [[nodiscard]] virtual size_t draw_pile_size() const = 0;

    [[nodiscard]] virtual std::uint64_t enumerated_states() const = 0;
    [[nodiscard]] virtual std::unordered_map<unsigned long, double> visited_states() const = 0;

    virtual void normalize_draw_and_positions() = 0;

    virtual ~HanabiStateIF() = default;

protected:
    virtual void print(std::ostream& os) const = 0;

    friend std::ostream& operator<<(std::ostream&, HanabiStateIF const&);
};

template <suit_t num_suits, player_t num_players, hand_index_t hand_size>
class HanabiState : public HanabiStateIF {
public:
    HanabiState() = default;
    explicit HanabiState(const std::vector<Card>& deck);

    double backtrack(size_t depth) final;

    void clue() final;
    BacktrackAction play(hand_index_t index) final;
    BacktrackAction discard(hand_index_t index) final;

    void revert_clue();
    void revert_play(const BacktrackAction &action, bool was_on_8_clues);
    void revert_discard(const BacktrackAction &action);

    [[nodiscard]] hand_index_t find_card_in_hand(const Card& card) const final;
    [[nodiscard]] bool is_trash(const Card& card) const final;
    [[nodiscard]] bool is_playable(const Card& card) const final;
    [[nodiscard]] size_t draw_pile_size() const final;

    [[nodiscard]] std::uint64_t enumerated_states() const final;
    [[nodiscard]] std::unordered_map<unsigned long, double> visited_states() const final;

    void normalize_draw_and_positions() final;

    auto operator<=>(const HanabiState &) const = default;

protected:
    void print(std::ostream& os) const final;

private:
    template<bool update_card_positions> BacktrackAction play_and_potentially_update(hand_index_t index);
    template<bool update_card_positions> BacktrackAction discard_and_potentially_update(hand_index_t index);

    template<bool update_card_positions> hand_index_t draw(hand_index_t index);
    void revert_draw(hand_index_t index, Card discarded_card);

    void incr_turn();
    void decr_turn();

    static constexpr uint8_t no_endgame = std::numeric_limits<uint8_t>::max();
    static constexpr player_t draw_pile = num_players;
    static constexpr player_t trash_or_play_stack = num_players + 1;

    std::uint64_t unique_id() const;

    player_t _turn{};
    clue_t _num_clues{};
    std::uint8_t _weighted_draw_pile_size{};
    Stacks<num_suits> _stacks{};
    std::array<std::array<Card, hand_size>, num_players> _hands{};
    std::list<CardMultiplicity> _draw_pile{};
    std::uint8_t _endgame_turns_left{};

    // This will save the card positions of all cards that are in the draw pile when we start backtracking
    // TODO: currently, we support at most 10 useful different cards in draw pile
    boost::container::static_vector<boost::container::static_vector<player_t, max_card_duplicity>, 10> _card_positions_draw;

    // This will indicate whether cards that were in hands initially still are in hands
    std::bitset<num_players * hand_size> _card_positions_hands;

    size_t _num_useful_cards_in_starting_hands;
    size_t _initial_draw_pile_size;

    // further statistics that we might want to keep track of
    uint8_t _pace{};
    uint8_t _score{};

    std::uint64_t _enumerated_states {};

    std::unordered_map<unsigned long, double> _position_tablebase;
};

template <std::size_t num_suits, player_t num_players, std::size_t hand_size>
bool same_up_to_discard_permutation(HanabiState<num_suits, num_players, hand_size> state1, HanabiState<num_suits, num_players, hand_size> state2) {
    auto comp = [](CardMultiplicity &m1, CardMultiplicity &m2) -> bool {
        return m1.card.suit < m2.card.suit || (m1.card.suit == m2.card.suit and m1.card.rank < m2.card.rank) ||
               (m1.card.suit == m2.card.suit and m1.card.rank == m2.card.rank and m1.multiplicity < m2.multiplicity);
    };
   state1._draw_pile.sort(comp);
   state2._draw_pile.sort(comp);
   return state1 == state2;
}


}

#include "game_state.hpp"

#endif // DYNAMIC_PROGRAM_GAME_STATE_H