#include <lib94/lib94.hpp>
#include <filesystem>
#include <algorithm>
#include <iostream>
#include <fstream>
#include <random>
#include <vector>
#include <ctime>

const lib94::warrior *warrior_to_beat;

constexpr int score_rounds = 50;
constexpr int rounds_to_tie = 1000000;
constexpr int warriors_per_generation = 100;

struct winfo {
  lib94::warrior w;
  int score;
};

std::vector<winfo> current_generation;
int generation_number;

std::default_random_engine prng;

//sets random opcode, modifier, and modes
void random_instruction(lib94::instruction &instr) {
  std::uniform_int_distribution<int> opcode_dist(lib94::DAT, lib94::NOP);
  std::uniform_int_distribution<int> modifier_dist(lib94::A, lib94::I);
  std::uniform_int_distribution<int> mode_dist(lib94::IMMEDIATE, lib94::B_INCREMENT);
  instr.op = (lib94::opcode)opcode_dist(prng);
  instr.mod = (lib94::modifier)modifier_dist(prng);
  instr.amode = (lib94::mode)mode_dist(prng);
  instr.bmode = (lib94::mode)mode_dist(prng);
}

//sets random numbers
void random_numbers(lib94::instruction &instr) {
  std::uniform_int_distribution<lib94::number_t> number_dist(0, LIB94_CORE_SIZE - 1);
  instr.anumber = number_dist(prng);
  instr.bnumber = number_dist(prng);
}

//sets random numbers within a range; left_limit and right_limit are assumed to be >= -LIB94_CORE_SIZE
void random_numbers(lib94::instruction &instr, std::make_signed_t<lib94::number_t> left_limit, std::make_signed_t<lib94::number_t> right_limit) {
  std::uniform_int_distribution<std::make_signed_t<lib94::number_t>> number_dist(left_limit, right_limit);
  instr.anumber = (number_dist(prng) + LIB94_CORE_SIZE) % LIB94_CORE_SIZE;
  instr.bnumber = (number_dist(prng) + LIB94_CORE_SIZE) % LIB94_CORE_SIZE;
}

void random_warrior(lib94::warrior &w) {

  std::uniform_int_distribution<int> length_dist(1, 10);
  int length = length_dist(prng);
  w.instructions.resize(length);

  for (int i = 0; i < length; ++i) {
    random_instruction(w.instructions[i]);
    //allow numbers to point up to one instruction before or after warrior
    random_numbers(w.instructions[i], -i - 1, length - i);
  }

  std::uniform_int_distribution<int> org_dist(0, length - 1);
  w.org = org_dist(prng);

}

void mutate(lib94::warrior &w) {

  std::uniform_int_distribution<int> case_dist(0, 9);

  while (true) {

    std::uniform_int_distribution<int> instr_dist(0, w.instructions.size() - 1);
    int instr = instr_dist(prng);

    lib94::instruction i;

    switch (case_dist(prng)) {

      case 0:
        return;

      case 1:
      case 2:
        random_instruction(w.instructions[instr]);
        break;

      case 3:
      case 4:
        random_numbers(w.instructions[instr], -instr - 1, w.instructions.size() - instr);
        break;

      case 5:
        random_numbers(w.instructions[instr]);
        break;

      case 6:
        random_instruction(i);
        random_numbers(i, -instr - 1, w.instructions.size() + 1 - instr);
        w.instructions.insert(w.instructions.cbegin() + instr, i);
        break;

      case 7:
        random_instruction(i);
        random_numbers(i, -instr - 2, w.instructions.size() - instr);
        w.instructions.insert(w.instructions.cbegin() + instr + 1, i);
        break;

      case 8:
      case 9:
        if (w.instructions.size() > 1)
          w.instructions.erase(w.instructions.cbegin() + instr);
        break;

    }

  }

}

void crossover(const lib94::warrior &a, const lib94::warrior &b, lib94::warrior &c) {

  std::uniform_int_distribution<int> b_offset_dist(0, a.instructions.size());
  int b_offset = b_offset_dist(prng);

  std::uniform_int_distribution<int> b_start_dist(b_offset, std::min(a.instructions.size(), b_offset + b.instructions.size()));
  int b_start = b_start_dist(prng);

  std::uniform_int_distribution<int> b_end_dist(b_start, std::min(a.instructions.size(), b_offset + b.instructions.size()));
  int b_end = b_end_dist(prng);

  c.instructions.resize(a.instructions.size());
  std::copy(a.instructions.cbegin(), a.instructions.cbegin() + b_start, c.instructions.begin());
  std::copy(b.instructions.cbegin() + b_start - b_offset, b.instructions.cbegin() + b_end - b_offset, c.instructions.begin() + b_start);
  std::copy(a.instructions.cbegin() + b_end, a.instructions.cend(), c.instructions.begin() + b_end);

  c.org = a.org;

}

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

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

  prng.seed(time(0));

  if (argc != 3) {
    std::cerr << "please call this with the path to a warrior to beat followed by the path to a directory to store the outputs." << std::endl;
    return 1;
  }

  try {
    warrior_to_beat = lib94::compile_warrior_from_file(argv[1]);
  }
  catch (const lib94::compiler_exception &ex) {
    std::cerr << "error compiling " << argv[1] << ": " << ex.message << " on line " << ex.source_line_number << '.' << std::endl;
    return 1;
  }

  if (!warrior_to_beat) {
    std::cerr << "could not open " << argv[1] << '.' << std::endl;
    return 1;
  }

  std::filesystem::path output_dir(argv[2]);

  if (std::filesystem::exists(output_dir)) {
    std::cerr << output_dir << " already exists." << std::endl;
    return 1;
  }

  if (!std::filesystem::create_directories(output_dir)) {
    std::cerr << "could not create " << output_dir << '.' << std::endl;
    return 1;
  }

  generation_number = 1;
  current_generation.resize(warriors_per_generation);
  for (int i = 0; i < warriors_per_generation; ++i) {
    random_warrior(current_generation[i].w);
    current_generation[i].w.name = "g1w" + std::to_string(i + 1);
    current_generation[i].score = -1;
  }

  int last_best_score = 0;

  while (true) {

    std::shuffle(current_generation.begin(), current_generation.end(), prng);

    std::cout << "generation " << generation_number << ". " << std::flush;

    int best_score = 0;
    const lib94::warrior *best_w = &current_generation[0].w;

    int total_score = 0;

    for (winfo &wi : current_generation) {

      if (wi.score == -1) {

        const lib94::warrior *ws[2] = { warrior_to_beat, &wi.w };
        int score = score_rounds;

        for (int i = 0; i < score_rounds; ++i) {
          const lib94::warrior *winner = lib94::do_round(background, ws, 2, 0, true, rounds_to_tie);
          if (winner == warrior_to_beat)
            --score;
          else if (winner != 0)
            ++score;
        }

        wi.score = score;

      }

      if (wi.score > best_score) {
        best_score = wi.score;
        best_w = &wi.w;
      }

      total_score += wi.score;

    }

    std::cout << "average / maximum score " << (double)total_score / current_generation.size() << " / " << best_score << ". " << std::endl;

    if (best_score > last_best_score) {

      last_best_score = best_score;

      std::filesystem::path file_path = output_dir / (std::to_string(generation_number) + ".red");
      std::ofstream file(file_path);

      if (!file) {
        std::cerr << "could not open " << file_path << '.' << std::endl;
        return 1;
      }

      file << ";author evolver\n;name " << best_w->name << "\n;score = " << best_score << " / " << score_rounds * 2 << "\norg " << best_w->org << '\n';
      for (const lib94::instruction &i : best_w->instructions)
        file << lib94::instruction_to_string(i) << '\n';

      file.close();

      if (best_score == score_rounds * 2) {
        std::cout << "perfect score reached." << std::endl;
        return 0;
      }

    }

    ++generation_number;

    std::vector<winfo> new_warriors;
    new_warriors.reserve(current_generation.size() + 2);

    new_warriors.push_back({.w = *best_w, .score = best_score});

    std::uniform_int_distribution<int> score_at_least(best_score > score_rounds ? score_rounds : 0, best_score > score_rounds ? score_rounds * 2 + 1 : score_rounds);

    for (const winfo &wi : current_generation)
      if (wi.score >= score_at_least(prng))
        new_warriors.push_back(wi);

    std::uniform_int_distribution<int> warrior_dist(0, new_warriors.size() - 1);

    do {

      lib94::warrior w1 = new_warriors[warrior_dist(prng)].w;
      lib94::warrior w2 = new_warriors[warrior_dist(prng)].w;

      mutate(w1);
      mutate(w2);

      lib94::warrior w;
      crossover(w1, w2, w);
      w.name = "g" + std::to_string(generation_number) + "w" + std::to_string(new_warriors.size() + 1);

      new_warriors.push_back({.w = w, .score = -1});

    } while (new_warriors.size() < warriors_per_generation);

    current_generation = std::move(new_warriors);

  }

}