#include <lib94/lib94.hpp>
#include <filesystem>
#include <algorithm>
#include <iostream>
#include <cstdlib>
#include <fstream>
#include <sstream>
#include <mpi.h>

constexpr size_t max_good_warriors = 100;
constexpr int rounds_per_matchup = 100;
constexpr int steps_to_tie = 1000000;

constexpr lib94::instruction background_instruction = {
  .op = lib94::DAT, .mod = lib94::F,
  .amode = lib94::DIRECT, .bmode = lib94::DIRECT,
  .anumber = 0, .bnumber = 0
};

lib94::warrior *warrior_to_beat;
int comm_rank, comm_size;

std::filesystem::path output_dir;

std::vector<lib94::warrior *> seed_warriors;

void head_main();
void child_main();

int main(int argc, char **argv) {

  MPI_Init(&argc, &argv);

  MPI_Comm_rank(MPI_COMM_WORLD, &comm_rank);
  MPI_Comm_size(MPI_COMM_WORLD, &comm_size);

  if (comm_size < 2) {
    std::cerr << "you need at least two processes." << std::endl;
    return 1;
  }

  if (argc < 4) {
    if (comm_rank == 0)
      std::cerr << "usage: " << argv[0] << " <warrior to beat> <output directory> <seed 1> <seed 2> ..." << std::endl;
    return 2;
  }

  try {
    warrior_to_beat = lib94::compile_warrior_from_file(argv[1]);
  }
  catch (const lib94::compiler_exception &ex) {
    if (comm_rank == 0)
      std::cerr
        << "syntax error on line " << ex.source_line_number
        << " of " << argv[1] << ":\n  " << ex.message << std::endl;
    return 3;
  }

  if (warrior_to_beat == 0) {
    if (comm_rank == 0)
      std::cerr << "could not read file " << argv[1] << std::endl;
    return 4;
  }

  output_dir = argv[2];

  for (int i = 3; i < argc; ++i) {

    try {
      seed_warriors.push_back(lib94::compile_warrior_from_file(argv[i]));
    }
    catch (const lib94::compiler_exception &ex) {
      if (comm_rank == 0)
        std::cerr
          << "syntax error on line " << ex.source_line_number
          << " of " << argv[i] << ":\n  " << ex.message << std::endl;
      return 3;
    }

    if (seed_warriors.back() == 0) {
      if (comm_rank == 0)
        std::cerr << "could not read file " << argv[i] << std::endl;
      return 4;
    }

  }

  if (comm_rank == 0)
    head_main();

  else
    child_main();

  return 0;

}

struct scored_warrior {
  int score;
  lib94::warrior *w;
  bool operator <(const scored_warrior &other) const {
    return score > other.score || (score == other.score && w->instructions.size() < other.w->instructions.size());
  }
};

void head_main() {

  std::cout << "\x1b""7\x1b[?47h\x1b[2J" << std::flush;

  std::vector<scored_warrior> warriors;
  for (lib94::warrior *const &w : seed_warriors)
    warriors.push_back({.score = -1, .w = w});

  int *warrior_counts = new int[comm_size];

  int round = 1;

  while (true) {

    std::cout << "\x1b[1;1Hhead node: sending warriors\x1b[0K" << std::flush;

    int warrior_count = warriors.size();
    MPI_Bcast(&warrior_count, 1, MPI_INT, 0, MPI_COMM_WORLD);

    for (scored_warrior &sw : warriors) {
      std::vector<uint8_t> serialization;
      lib94::serialize_warrior(sw.w, serialization);
      int serialization_length = serialization.size();
      MPI_Bcast(&serialization_length, 1, MPI_INT, 0, MPI_COMM_WORLD);
      MPI_Bcast(serialization.data(), serialization_length, MPI_UINT8_T, 0, MPI_COMM_WORLD);
    }

    if (round == 1) {
      for (scored_warrior &sw : warriors)
        delete sw.w;
      warriors.clear();
    }

    std::cout << "\x1b[1;1Hhead node: waiting for child nodes\x1b[0K" << std::flush;

    MPI_Gather(warrior_counts, 1, MPI_INT, warrior_counts, 1, MPI_INT, 0, MPI_COMM_WORLD);

    for (int i = 1; i < comm_size; ++i) {
      std::cout << "\x1b[1;1Hhead node: receiving warriors from child node " << i << "\x1b[0K" << std::flush;
      for (int j = 0; j < warrior_counts[i]; ++j) {
        int score;
        MPI_Recv(&score, 1, MPI_INT, i, 0, MPI_COMM_WORLD, 0);
        int serialization_length;
        MPI_Recv(&serialization_length, 1, MPI_INT, i, 0, MPI_COMM_WORLD, 0);
        uint8_t *buffer = new uint8_t[serialization_length];
        MPI_Recv(buffer, serialization_length, MPI_UINT8_T, i, 0, MPI_COMM_WORLD, 0);
        warriors.push_back({.score = score, .w = lib94::deserialize_warrior(buffer)});
        delete[] buffer;
      }
    }

    std::sort(warriors.begin(), warriors.end());

    std::filesystem::path out_path = output_dir / (std::to_string(round) + ".red");

    int total_score = 0;
    for (const scored_warrior &sw : warriors)
      total_score += sw.score;

    std::cout << "\x1b[" << (comm_size + 2) << ";1Hmin/avg/max score from round " << round << ":\n  "
      << (float)(warriors[warriors.size() - 1].score * 100) / (rounds_per_matchup * 2) << "% / "
      << (float)(total_score * 100) / (warriors.size() * rounds_per_matchup * 2) << "% / "
      << (float)(warriors[0].score * 100) / (rounds_per_matchup * 2) << "%\x1b[0K\n\nbest was:\n";

    std::ofstream out_file(out_path);
    if (!out_file) {
      std::cout << "\x1b[?47l\x1b""8" << std::flush;
      std::cerr << "could not open file " << out_path << std::endl;
      int zero = 0;
      MPI_Bcast(&zero, 1, MPI_INT, 0, MPI_COMM_WORLD);
      MPI_Finalize();
      exit(5);
    }

    out_file << ";author evolver\n;name evolved\n";
    out_file << ";score was " << (float)(warriors[0].score * 100) / (rounds_per_matchup * 2) << "%\n";
    out_file << "org " << warriors[0].w->org << '\n';
    std::cout << "  org " << warriors[0].w->org << "\x1b[0K\n";
    for (const lib94::instruction &i : warriors[0].w->instructions) {
      out_file << lib94::instruction_to_string(i) << '\n';
      std::cout << "  " << lib94::instruction_to_string(i) << "\x1b[0K\n";
    }

    std::cout << "\x1b[0J" << std::flush;

    out_file.close();

    if (warriors[0].score == rounds_per_matchup * 2) {
      std::cout << "\x1b[?47l\x1b""8completed in " << round << " rounds\n";
      int zero = 0;
      MPI_Bcast(&zero, 1, MPI_INT, 0, MPI_COMM_WORLD);
      MPI_Finalize();
      exit(0);
    }

    size_t good_warriors = std::min(max_good_warriors, warriors.size());

    for (size_t i = good_warriors; i < warriors.size(); ++i)
      delete warriors[i].w;

    warriors.resize(good_warriors);

    ++round;

  }

}

void randomize_in_range(lib94::instruction &instr, int min, int max) {
  instr.op = (lib94::opcode)(rand() % 16);
  instr.mod = (lib94::modifier)(rand() % 7);
  instr.amode = (lib94::mode)(rand() % 8);
  instr.bmode = (lib94::mode)(rand() % 8);
  instr.anumber = (rand() % (max - min + 1) + min + LIB94_CORE_SIZE) % LIB94_CORE_SIZE;
  instr.bnumber = (rand() % (max - min + 1) + min + LIB94_CORE_SIZE) % LIB94_CORE_SIZE;
}

void randomize(lib94::instruction &instr) {
  instr.op = (lib94::opcode)(rand() % 16);
  instr.mod = (lib94::modifier)(rand() % 7);
  instr.amode = (lib94::mode)(rand() % 8);
  instr.bmode = (lib94::mode)(rand() % 8);
  instr.anumber = rand() % LIB94_CORE_SIZE;
  instr.bnumber = rand() % LIB94_CORE_SIZE;
}

void mutate_step(lib94::warrior *w) {

  if (w->instructions.size() == 0) {
    bool in_range = rand() % 2 == 0;
    lib94::instruction instruction;
    if (in_range)
      randomize_in_range(instruction, 0, 0);
    else
      randomize(instruction);
    w->instructions.push_back(instruction);
    return;
  }

  int rand_inclusive = rand() % (w->instructions.size() + 1);
  int rand_exclusive = rand() % w->instructions.size();
  bool in_range = rand() % 2 == 0;

  switch (rand() % 4) {

  //org
  case 0:
    w->org = rand_exclusive;
    return;

  //insert
  case 1:
    w->instructions.insert(w->instructions.cbegin() + rand_inclusive, (lib94::instruction){});
    if (w->org > rand_inclusive)
      ++w->org;
    rand_exclusive = rand_inclusive;
    //FALLTHRU

  //modify
  case 2:
    if (in_range)
      randomize_in_range(w->instructions[rand_exclusive], -rand_exclusive, w->instructions.size() - rand_exclusive - 1);
    else
      randomize(w->instructions[rand_exclusive]);
    return;

  //remove
  case 3:
    w->instructions.erase(w->instructions.cbegin() + rand_exclusive);
    if (w->org > rand_exclusive)
      --w->org;
    return;

  }

}

void mutate(lib94::warrior *w) {
  do
    mutate_step(w);
  while (rand() % 2 == 0);
}

int score(const lib94::warrior *w) {
  int the_score = 0;
  const lib94::warrior *const warriors[2] = {warrior_to_beat, w};
  for (int i = 0; i < rounds_per_matchup; ++i) {
    const lib94::warrior *winner = lib94::do_round(background_instruction, warriors, 2, 0, true, steps_to_tie);
    if (winner == 0)
      ++the_score;
    else if (winner == w)
      the_score += 2;
  }
  return the_score;
}

void child_main() {

  srand(time(0));

  std::vector<scored_warrior> warriors;

  while (true) {

    int warrior_count;
    MPI_Bcast(&warrior_count, 1, MPI_INT, 0, MPI_COMM_WORLD);

    if (warrior_count == 0) {
      MPI_Finalize();
      exit(0);
    }

    warriors.resize(warrior_count);

    std::cout << "\x1b[" << (comm_rank + 1) << ";1Hchild node " << comm_rank << ": receiving warriors\x1b[0K" << std::flush;

    for (int i = 0; i < warrior_count; ++i) {
      int serialization_length;
      MPI_Bcast(&serialization_length, 1, MPI_INT, 0, MPI_COMM_WORLD);
      uint8_t *buffer = new uint8_t[serialization_length];
      MPI_Bcast(buffer, serialization_length, MPI_UINT8_T, 0, MPI_COMM_WORLD);
      warriors[i].w = lib94::deserialize_warrior(buffer);
      delete[] buffer;
    }

    std::cout << "\x1b[" << (comm_rank + 1) << ";1Hchild node " << comm_rank << ": mutating warriors\x1b[0K" << std::flush;

    for (scored_warrior &sw : warriors)
      mutate(sw.w);

    for (size_t i = 0; i < warriors.size(); ++i) {
      std::cout << "\x1b[" << (comm_rank + 1) << ";1Hchild node " << comm_rank << ": scoring warriors " << (float)(i * 100) / warriors.size() << "%\x1b[0K" << std::flush;
      warriors[i].score = score(warriors[i].w);
    }

    std::cout << "\x1b[" << (comm_rank + 1) << ";1Hchild node " << comm_rank << ": waiting for other child nodes\x1b[0K" << std::flush;

    MPI_Gather(&warrior_count, 1, MPI_INT, 0, 0, 0, 0, MPI_COMM_WORLD);

    for (scored_warrior &sw : warriors) {
      std::vector<uint8_t> serialization;
      lib94::serialize_warrior(sw.w, serialization);
      int serialization_length = serialization.size();
      MPI_Send(&sw.score, 1, MPI_INT, 0, 0, MPI_COMM_WORLD);
      MPI_Send(&serialization_length, 1, MPI_INT, 0, 0, MPI_COMM_WORLD);
      MPI_Send(serialization.data(), serialization_length, MPI_UINT8_T, 0, 0, MPI_COMM_WORLD);
    }

    for (scored_warrior &sw : warriors)
      delete sw.w;

    warriors.clear();

  }

}