main.cpp

#include "algorithm_parser/ast.hpp"
#include "algorithm_parser/ast_printer.hpp"
#include "algorithm_parser/parser.hpp"
#include "algorithm_parser/type_check.hpp"
#include "spec_impl.hpp"

#include <fstream>
#include <iostream>
#include <ranges>
#include <set>
#include <utility>
#include <vector>

constexpr auto &document = compiled_json::spec_json::impl::document;

using namespace AST;
using namespace std::literals::string_view_literals;
using namespace json2cpp;

std::optional<Id> get_autovar(const Algorithm &algo, const LVal &lvl) {
  switch (lvl.val_kind) {
  case LValKind::Var:
    if (algo.auto_decls.contains(dynamic_cast<const Var &>(lvl).id)) {
      return dynamic_cast<const Var &>(lvl).id;
    }
    /* code */
    break;

  case LValKind::Idx:
    return get_autovar(algo, *dynamic_cast<const Idx &>(lvl).arr);

  case LValKind::Fld:
    return get_autovar(algo, *dynamic_cast<const Fld &>(lvl).rcrd);

  default:
    throw std::invalid_argument("Unknown LVal");
  }
  return std::nullopt;
}

std::optional<Id> get_datavar(const Algorithm &algo, const LVal &lvl) {
  switch (lvl.val_kind) {
  case LValKind::Var:
    if (algo.data_decls.contains(dynamic_cast<const Var &>(lvl).id)) {
      return dynamic_cast<const Var &>(lvl).id;
    }
    /* code */
    break;

  case LValKind::Idx:
    return get_datavar(algo, *dynamic_cast<const Idx &>(lvl).arr);

  case LValKind::Fld:
    return get_datavar(algo, *dynamic_cast<const Fld &>(lvl).rcrd);

  default:
    throw std::invalid_argument("Unknown LVal");
  }
  return std::nullopt;
}

std::string to_str(const LVal &lvl) {
  switch (lvl.val_kind) {
  case LValKind::Var:
    return std::string(dynamic_cast<const Var &>(lvl).id);
  case LValKind::Idx: {

    switch (auto &idx = dynamic_cast<const Idx &>(lvl); idx.idx->l_kind) {
    case LiteralKind::Int:
      return to_str(*idx.arr) + "." +
             std::to_string(dynamic_cast<const IntegerLit &>(*idx.idx).value);

    case LiteralKind::Str:
      return to_str(*idx.arr) + "." +
             std::string{dynamic_cast<const StringLit &>(*idx.idx).str};

    default:
      throw std::invalid_argument("Invalid literal kind");
    }
  }

  case LValKind::Fld: {
    auto &fld = dynamic_cast<const Fld &>(lvl);
    return to_str(*fld.rcrd) + "." + to_str(*fld.field_name);
  }
  default:
    throw std::invalid_argument("Invalid LVal kind");
  }
}

std::ostream &operator<<(std::ostream &o, const json2cpp::json &json) {
  switch (json.data.selected) {
  case json2cpp::json::data_t::selected_type::integer:
    return o << json.get<std::int64_t>();
  case json2cpp::json::data_t::selected_type::uinteger:
    return o << json.get<std::uint64_t>();
  case json2cpp::json::data_t::selected_type::floating_point:
    return o << json.get<double>();
  case json2cpp::json::data_t::selected_type::string:
    return o << json.get<std::string_view>();

  default:
    throw std::invalid_argument("Invalid json value");
  }
}

const json2cpp::json &spec_lookup(std::string_view key,
                                  json2cpp::json::const_iterator spec) {

  json2cpp::json::const_iterator ret = spec;

  for (auto subkey : std::views::split(key, "."sv)) {
    const json2cpp::json &curr = *ret.parent_value_;
    std::string_view s_key{subkey.begin(), subkey.end()};
    if (curr.is_array()) {
      for (auto &field : curr) {
        auto field_str = field.get<std::string_view>();
        if (s_key == field_str) {
          ret = field.begin();
          break;
        }
      }
      continue;
    }
    ret = curr[s_key].begin();
  }

  return *ret.parent_value_;
}

std::optional<const AST::Type *> data_lookup_aux(const Algorithm &algo,
                                                 const Fields **fields_p,
                                                 const LVal &lvl) {

  switch (lvl.val_kind) {
  case LValKind::Var: {
    const Var &var = dynamic_cast<const Var &>(lvl);
    if (*fields_p == nullptr) {
      const AST::Record &rcrd = algo.data_decls.at(var.id);
      *fields_p = &algo.type_defs.at(rcrd);
      return &algo.data_decls.at(var.id);
    }

    return (*fields_p)->at(var.id).get();
  }

  case LValKind::Idx: {
    const Idx &idx = dynamic_cast<const Idx &>(lvl);
    auto typ = data_lookup_aux(algo, fields_p, *idx.arr);
    if (!typ || !*fields_p)
      return std::nullopt;

    auto typ_val = typ.value();
    if (typ_val->kind != TypeKind::Array)
      return std::nullopt;

    const Array *arr = dynamic_cast<const Array *>(typ_val);
    const IntegerLit &lit = dynamic_cast<const IntegerLit &>(*idx.idx);

    if (arr->size < lit.value)
      return std::nullopt;

    switch ((arr->sub_type)->kind) {
    case TypeKind::Record:
      *fields_p =
          &algo.type_defs.at(dynamic_cast<const Record &>(*arr->sub_type).id);
      break;
    default:
      break;
    }

    return arr;
  }
  case LValKind::Fld: {

    const Fld &fld = dynamic_cast<const Fld &>(lvl);

    auto typ = data_lookup_aux(algo, fields_p, *fld.rcrd);

    if (!typ || *fields_p == nullptr)
      return std::nullopt;

    if (typ.value()->kind != TypeKind::Record)
      return std::nullopt;

    auto *field_typ = (*fields_p)->at(fld.field_name->id).get();

    if (field_typ->kind == TypeKind::Record) {
      *fields_p =
          &algo.type_defs.at(dynamic_cast<const AST::Record *>(field_typ)->id);
    }

    return field_typ;
  }

  default:
    throw std::invalid_argument("Unexpected lval kind");
  }
}

std::optional<const AST::Type *> data_lookup(const Algorithm &algo,
                                             const LVal &lvl) {
  const Fields *p = nullptr;
  return data_lookup_aux(algo, &p, lvl);
}

struct Step1 {

  struct Results {
    const Algorithm &algo;
    const json2cpp::json::const_iterator &spec;

    struct AutoVar {
      Id var;

      bool operator<(const AutoVar &other) const { return var < other.var; }
    };
    struct DataVar {
      Id var;
      bool operator<(const DataVar &other) const { return var < other.var; }
    };

    struct Operation {
      AutoVar var;
      json::const_iterator chosen_peripheral;
      json::const_iterator operation;
      const std::vector<std::unique_ptr<AST::Expr>> &operands;

      bool operator<(const Operation &other) const {
        return chosen_peripheral < other.chosen_peripheral ||
               (chosen_peripheral == other.chosen_peripheral &&
                operation < other.operation);
      }
    };

    struct DataAssignment {
      enum struct ResourceKind { Spec, Data };

      struct Resource {

        Resource(AutoVar var, const Fld &field)
            : kind(ResourceKind::Spec), id(var.var), lvl(field) {}

        Resource(DataVar var, const LVal &field)
            : kind(ResourceKind::Data), id(var.var), lvl(field) {}

        ResourceKind kind;
        Id id;
        const LVal &lvl;

        friend std::ostream &operator<<(std::ostream &o, const Resource &r) {
          return o << to_str(r.lvl);
        }
      };

      Resource source;
      Resource dest;
    };

    std::map<AutoVar, std::string_view> auto_literals = {};
    std::vector<DataAssignment> data_assignments = {};
    std::map<AutoVar, std::set<json::const_iterator>> output_aliases = {};
    std::vector<std::set<Operation>> compatible_operations = {};
  };

  Step1(const Algorithm &algo)
      : results({.algo = algo, .spec = document["peripherals"sv].begin()}) {}

  const Results &get() {
    if (finalized)
      return results;
    finalized = true;

    for (auto &stmt : results.algo.statements) {
      switch (stmt->kind) {
      case StmtKind::CallStmt:

        callstmt(dynamic_cast<const CallStmt &>(*stmt));
        break;

      case StmtKind::AssignStmt:
        assignstmt(dynamic_cast<const AssignStmt &>(*stmt));
        break;

      default:
        throw std::invalid_argument("Unknown statement type");
      }
    }

    return results;
  }

private:
  Results results;
  bool finalized = false;

  json2cpp::json::const_iterator
  spec_lookup(const LVal &lvl, json2cpp::json::const_iterator curr) const {

    switch (lvl.val_kind) {
    case LValKind::Var: {
      auto &var = dynamic_cast<const Var &>(lvl);
      if (auto var_id = (*curr).find(var.id); var_id != (*curr).end()) {
        return var_id;
      }

      if (results.algo.auto_decls.contains(var.id)) {
        return curr;
      }

      throw std::invalid_argument("Unfindable variable name");
    }

    case LValKind::Idx: {
      const Idx &idx = dynamic_cast<const Idx &>(lvl);

      auto base = spec_lookup(*idx.arr, curr);
      const Literal &lit = *idx.idx;

      switch (lit.l_kind) {
      case LiteralKind::Int:
        return (*base)[dynamic_cast<const IntegerLit &>(lit).value].begin();
      case LiteralKind::Str:
        return (*base)[dynamic_cast<const StringLit &>(lit).str].begin();

      default:
        throw std::logic_error("Invalid lookup literal");
      }
    }

    case LValKind::Fld: {
      const Fld &fld = dynamic_cast<const Fld &>(lvl);
      return spec_lookup(*fld.field_name, spec_lookup(*fld.rcrd, curr));
    }
    default:
      throw std::invalid_argument("Unknown literal");
    }
  }

  void callexpr(const Call &expr) {
    const LVal &lhs = *expr.id;

    switch (lhs.val_kind) {
    case LValKind::Var:
      throw std::logic_error("Unimplemented and unexpected");

    case LValKind::Idx:
      throw std::logic_error("Unimplemented and unexpected");

    case LValKind::Fld: {
      std::set<Results::Operation> co_ops;
      const Fld &fld = dynamic_cast<const Fld &>(lhs);
      for (auto it = document["peripherals"sv].begin();
           it != document["peripherals"sv].end(); ++it) {

        auto periph = spec_lookup(*fld.rcrd, it);

        auto &ops = (*periph)["operations"sv];

        for (auto op_it = ops.begin(); op_it != ops.end(); ++op_it) {
          if (op_it.key() == fld.field_name->id) {
            co_ops.emplace(Results::Operation{
                .var = {get_autovar(results.algo, lhs).value()},
                .chosen_peripheral = periph,
                .operation = op_it,
                .operands = expr.args});
            continue;
          }

          auto op_alias = (*op_it).find("alias"sv);
          if (op_alias == (*op_it).end())
            continue;

          for (auto alias : op_alias.value()) {
            if (!alias.data.is_string())
              throw std::invalid_argument("Invalid operation alias");

            if (*alias.data.get_if_string() != fld.field_name->id)
              continue;

            co_ops.emplace(Results::Operation{
                .var = get_autovar(results.algo, lhs).value(),
                .chosen_peripheral = periph,
                .operation = op_it,
                .operands = expr.args});
          }
        }
      }
      results.compatible_operations.emplace_back(std::move(co_ops));
      return;
    }

    default:
      break;
    }

    throw std::logic_error("Unimplemented");
  }

  void callstmt(const CallStmt &stmt) { return callexpr(*stmt.expr); }

  void assignstmt(const AssignStmt &stmt) {

    const Assign &expr = *stmt.expr;

    switch (expr.rhs->kind) {
    case ExprKind::Call:
      callexpr(dynamic_cast<const Call &>(*expr.rhs));
      switch (expr.lhs->val_kind) {
      case LValKind::Var: {
        std::set<json::const_iterator> iterators;
        for (const auto &op : results.compatible_operations.back()) {
          iterators.emplace(op.operation);
        }

        results.output_aliases.emplace(
            std::make_pair(dynamic_cast<const Var &>(*expr.lhs).id, iterators));
      } break;
      default:
        throw std::logic_error("Unexpected lhs of assignstmt");
      }
      break;

    case ExprKind::Literal:
      switch (expr.lhs->val_kind) {
      case LValKind::Var: {
        auto &lit = dynamic_cast<const Literal &>(*expr.rhs);
        results.auto_literals.emplace(
            std::make_pair(dynamic_cast<const Var &>(*expr.lhs).id,
                           dynamic_cast<const StringLit &>(lit).str));
      } break;
      default:
        throw std::logic_error("Unexpected lhs of assignstmt");
      }
      break;

    case ExprKind::LVal: {

      auto get_resource =
          [&](const LVal &lval) -> Results::DataAssignment::Resource {
        switch (lval.val_kind) {
        case LValKind::Fld: {
          if (auto autovar = get_autovar(results.algo, lval); autovar) {
            return {Results::AutoVar{autovar.value()},
                    dynamic_cast<const Fld &>(lval)};
          }

          auto datavar = get_datavar(results.algo, lval).value();
          return {Results::DataVar{datavar}, lval};
        }

        default:
          throw std::logic_error("Unexpected lhs of assignstmt");
        }
      };

      results.data_assignments.emplace_back(
          get_resource(*expr.lhs),
          get_resource(dynamic_cast<const LVal &>(*expr.rhs)));
      break;
    }

    default:
      throw std::logic_error("Unexpected rhs of assignstmt");
    }
  }
};

struct Step2 {

  Step2(Step1::Results &prev_) : prev(prev_) {}
  Step1::Results &prev;

  enum struct CompareOp { Eq, Le };

  enum struct ArgKind { It, Int, None, Data, Compare };
  struct Arg {

    constexpr Arg(json2cpp::json::const_iterator it_)
        : kind(ArgKind::It), it(it_) {}

    constexpr Arg(std::uint64_t val_) : kind(ArgKind::Int), as_uint(val_) {}

    constexpr Arg(std::nullopt_t) : kind(ArgKind::None), null(std::nullopt) {}

    constexpr Arg(CompareOp op) : kind(ArgKind::Compare), cmp_op(op) {}

    constexpr Arg(const AST::LVal *data_loc_)
        : kind(ArgKind::Data), data_loc(data_loc_) {}

    ArgKind kind;
    union {
      json2cpp::json::const_iterator it;
      std::uint64_t as_uint;
      std::nullopt_t null;
      const AST::LVal *data_loc;
      CompareOp cmp_op;
    };

    bool operator<(const Arg &other) const {
      auto cmp = [&]() {
        switch (kind) {
        case ArgKind::It:
          return it < other.it;
        case ArgKind::Int:
          return as_uint < other.as_uint;

        default:
          throw std::invalid_argument("Unknown arg type");
        }
      };
      return kind < other.kind || (kind == other.kind && cmp());
    }

    friend std::ostream &operator<<(std::ostream &o, const Arg &arg) {
      switch (arg.kind) {
      case ArgKind::It:
        return o << *arg.it;
      case ArgKind::Int:
        return o << std::to_string(arg.as_uint);
      case ArgKind::None:
        return o << "None";
      case ArgKind::Data:
        return o << to_str(*arg.data_loc);

      default:
        throw std::invalid_argument("Unknown arg kind");
      }
    }
  };

  using vArgs_t = std::set<Arg>;

  struct OpAlloc {
    std::map<Step1::Results::AutoVar, json2cpp::json::const_iterator>
        auto_allocs;
    std::vector<Step1::Results::Operation> operations;
    std::vector<std::vector<vArgs_t>> args;
  };

  using OpAllocs = std::vector<OpAlloc>;

  OpAllocs operator()() {
    auto all_opts = all_options(0, {});
    return all_opts;
  }

private:
  OpAllocs all_options(std::size_t i, OpAlloc op_alloc) const {
    if (i == prev.compatible_operations.size())
      return {op_alloc};

    OpAllocs ret;

    for (const Step1::Results::Operation &option :
         prev.compatible_operations[i]) {

      auto auto_alloc = op_alloc.auto_allocs.find(option.var);

      // Filter incompatible autovar allocations
      if (auto_alloc != op_alloc.auto_allocs.end() &&
          auto_alloc->second != option.chosen_peripheral) {
        continue;
      }

      auto vopArgs = valid_operation_args(option, op_alloc);
      if (vopArgs.empty())
        continue;

      OpAlloc local_copy = op_alloc;

      if (auto_alloc == op_alloc.auto_allocs.end()) {
        local_copy.auto_allocs.emplace(option.var, option.chosen_peripheral);
      }

      local_copy.operations.emplace_back(option);
      local_copy.args.emplace_back(vopArgs);

      auto sub_opts = all_options(i + 1, local_copy);

      for (auto &&opt : sub_opts) {
        ret.emplace_back(std::move(opt));
      }
    }

    return ret;
  }

  std::vector<vArgs_t>
  valid_operation_args(const Step1::Results::Operation &option,
                       const OpAlloc &allocs) const {

    std::vector<vArgs_t> ret{};

    const json2cpp::json &inputs = (*option.operation)["inputs"sv];
    [[maybe_unused]] const json2cpp::json &outputs =
        (*option.operation)["outputs"sv];
    if (option.operands.size() != inputs.size())
      return {};

    auto add_to_ret = [&](auto arg) { ret.emplace_back(vArgs_t{arg}); };

    CompareOp op = CompareOp::Eq;

    for (std::size_t i = 0; i < inputs.size(); ++i) {
      json2cpp::json::const_iterator inputs_it = inputs[i].begin();

      auto &curr_arg = option.operands[i];

      if (inputs[i].is_object()) {
        if (inputs[i].size() != 1)
          throw std::invalid_argument("Object size not supported");

        if (auto one_of = inputs[i].find("one_of"sv);
            one_of != inputs[i].end()) {
          inputs_it = inputs[i]["one_of"sv].begin();
        } else if (auto optional_inp = inputs[i].find("optional"sv);
                   optional_inp != inputs[i].end()) {

          if (curr_arg->kind == ExprKind::Literal &&
              dynamic_cast<const Literal &>(*curr_arg).l_kind ==
                  LiteralKind::None) {
            add_to_ret(Arg{std::nullopt});
            continue;
          }
          inputs_it = optional_inp;

        } else {

          throw std::logic_error("Unimplemented");
        }
      }
      auto curr_input = inputs_it->get<std::string_view>();

      // Built in operators
      if (curr_input == "operator.compare") {
        const AST::StringLit &lit =
            dynamic_cast<const AST::StringLit &>(*curr_arg);
        if (lit.str == "==") {
          op = CompareOp::Eq;
        } else if (lit.str == "<=") {
          op = CompareOp::Le;
        } else {
          throw std::invalid_argument("Invalid compare operator");
        }
        add_to_ret(op);
        continue;
      } else if (curr_input == "data") {

        const LVal &lvl = dynamic_cast<const LVal &>(*curr_arg);
        if (auto datavar = get_datavar(prev.algo, lvl); datavar) {

          auto lu = data_lookup(prev.algo, lvl);
          if (!lu)
            throw std::logic_error("Data not found");

          add_to_ret(Arg{dynamic_cast<const LVal *>(curr_arg.get())});
          continue;
        }

        auto auto_var_opt = get_autovar(prev.algo, lvl);
        if (!auto_var_opt)
          throw std::invalid_argument("Not an autovar");
        auto auto_var = auto_var_opt.value();

        auto var_it = allocs.auto_allocs.find({auto_var});
        if (var_it == allocs.auto_allocs.end())
          throw std::invalid_argument("AutoVar Not allocated");

        auto as_str = to_str(lvl);
        std::string_view sview{as_str};

        auto &spec = spec_lookup(sview.substr(sview.find_first_of('.') + 1),
                                 (*var_it->second)["resources"sv].begin());

        add_to_ret(Arg{spec.begin()});
        continue;
      }

      json2cpp::json::const_iterator spec_field_it;
      bool found = false;

      try {
        spec_field_it =
            spec_lookup(curr_input, option.chosen_peripheral->begin()).begin();
        found = true;
      } catch (...) {
      }
      if (!found) {
        try {
          spec_field_it =
              spec_lookup(curr_input, document["peripherals"sv].begin())
                  .begin();
        } catch (...) {
          return {};
        }
      }
      const json2cpp::json &spec_field = *spec_field_it.parent_value_;

      switch (curr_arg->kind) {
      case ExprKind::Literal:
        switch (auto &literal = dynamic_cast<const Literal &>(*curr_arg);
                literal.l_kind) {
        case LiteralKind::NS: {
          const NSLit &nslit = dynamic_cast<const NSLit &>(*curr_arg);
          double arg = nslit.ns * 1e-9;
          vArgs_t v_args;
          for (auto it = spec_field.begin(); it != spec_field.end(); ++it) {
            switch (op) {
            case CompareOp::Eq:
              if (arg == it->get<double>()) {
                v_args.emplace(it);
              }
              break;
            case CompareOp::Le:
              if (it->get<double>() <= arg) {
                v_args.emplace(it);
              }
              break;

            default:
              throw std::invalid_argument("Invalid compare op");
            }
          }
          add_to_ret(v_args);
          op = CompareOp::Eq;
        }

        break;

        case LiteralKind::Int: {
          const IntegerLit &lit = dynamic_cast<const IntegerLit &>(*curr_arg);

          if (spec_field.is_array() && spec_field.size() == 2 &&
              spec_field[0].is_number_integer() &&
              spec_field[1].is_number_integer()) {
            if (spec_field[0].get<std::uint64_t>() <= lit.value &&
                lit.value <= spec_field[1].get<std::uint64_t>()) {
              add_to_ret(Arg{lit.value});
            }
          } else {
            vArgs_t v_args;
            for (auto it = spec_field.begin(); it != spec_field.end(); ++it) {
              switch (op) {
              case CompareOp::Eq:
                if (lit.value == it->get<double>()) {
                  v_args.emplace(it);
                }
                break;
              case CompareOp::Le:
                if (it->get<double>() <= lit.value) {
                  v_args.emplace(it);
                }
                break;

              default:
                throw std::invalid_argument("Invalid compare op");
              }
            }
            add_to_ret(v_args);
          }
          op = CompareOp::Eq;
        }

        break;

        case LiteralKind::Str: {
          if (op != CompareOp::Eq)
            throw std::invalid_argument(
                "Invalid comparison operator for strings");

          const StringLit &slit = dynamic_cast<const StringLit &>(*curr_arg);

          if (spec_field.is_object()) {
            auto &value = spec_field[slit.str];
            add_to_ret(value.begin());
            break;
          }

          for (auto it = spec_field.begin(); it != spec_field.end(); ++it) {
            if (it->get<std::string_view>() == slit.str) {
              add_to_ret(it);
              break;
            }
          }
        }

        break;

        default:
          throw std::invalid_argument("Invalid literal kind");
        }

        break;

      case ExprKind::LVal: {
        const LVal &lval = dynamic_cast<const LVal &>(*curr_arg);

        if (curr_input == "data")
          throw std::logic_error("TODO: Implement");

        auto a_var = get_autovar(prev.algo, lval).value();

        std::string full_name;

        auto alloced = allocs.auto_allocs.find({a_var});

        if (alloced != allocs.auto_allocs.end()) {
          auto sview = to_str(lval);
          sview = sview.substr(sview.find('.'));
          full_name = std::string{alloced->second.key()} + ".resources" +
                      std::string{sview};
        } else if (auto auto_lit = prev.auto_literals.find({a_var});
                   auto_lit != prev.auto_literals.end()) {

          auto &[_, slit] = *auto_lit;
          full_name = std::string{slit.substr(0, slit.find_first_of('.'))} +
                      ".resources" +
                      std::string{slit.substr(slit.find_first_of('.'))};
        } else if (auto output_var = prev.output_aliases.find({a_var});
                   output_var != prev.output_aliases.end()) {

          auto output_op = output_var->second;
          const Step1::Results::Operation *alloced_op = nullptr;

          for (auto j = allocs.operations.size();
               alloced_op == nullptr && j != 0; --j) {
            for (auto &o_op : output_op) {
              if (o_op == allocs.operations[j - 1].operation) {
                alloced_op = &allocs.operations[j - 1];
                break;
              }
            }
          }
          if (alloced_op == nullptr)
            return {};

          const json2cpp::json &output_json =
              (*(alloced_op->operation))["outputs"sv];
          if (auto sz = output_json.size(); sz != 1)
            throw std::invalid_argument("Multiple outputs are not supported");

          for (auto outp : output_json) {
            full_name = std::string{alloced_op->chosen_peripheral.key()} + "." +
                        std::string{outp.get<std::string_view>()};
          }

          if (alloced_op->chosen_peripheral == option.chosen_peripheral) {
            std::string_view sv{full_name};
            sv = sv.substr(sv.find_last_of('.') + 1);
            full_name = std::string{sv};
          }

        } else {
          throw std::invalid_argument("Invalid arg");
        }

        std::string_view path{full_name};

        json2cpp::json::const_iterator spec_lu_it;
        found = false;

        try {
          spec_lu_it =
              spec_lookup(curr_input, option.chosen_peripheral->begin())
                  .begin();
          found = true;
        } catch (...) {
        }

        if (!found) {
          try {
            spec_lu_it =
                spec_lookup(curr_input, document["peripherals"sv].begin())
                    .begin();
          } catch (...) {
            return {};
          }
        }

        const json2cpp::json &spec_lu = *spec_lu_it.parent_value_;

        if (alloced->second == option.chosen_peripheral) {
          path = path.substr(path.find_last_of('.') + 1);
        } else if (path == curr_input) {
          add_to_ret(spec_lu.begin());
          break;
        } else if (path.starts_with(curr_input)) {
          path = path.substr(curr_input.size() + 1);
        }

        if (spec_lu.begin()->is_object()) {
          auto &value = spec_lu[path];
          add_to_ret(value.begin());
          break;
        }

        for (auto it = spec_lu.begin(); it != spec_lu.end(); ++it) {
          if (it->get<std::string_view>() == path) {
            add_to_ret(it);
            break;
          }
        }

      }

      break;

      default:
        throw std::invalid_argument("Invalid expression kind");
      }

      if (ret.size() != i + 1)
        return {};
    }

    return ret;
  }
};

std::ostream &operator<<(std::ostream &o, const Step2::OpAllocs &all_opts) {
  std::size_t i = 0;
  for (auto &opt : all_opts) {
    o << "\n\n";
    o << "auto allocs[" << i << "]:\n";
    for (auto &[var_name, spec_it] : opt.auto_allocs) {
      o << "\t" << var_name.var << ": " << spec_it.key() << "\n";
    }
    ++i;
  }
  return o;
}

bool operator==(const json2cpp::json &a, const json2cpp::json &b) {
  if (a.data.selected != b.data.selected)
    throw std::invalid_argument("a and b have different types");

  switch (a.data.selected) {
  case json2cpp::json::data_t::selected_type::integer:
    return a.get<std::int64_t>() == b.get<std::int64_t>();

  case json2cpp::json::data_t::selected_type::uinteger:
    return a.get<std::uint64_t>() == b.get<std::uint64_t>();

  case json2cpp::json::data_t::selected_type::floating_point:
    return a.get<double>() == b.get<double>();

  default:
    throw std::invalid_argument("Unexpected datatype");
  }
}

bool operator<(const json2cpp::json &a, const json2cpp::json &b) {
  if (a.data.selected != b.data.selected)
    throw std::invalid_argument("a and b have different types");

  switch (a.data.selected) {
  case json2cpp::json::data_t::selected_type::integer:
    return a.get<std::int64_t>() < b.get<std::int64_t>();

  case json2cpp::json::data_t::selected_type::uinteger:
    return a.get<std::uint64_t>() < b.get<std::uint64_t>();

  case json2cpp::json::data_t::selected_type::floating_point:
    return a.get<double>() < b.get<double>();

  default:
    throw std::invalid_argument("Unexpected datatype");
  }
}

bool operator<=(const json2cpp::json &a, const json2cpp::json &b) {
  if (a.data.selected != b.data.selected)
    throw std::invalid_argument("a and b have different types");

  switch (a.data.selected) {
  case json2cpp::json::data_t::selected_type::integer:
    return a.get<std::int64_t>() <= b.get<std::int64_t>();

  case json2cpp::json::data_t::selected_type::uinteger:
    return a.get<std::uint64_t>() <= b.get<std::uint64_t>();

  case json2cpp::json::data_t::selected_type::floating_point:
    return a.get<double>() <= b.get<double>();

  default:
    throw std::invalid_argument("Unexpected datatype");
  }
}

int main(int argc, const char **argv) {
  if (argc != 2 && argc != 3) {
    std::cerr << "Invalid arguments";
    return -1;
  }

  bool verbose = false;

  if (argc == 3) {
    if ("--verbose" != std::string_view{argv[2]}) {
      std::cerr << "Invalid argument";
      return -2;
    }
    verbose = true;
  }

  auto ifs = std::ifstream(argv[1]);
  std::string file_str{std::istreambuf_iterator<char>(ifs),
                       std::istreambuf_iterator<char>()};

  auto algorithm = parse(std::string_view{file_str});
  if (!tt_algorithm(algorithm))
    throw std::logic_error("Type checking failed");

  if (verbose)
    std::cout << "\n\n" << algorithm;

  Step1 step1{algorithm};
  auto results1 = step1.get();

  Step2 step2(results1);
  auto results2 = step2();

  while (results2.size() != 1) {

    if (verbose) {
      std::cout << "-------------------------------\n";
      std::cout << "Possible allocations:\n";
      std::cout << results2;
    }

    std::map<Step1::Results::AutoVar, std::set<json2cpp::json::const_iterator>>
        periph_options;
    for (auto &alloc : results2) {
      for (auto &auto_alloc : alloc.auto_allocs) {
        auto it = periph_options.find(auto_alloc.first);
        if (it == periph_options.end()) {
          periph_options[auto_alloc.first] = {};
        }

        periph_options[auto_alloc.first].emplace(auto_alloc.second);
      }
    }

    for (auto &[var, opts] : periph_options) {
      if (opts.size() == 1)
        continue;

      std::size_t min_size = -1;
      std::set<json2cpp::json::const_iterator> chosen;

      for (auto &it : opts) {
        const auto &resources = (*it)["resources"sv];
        if (resources.size() < min_size) {
          min_size = resources.size();
          chosen = {it};
        } else if (resources.size() == min_size) {
          chosen.emplace(it);
        }
      }

      if (opts.size() != chosen.size()) {
        opts = chosen;
        goto update;
      }

      min_size = -1;
      chosen = {};

      for (auto &it : opts) {
        const auto &cap = (*it)["capacity"sv];
        if (cap.size() < min_size) {
          min_size = cap.size();
          chosen = {it};
        } else if (cap.size() == min_size) {
          chosen.emplace(it);
        }
      }

      if (chosen.size() == 1) {
        opts = chosen;
        goto update;
      }

      std::map<std::string_view, std::pair<json2cpp::json::const_iterator,
                                           json2cpp::json::const_iterator>>
          cap_min_max;

      for (auto opts_it = opts.begin(); opts_it != opts.end(); ++opts_it) {
        auto &it = *opts_it;
        const auto &cap = (*it)["capacity"sv];
        for (auto cap_it = cap.begin(); cap_it != cap.end(); ++cap_it) {

          auto key = cap_it.key();
          auto &values = *cap_it;

          json2cpp::json::const_iterator min = values.begin(),
                                         max = values.begin();

          for (auto v_it = values.begin(); v_it != values.end(); ++v_it) {
            if (*v_it < *min) {
              min = v_it;
            } else if (*max < *v_it) {
              max = v_it;
            }
          }

          auto curr = cap_min_max.find(key);
          if (curr == cap_min_max.end()) {
            cap_min_max[key] = std::make_pair(min, max);
            continue;
          }
          const auto &[c_min, c_max] = cap_min_max[key];

          if (*c_min == *min && *c_max == *max) {
            opts.erase(it);
            goto update;
          }
          // [c_min < min] ==> keep this
          if (*c_min < *min && *max <= *c_max) {
            opts.erase(*--opts_it);
            goto update;
          }
          // [max < c_max] ==> keep this
          if (*max < *c_max && *c_min <= *min) {
            opts.erase(*--opts_it);
            goto update;
          }

          // [min < c_min] ==> keep prev
          if (*c_min < *min && *max <= *c_max) {
            opts.erase(it);
            goto update;
          }
          // [c_max < max] ==> keep prev
          if (*c_max < *max && *c_min <= *min) {
            opts.erase(it);
            goto update;
          }

          throw std::logic_error("Invalid comparison");
        }
      }
    }

    throw std::logic_error("Insufficient strategy");

  update:

    for (auto i = results2.size(); i != 0; --i) {
      const auto &alloc = results2[i - 1];
      for (auto &[key, periph] : alloc.auto_allocs) {
        auto it = periph_options.at(key);
        if (it.contains(periph))
          continue;

        results2.erase(results2.begin() + i - 1);
        break;
      }
    }
  }

  auto &alloc = results2[0];

  for (auto &op : alloc.args) {
    for (auto &val_set : op) {
      if (val_set.size() == 1)
        continue;

      val_set = {*val_set.rbegin()};
    }
  }

  if (verbose)
    std::cout << "\n-------------------------------\n\n\n";

  std::cout << "Chosen alloccation:\n";
  std::cout << "\tAllocs:\n";
  for (auto &var : alloc.auto_allocs) {
    std::cout << "\t\t" << var.first.var << ": " << var.second.key() << "\n";
  }

  std::cout << "\n\tData:\n";
  for (auto &data : results1.algo.data_decls) {
    std::cout << "\t\t" << data.second.id << " " << data.first << "\n";
  }

  std::cout << "\n\tInitial assignments:\n";
  for (auto &data_assign : results1.data_assignments) {
    std::cout << "\t\t" << data_assign.dest << " = " << data_assign.source
              << "\n";
  }

  std::cout << "\n\tOperations:\n";
  for (std::size_t i = 0; i < alloc.operations.size(); ++i) {
    std::cout << "\t\t" << alloc.operations[i].chosen_peripheral.key() << "."
              << alloc.operations[i].operation.key() << "(";
    bool first = true;
    for (const auto &arg : alloc.args[i]) {
      if (arg.begin()->kind == Step2::ArgKind::Compare)
        continue;
      if (!first) {
        std::cout << ", ";
      }
      std::cout << *arg.begin();
      first = false;
    }
    std::cout << ")\n";
  }

  return 0;
}