#include <iostream>
#include <vector>
#include <map>
#include <cstdlib>
#include <algorithm>
#include <string>

#include "util.hpp"
#include "population.hpp"
#include "econ.hpp"

struct PopSimpleResult {
    std::vector<int> ages;
    std::vector<int> balances;
};

enum EstateMode {
    EVEN_REDIST,
    INHERITANCE
};

PopSimpleResult pop_simple(
    int duration,
    int max_age,
    int init_balance_each,
    int starting_balance_new,

    // logistic function params
    int N_o,
    int N_f,
    float k,

    // estate mode
    EstateMode estate_mode
) {

    std::vector<int> ages;
    std::vector<int> balances;
    ages.assign(N_o, 1);
    balances.assign(N_o, init_balance_each);

    // Changed/used only if emode=redist
    int redist_total = 0;
    int redist_total_next = 0;

    for (int year = 0; year < duration; year++) {
        // std::cout << year << "\t" << balances.size() << std::endl;

        // age all
        for (int id = 0; id < ages.size(); id++) {
            ages.at(id) += 1;
            if (ages.at(id) > max_age) {
                int bal = balances.at(id);
                ages.erase(ages.begin()+id);
                balances.erase(balances.begin()+id);
                switch (estate_mode) {
                    case INHERITANCE: {
                        int recipient_id = random_item(balances, true);
                        balances.at(recipient_id) += bal;
                        break;
                    }
                    case EVEN_REDIST: {
                        redist_total_next += bal;
                        break;
                    }
                }
            }
        }

        // (roughly) simulate logistic population growth
        int diff = logistic_population_func(
        year, N_o, N_f, k, duration) - balances.size();
        if (diff > 0) {
            for (int i = 0; i < diff; i++) {
                ages.push_back(1);
                balances.push_back(starting_balance_new);
            }
        }

        for (int j = 0; j < balances.size(); j++) {

            // redistribute estates
            balances.at(j) += redist_total / balances.size();

            // make 3 arbitrary spending decisions -> 30% of savings
            for (int k = 0; k < 3; k++) {
                economic_decision_simple(balances, j, 0.1);
            }
        }

        // reset redist tally
        redist_total = redist_total_next;
        redist_total_next = 0;
    }

    // std::sort(balances.begin(), balances.end(), std::greater<int>());
    std::string ext;
    switch (estate_mode) {
        case EVEN_REDIST:
            ext = "even";
            break;
        case INHERITANCE:
            ext = "inheritance";
            break;
    }
    // dump_balances_and_ages_csv(balances, ages, "pop-simple-results-"+ext);
    std::cout << "Finished with " << balances.size() << " balances\n";

    PopSimpleResult res;
    res.ages = ages;
    res.balances = balances;
    return res;
}