/**
 * Main author: Jérôme Pouiller <jerome.pouiller@sysmic.org>
 * Last update: Fri, 13 May 2011 10:59:35 +0200
 * Licence: GPL
 */
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <time.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <errno.h>
#include <string.h>
#include <sched.h>

static float compute_correction(float objective);
static void consume(float objective, int period);
static long timediff(const struct timespec *ta, const struct timespec *tb);
static void pr_stats(long time_cur, struct rusage *usage_cur);

// Main work
static void consume(float objective, int period) {
    float out = objective;
    for(;;) {
        struct timespec startwork, endwork;
        clock_gettime(CLOCK_REALTIME, &startwork);

        //estimate how much the controlled process is using the cpu in its working interval
        out += compute_correction(objective);

        long twork, tsleep;   //working and sleeping intervals
        //adjust work and sleep time slices
        twork = period * out;
        if (twork > period)
            twork = period;
        if (twork < 0)
            twork = 0;
        tsleep = period - twork;
        //printf("s: %ld, w: %ld, l: %f\n", tsleep, twork, out);

        do {
            clock_gettime(CLOCK_REALTIME, &endwork);
        }  while (timediff(&endwork, &startwork) < twork);
        usleep(tsleep);
    }
}

// Compute correction using PID regulator
static float compute_correction(float objective) {
    struct timespec w;
    clock_gettime(CLOCK_REALTIME, &w);
    struct rusage u;
    getrusage(RUSAGE_SELF, &u);    

    static long time_prev = 0;
    long time_cur = w.tv_sec * 1000000  + w.tv_nsec / 1000;
    static long usage_prev = 0;
    long usage_cur = (u.ru_utime.tv_sec + u.ru_stime.tv_sec) * 1000000  + (u.ru_utime.tv_usec + u.ru_stime.tv_usec);

    pr_stats(time_cur, &u);

    static float Ep = 0.;
    static float Ei = 0.;
    static float Ed = 0.;
    static const float Kp = 0.2;
    static const float Ki = 0.2;
    static const float Kd = 0.0;

    // It is the first call
    if (time_prev == 0) {
        usage_prev = usage_cur;
        time_prev = time_cur;
        return 0.;
    }

    // Frequency of measure
    if (time_prev + 25000 > time_cur) 
        return  0.;

    Ed = (objective - (float) (usage_cur - usage_prev) / (float) (time_cur - time_prev)) - Ep;
    Ep += Ed;
    Ei += Ep;
    //printf("(p:%f i:%f d:%f), usage: %ld / %ld\n", Ep, Ei, Ed, (usage_cur - usage_prev), (time_cur - time_prev));

    usage_prev = usage_cur;
    time_prev = time_cur;

    return  Kp * (Ep + Ki * Ei + Kd * Ed);
}

static void pr_stats(long time_cur, struct rusage *usage_cur) {
    static long time_prev = 0;
    static struct rusage usage_prev;

    if (time_prev == 0) {
        memcpy(&usage_prev, usage_cur, sizeof(usage_prev));
        time_prev = time_cur;
        return ;
    }
    // Display stats only each second
    if (time_prev + 1000000 > time_cur) 
        return ;

    float user = ((usage_cur->ru_utime.tv_sec - usage_prev.ru_utime.tv_sec) * 1000000 
            + (usage_cur->ru_utime.tv_usec - usage_prev.ru_utime.tv_usec)) / (float) (time_cur - time_prev);
    float system = ((usage_cur->ru_stime.tv_sec - usage_prev.ru_stime.tv_sec) * 1000000 
            + (usage_cur->ru_stime.tv_usec - usage_prev.ru_stime.tv_usec)) / (float) (time_cur - time_prev);
    float total = user + system;
    printf("Usage: %02.1f%% (u: %02.1f%% s: %02.1f%%)\n", total * 100, user * 100, system * 100);

    memcpy(&usage_prev, usage_cur, sizeof(usage_prev));
    time_prev = time_cur;
}

// Return ta-tb in microseconds (no overflow checks!)
static long timediff(const struct timespec *ta, const struct timespec *tb) {
    return (ta->tv_sec - tb->tv_sec) * 1000000 + (ta->tv_nsec - tb->tv_nsec) / 1000;
}

int main(int argc, char **argv) {
    // Period of control output in microseconds
    // Should be between 1000us and 1000000us
    // Should be at least 2 times 1/HZ value of kernel
    int period = 20000;
    // CPU load
    float objective;
    char *end;

    if (argc < 2) {
        fprintf(stderr, "Usage: %s CPU_PERCENT [PERIOD]\n", argv[0]);
        fprintf(stderr, "  CPU_PERCENT is cpu load we try to use in %%\n");
        fprintf(stderr, "  PERIOD is time quantum in microseconds. Should be at least 2 time quantum period of kernel scheduler et must be less than 1s\n");
        fprintf(stderr, "  Notice you will have better results if you launch this program as root\n");
        return 1; 
    }
    objective = strtol(argv[1], &end, 10) / 100.;
    if (!*argv[1] || *end || objective < 0. || objective > 1.) {
        fprintf(stderr, "Invalid cpu load\n");
        return 1;
    }
    if (argc > 2) {
        period = strtol(argv[2], &end, 10);
        if (!*argv[2] || *end || period < 1000 || period > 1000000) {
            fprintf(stderr, "Invalid period\n");
            return 1;
        }
    }

    // Get one of biggest priority for tim precision
    struct sched_param tSp;
    tSp.sched_priority = 90;
    if (sched_setscheduler(0, SCHED_RR, &tSp) < 0) 
        fprintf(stderr, "Warning: Unable to set Scheduler: %s (Are you root?)\n", strerror(errno));

    consume(objective, period);
    return 0;
}