// Copyright 2019 DeepMind Technologies Ltd. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <iostream>
#include <filesystem>
#include <fstream>

#include "open_spiel/algorithms/kriegspiel/stats_collector.h"

#include "open_spiel/games/kriegspiel.h"
#include "open_spiel/games/chess/pgn.h"
#include "open_spiel/games/chess/pgn_parser.h"
#include "open_spiel/abseil-cpp/absl/random/distributions.h"

namespace open_spiel {
namespace algorithms {


void UpdateStatistics(const chess::ChessBoard &board, int t,
                      std::array<std::array<uint, kStatsLength>, 2> &counts,
                      StatisticsTable &king_counts,
                      StatisticsTable &piece_counts,
                      StatisticsTable &pawn_counts,
                      StatisticsTable &rook_counts,
                      StatisticsTable &knight_counts,
                      StatisticsTable &bishop_counts,
                      StatisticsTable &queen_counts) {

  for (int player = 0; player < 2; ++player) {
    counts[player][t]++;
    for (int8_t x = 0; x < chess::kMaxBoardSize; ++x) {
      for (int8_t y = 0; y < chess::kMaxBoardSize; ++y) {
        chess::Piece piece = board.at({x, y});

        if (piece.color == chess::Color::kEmpty || chess::ColorToPlayer(piece.color) != player) {
          continue;
        }

        switch(piece.type) {
          case chess::PieceType::kKing:
            king_counts[player][t][x][y]++;
            break;
          case chess::PieceType::kPawn:
            pawn_counts[player][t][x][y]++;
            break;
          case chess::PieceType::kRook:
            rook_counts[player][t][x][y]++;
            piece_counts[player][t][x][y]++;
            break;
          case chess::PieceType::kKnight:
            knight_counts[player][t][x][y]++;
            piece_counts[player][t][x][y]++;
            break;
          case chess::PieceType::kBishop:
            bishop_counts[player][t][x][y]++;
            piece_counts[player][t][x][y]++;
            break;
          case chess::PieceType::kQueen:
            queen_counts[player][t][x][y]++;
            piece_counts[player][t][x][y]++;
            break;
        }
      }
    }
  }

}

void ComputeStatistics() {

  std::ifstream is = std::ifstream ("/hdd/diplomka/KS_databases/kriegbase.unix.pgn", std::ios::in);

  std::string pgndata;
  std::copy(std::istreambuf_iterator<char>(is), std::istreambuf_iterator<char>(), std::inserter(pgndata, pgndata.end()));

  std::string::const_iterator itr1 = pgndata.begin();
  std::string::const_iterator itr2 = pgndata.end();

  std::vector<pgn::ChessPGN> games;
  pgn::Parser::GetGames(itr1, itr2, games);

  std::array<std::array<uint, kStatsLength>, 2> counts{};
  StatisticsTable king_counts{};
  StatisticsTable piece_counts{};
  StatisticsTable pawn_counts{};
  StatisticsTable rook_counts{};
  StatisticsTable knight_counts{};
  StatisticsTable bishop_counts{};
  StatisticsTable queen_counts{};

  for (const pgn::ChessPGN& chess_pgn : games) {
    chess::ChessState chess_state = chess_pgn.state_;

    chess::ChessBoard board = chess_state.StartBoard();

    int t = 0;
    UpdateStatistics(board, t, counts, king_counts, piece_counts, pawn_counts, rook_counts, knight_counts, bishop_counts, queen_counts);

    for (const chess::Move move : chess_state.MovesHistory()) {

      board.ApplyMove(move);
      t++;

      UpdateStatistics(board, t, counts, king_counts, piece_counts, pawn_counts, rook_counts, knight_counts, bishop_counts, queen_counts);
    }

  }
  for (int player = 0; player < 2; ++player) {
    std::ofstream ofstream_counts("/hdd/stats_counts_" + std::to_string(player));
    std::ofstream ofstream_king("/hdd/stats_king_" + std::to_string(player));
    std::ofstream ofstream_piece("/hdd/stats_piece_" + std::to_string(player));
    std::ofstream ofstream_pawn("/hdd/stats_pawn_" + std::to_string(player));
    std::ofstream ofstream_rook("/hdd/stats_rook_" + std::to_string(player));
    std::ofstream ofstream_knight("/hdd/stats_knight_" + std::to_string(player));
    std::ofstream ofstream_bishop("/hdd/stats_bishop_" + std::to_string(player));
    std::ofstream ofstream_queen("/hdd/stats_queen_" + std::to_string(player));
    for (int i = 0; i < kStatsLength; ++i) {
      ofstream_counts << counts[player][i] << std::endl;
      for (int8_t x = 0; x < 8; ++x) {
        for (int8_t y = 0; y < 8; ++y) {
          ofstream_king << king_counts[player][i][x][y] << " ";
          ofstream_piece << piece_counts[player][i][x][y] << " ";
          ofstream_pawn << pawn_counts[player][i][x][y] << " ";
          ofstream_rook << rook_counts[player][i][x][y] << " ";
          ofstream_knight << knight_counts[player][i][x][y] << " ";
          ofstream_bishop << bishop_counts[player][i][x][y] << " ";
          ofstream_queen << queen_counts[player][i][x][y] << " ";
        }
        ofstream_king << std::endl;
        ofstream_piece << std::endl;
        ofstream_pawn << std::endl;
        ofstream_rook << std::endl;
        ofstream_knight << std::endl;
        ofstream_bishop << std::endl;
        ofstream_queen << std::endl;
      }

      ofstream_king << std::endl;
      ofstream_piece << std::endl;
      ofstream_pawn << std::endl;
      ofstream_rook << std::endl;
      ofstream_knight << std::endl;
      ofstream_bishop << std::endl;
      ofstream_queen << std::endl;
    }

  }
}

Stats ReadStatistics() {

  std::array<std::array<uint, kStatsLength>, 2> counts{};
  StatisticsTable king_counts{};
  StatisticsTable piece_counts{};
  StatisticsTable pawn_counts{};
  StatisticsTable rook_counts{};
  StatisticsTable knight_counts{};
  StatisticsTable bishop_counts{};
  StatisticsTable queen_counts{};

  for (int player = 0; player < 2; ++player) {
    std::ifstream ifstream_counts("/hdd/stats_counts_" + std::to_string(player));
    std::ifstream ifstream_king("/hdd/stats_king_" + std::to_string(player));
    std::ifstream ifstream_piece("/hdd/stats_piece_" + std::to_string(player));
    std::ifstream ifstream_pawn("/hdd/stats_pawn_" + std::to_string(player));
    std::ifstream ifstream_rook("/hdd/stats_rook_" + std::to_string(player));
    std::ifstream ifstream_knight("/hdd/stats_knight_" + std::to_string(player));
    std::ifstream ifstream_bishop("/hdd/stats_bishop_" + std::to_string(player));
    std::ifstream ifstream_queen("/hdd/stats_queen_" + std::to_string(player));
    for (int i = 0; i < kStatsLength; ++i) {
      ifstream_counts >> counts[player][i];
      for (int8_t x = 0; x < 8; ++x) {
        for (int8_t y = 0; y < 8; ++y) {
          ifstream_king >> king_counts[player][i][x][y];
          ifstream_piece >> piece_counts[player][i][x][y];
          ifstream_pawn >> pawn_counts[player][i][x][y];
          ifstream_rook >> rook_counts[player][i][x][y];
          ifstream_knight >> knight_counts[player][i][x][y];
          ifstream_bishop >> bishop_counts[player][i][x][y];
          ifstream_queen >> queen_counts[player][i][x][y];
        }
      }
    }
  }

  return {counts, king_counts, piece_counts, pawn_counts, rook_counts, knight_counts, bishop_counts, queen_counts};
}

struct Result {
  int wins = 0;
  int draws = 0;
  int losses = 0;
};

void ReadResults() {
  std::string dir = "/hdd/diplomka/experiments/results/rr";
//  std::string dir = "/hdd/neco";

  std::vector<open_spiel::pgn::ChessPGN> games;

  int i = 0;
  for (const auto& entry : std::filesystem::directory_iterator(dir)) {
    std::cout << entry.path() << "          " << i++ << std::endl;
    std::ifstream ifs = std::ifstream(entry.path(), std::ios::in);

    std::string pgndata;
    std::copy(std::istreambuf_iterator<char>(ifs), std::istreambuf_iterator<char>(), std::inserter(pgndata, pgndata.end()));

    std::string::const_iterator itr1 = pgndata.begin();
    std::string::const_iterator itr2 = pgndata.end();

    open_spiel::pgn::Parser::GetGames(itr1, itr2, games);

  }
  size_t b = games.size();
  int a = 5;

  std::map<std::pair<std::string, std::string>, Result> map;

  for (const auto& pgnGame : games) {
    std::pair<std::string, std::string> players = std::make_pair(pgnGame.white_, pgnGame.black_);
    bool reversed = players.first == "Uniform Random";
    if (reversed) {
      players = std::make_pair(players.second, players.first);
    }

    if (!map.count(players)) {
      map[players] = {};
    }

    Result result = map[players];

    auto returns = pgnGame.state_.Returns();
    if (returns[0] == 0.0) {
      result.draws++;
    }
    else if (returns[0] > 0.0) {
      if (reversed) {
        result.wins++;
      }
      else {
        result.losses++;
      }
    }
    else {
      if (reversed) {
        result.losses++;
      }
      else {
        result.wins++;
      }
    }
    map[players] = result;
    int c = 5;

  }

  for (auto &[players, result] : map) {
    int count = result.wins + result.draws + result.losses;
    double e = (((double)result.wins) + result.draws / 2) / (count);

    double sd = sqrt((1.0/(count - 1.0)) * (result.wins * (pow(1.0 - e, 2.0)) + result.draws * (pow(0.5 - e, 2.0)) + result.losses * (pow(0.0 - e, 2.0))));
    double ci = 1.96 * sd / sqrt(count);
    std::cout << players.first << " " << players.second << ": " << result.wins << " - " << result.draws << " - " << result.losses << " " << e << " +- " << ci << " " << sd << std::endl;
  }


}

}  // namespace algorithms
}  // namespace open_spiel
