#include "mes_add1.hpp"
#include "mes_utils.hpp"

#include <iostream>
#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <algorithm>
#include <numeric>
#include <stdexcept>
#include <string>
#include <optional>

using namespace std;

using VoterId = string;
using CandidateId = string;
using Utility = double;
using CostMap = unordered_map<CandidateId, double>;
using BudgetMap = unordered_map<VoterId, double>;
using ApproversMap = unordered_map<CandidateId, vector<pair<VoterId, Utility>>>;

vector<CandidateId> equal_shares_add1(
    const vector<VoterId>& voters,
    const vector<CandidateId>& projects,
    const CostMap& cost,
    ApproversMap approvers_utilities,
    double total_budget) {
    // Step 1: Start with a fixed budget method
    vector<CandidateId> mes = equal_shares_utils(voters, projects, cost, approvers_utilities, total_budget);

    // Step 2: Initialize integral per-voter budget
    double budget = static_cast<int>(total_budget / voters.size()) * voters.size();
    double current_cost = accumulate(mes.begin(), mes.end(), 0.0,
                                      [&cost](double sum, const string& c) { return sum + cost.at(c); });

    // Step 3: Add projects while the budget allows
    while (true) {
        bool is_exhaustive = true;
        for (const auto& extra : projects) {
            if (find(mes.begin(), mes.end(), extra) == mes.end() && current_cost + cost.at(extra) <= total_budget) {
                is_exhaustive = false;
                break;
            }
        }

        // If exhaustive, stop
        if (is_exhaustive) {
            break;
        }

        // Try the next higher budget
        double next_budget = budget + voters.size();
        vector<string> next_mes = equal_shares_utils(voters, projects, cost, approvers_utilities, next_budget);

        current_cost = accumulate(next_mes.begin(), next_mes.end(), 0.0,
                                  [&cost](double sum, const string& c) { return sum + cost.at(c); });
        if (current_cost <= total_budget) {
            // Update budget and project list if within budget
            budget = next_budget;
            mes = next_mes;
        } else {
            break;
        }
    }

    return mes;
}
