// model with logistic population growth & basic 
// inflation + estate redistribution systems

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

#include <argparse/argparse.hpp>

#include "econ.hpp"
#include "io.hpp"
#include "util.hpp"

// std::map<std::string, std::string> config;
argparse::ArgumentParser program("pop-simple");

#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);

std::vector<int> ages;
std::vector<int> balances;

enum estate_mode {
    EVEN_REDIST,
    INHERITANCE
};

// used by kill_even_redist

int redist_total = 0;
int redist_total_next = 0;

// pop modeling functions

/**
 * @brief kill participant of ID and redistribute their wealth 
 * to all other remaining participants evenly
 * 
 * @param id ID of participant to kill
 */
void kill_even_redist(int id) {
    int bal = balances.at(id);
    ages.erase(ages.begin()+id);
    balances.erase(balances.begin()+id);
    redist_total_next += bal;
}

/**
 * @brief kill participant of ID and give it to one random
 * participant as inheritance
 * 
 * @param id ID of participant to kill
 */
void kill_inheritance(int id) {
    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;
}

/**
 * @brief Age all participants by 1 year & kill any that are 
 * over the age limit
 * 
 * @param mode Redistribution mode
 */
void age_all(estate_mode mode) {
    for (int i = 0; i < ages.size(); i++) {
        ages.at(i) += 1;
        if (ages.at(i) > iarg("MAX_AGE")) {
            switch (mode) {
                case INHERITANCE:
                    kill_inheritance(i);
                    break;
                case EVEN_REDIST:
                    kill_even_redist(i);
                    break;
            }
        }
    }
}

int logistic_population_func(int year) {
    // https://en.wikipedia.org//wiki/Logistic_function
    return iarg("N_o") + (iarg("N_f") - iarg("N_o")) / 
        (1 + exp(-farg("GROWTH_RATE") * (year - ci("DURATION")/2)));
}

void adjust_population(int year) {
    int diff = logistic_population_func(year) - balances.size();
    if (diff > 0) {
        for (int i = 0; i < diff; i++) {
            ages.push_back(1);
            balances.push_back(iarg("STARTING_BALANCE_NEW"));
        }
    }
}

int main(int argc, char* argv[]) {
    // set config via cli args
    program.add_argument("estate", "-e")
        .help("Estate distribution mode (redistribute, inheritance)")
        .default_value(INHERITANCE);

    program.add_argument("time", "-t")
        .help("Duration in years of simulation.")
        .default_value(1000);

    try {
        program.parse_args(argc, argv);
    }
    catch (const std::runtime_error& err) {
        std::cerr << err.what() << std::endl;
        std::cerr << program;
        std::exit(1);
    }

    // model params
    csetn("MAX_AGE", 81);
    csetn("INIT_BALANCE_EACH", 5000);
    csetn("STARTING_BALANCE_NEW", 1000); // effectively inflation rate
    // logistic growth params
    csetn("N_o", 10);
    csetn("N_f", 10000);
    csetn("GROWTH_RATE", 0.04);

    // init state
    ages.assign(ci("N_o"), 1);
    balances.assign(ci("N_o"), ci("INIT_BALANCE_EACH"));

    // estate mode
    // estate_mode emode = INHERITANCE;
    estate_mode emode = EVEN_REDIST;

    for (int year = 0; year < ci("DURATION"); year++) {
        // std::cout << year << "\t" << balances.size() << std::endl;
        age_all(emode);
        adjust_population(year);
        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";

    return 0;
}