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

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

#define arg(a) program[a]
#define iarg(a) std::stoi(program[a])
#define farg(a) std::stof(program[a])

#define cset(k, v) config[k] = v
#define csetn(k, v) config[k] = std::to_string(v);

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

enum EstateMode {
    EVEN_REDIST,
    INHERITANCE
};

PopSimpleResult pop_simple(
    int duration = 1000,
    int max_age = 81,
    int init_balance_each = 5000,
    int starting_balance_new = 1000,

    // logistic function params
    int N_o = 10,
    int N_f = 10000,
    float k = 0.04,

    // estate mode
    std::string estate_mode = "inheritance"
) {

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

    // estate mode
    EstateMode emode = INHERITANCE;

    // 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) {
                switch (emode) {
                    case INHERITANCE: {
                        int bal = balances.at(id);
                        ages.erase(ages.begin()+id);
                        balances.erase(balances.begin()+id);
                        int recipient_id = random_item(balances, true);
                        balances.at(recipient_id) += bal;
                        break;
                    }
                    case EVEN_REDIST: {
                        int bal = balances.at(id);
                        ages.erase(ages.begin()+id);
                        balances.erase(balances.begin()+id);
                        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);
            }
        }

        if (year % 100 == 0) {
            int avg = 0;
            int avg_age = 0;
            for (int i = 0; i < balances.size(); i++) {
                avg += balances[i];
            }
            for (int i = 0; i < balances.size(); i++) {
                avg_age += ages[i];
            }
            avg /= balances.size();
            avg_age /= ages.size();
            std::cout << std::endl;
            std::cout << "year: " << year << std::endl;
            std::cout << "avg bal: " << avg << std::endl;
            std::cout << "avg age: " << avg_age << std::endl;
            std::cout << "players: " << balances.size() << std::endl;
        }

        // 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 (emode) {
        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 = ages;
    return res;
}