#include <iostream>

#include <string.h>
#include <stdlib.h>
#include <unistd.h>

#include "desired_resources.h"

NTL_CLIENT


// For a given dpfreq, a worker will always transmit a specific
// number of dps per second (this needs to be determined
// experimentally). A dpnode can only accept a specific number
// of dps per second.
// freq_reduction_threshold should be equal to the max number of
// workers that a dpnode can handle when the dpfreq == 4294967

// Input: order, total_workers, total_dpnodes,
//        GB_mem_per_node, freq_reduction_threshold
// Output: desired_dpnodes, max_workers, dpfreq
void desired_resources(const ZZ &order, unsigned short total_workers,
    unsigned short total_dpnodes, unsigned short GB_mem_per_node,
    unsigned short freq_reduction_threshold, unsigned short &desired_dpnodes,
    unsigned int &max_workers, unsigned int &dpfreq)
{
    // One point in how many is a DP by default?
    unsigned int dpscale = 1000;
    ZZ sorder = SqrRoot(order);

    // How many DPnodes should we use for a problem of this size?

    // 338 is bytes per DP in the table.  10 is a safety factor.
    ZZ dpnumerator = sorder * 338 * 10;
    ZZ dpdenominator;
    dpdenominator = GB_mem_per_node;
    dpdenominator *= 1000000000UL;  // Convert the above line to B
    ZZ dpnodes = (dpnumerator / (dpdenominator * dpscale)) + 1;
    if (dpnodes > total_dpnodes) {
	desired_dpnodes = total_dpnodes;
	ZZ zzdpscale = dpnumerator / (dpdenominator * total_dpnodes);
	if (NumBits(zzdpscale) > 31) {
	    dpscale = 4294967295U;
	} else {
	    dpscale = trunc_long(zzdpscale, 31);
	}
    } else {
	desired_dpnodes = total_dpnodes;
    }

    // How many workers would we like to use?
    ZZ sorder23 = sorder >> 23;
    if (NumBits(sorder23) > 30) {
	// Just use all the workers we can find
	max_workers = 4294967295U;  // 2^32 - 1
    } else {
	max_workers = trunc_long(sorder23,31) + 1;
    }
    if (max_workers > total_workers) {
	max_workers = total_workers;
    }

    // By default, 1 in dpscale points are distinguished points.
    dpfreq = 4294967295U / dpscale;

    if ((float)max_workers/desired_dpnodes > freq_reduction_threshold) {
	// if the ratio of workers to dpnodes is too high
	unsigned int max_dpfreq = (4294967295U/1000) / (((float)max_workers/desired_dpnodes)/freq_reduction_threshold);
	// max dpfreq for the given ratio of workers to dpnodes
	if (dpfreq > max_dpfreq) {
	    dpfreq = max_dpfreq;
	}
    }

#ifdef DPFREQ_DIVISOR
    // not a nice solution, but gives you the extra control if you need it
    dpfreq /= DPFREQ_DIVISOR;
#endif

    // Orders smaller than 100*scale^2 behave specially, in order to
    // avoid DP-free cycles
    ZZ orderlimit;
    orderlimit = 100;
    orderlimit *= dpscale;
    orderlimit *= dpscale;

    if (order < 1000) {
	// Just make every point a DP
	dpfreq = 4294967295U;
    } else if (order < orderlimit) {
	// The frequency of DPs should be 10/sqrt(order) to avoid
	// a DP-free cycle, so dpfreq = (10*2^32)/sqrt(order)
	ZZ f = (to_ZZ(10) << 32) / SqrRoot(order);
	dpfreq = trunc_long(f, 31);
    }
}

#ifdef TEST_DESIRED_RESOURCES
int main(int argc, char **argv)
{
    if (argc != 5) {
	std::cerr << "Usage: " << argv[0] << " num_workers num_dpnodes GB_mem_per_node freq_reduction_threshold\n";
	return 1;
    }

    unsigned short total_workers = strtoul(argv[1], NULL, 10);
    unsigned short total_dpnodes = strtoul(argv[2], NULL, 10);
    unsigned short GB_mem_per_node = strtoul(argv[3], NULL, 10);
    unsigned short freq_reduction_threshold = strtoul(argv[4], NULL, 10);

    cout << "# log_2(order) dpnodes workers dpfreq\n";
    for (int i=44; i<=92; ++i) {
	unsigned short desired_dpnodes;
	unsigned int max_workers;
	unsigned int dpfreq;
	ZZ order;
	order = 1;
	order <<= i;
	order += 1;
	desired_resources(order, total_workers, total_dpnodes, GB_mem_per_node, freq_reduction_threshold, desired_dpnodes, max_workers, dpfreq);
	cout << i << " " << desired_dpnodes << " " << max_workers << " "
		<< dpfreq << "\n";
    }
}
#endif