#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <math.h>

#include "ImC/imc_kernel.h"
#include "ImC/imc_extension.h"
#include "ImC/imc_api.h"

extern ADC_HandleTypeDef ADC_HANDLER_SBC;

float loop_pass_estimation[imcMAX_LOOP_IDS];
int imc_is_passing_loops_backup IMC_KERNEL_NVM;
int imc_is_passing_loops;
int imc_energy_at_start;

float estimation_backup[2][imcMAX_LOOP_IDS] IMC_KERNEL_NVM;
int imc_energy_estimation_initialized IMC_KERNEL_NVM;
int latest_buffer_index IMC_KERNEL_NVM;
int last_loop_skip_count;
int pass_due_to_low_energy;

#define imcLOOP_OPT_DEBUG_VERBOSE 0

#if(imcPRINT_STATS)
    int imc_checkpoint_triggered;
    int imc_checkpoint_executed;
    int imc_checkpoint_triggered_backup[2] IMC_KERNEL_NVM;
    int imc_checkpoint_executed_backup[2] IMC_KERNEL_NVM;
    int stat_buffer_index IMC_KERNEL_NVM;

    int first_checkpoint IMC_KERNEL_NVM;
    int no_forward_progress;
    int imc_no_forward_progress_backup[2] IMC_KERNEL_NVM;

    void imc_recover_stats() {
        imc_checkpoint_triggered = imc_checkpoint_triggered_backup[stat_buffer_index];
        imc_checkpoint_executed = imc_checkpoint_executed_backup[stat_buffer_index];
        no_forward_progress = imc_no_forward_progress_backup[stat_buffer_index];
    }

    void imc_backup_stats() {
        int nextBufferIndex = stat_buffer_index == 0 ? 1: 0;
        imc_checkpoint_triggered_backup[nextBufferIndex] = imc_checkpoint_triggered;
        imc_checkpoint_executed_backup[nextBufferIndex] = imc_checkpoint_executed;
        imc_no_forward_progress_backup[nextBufferIndex] = no_forward_progress;
        stat_buffer_index = nextBufferIndex;
    }
#endif

#if (imcPRINT_LATENCY_OVERHEAD)
    unsigned int imc_ticks_init;
    unsigned int imc_ticks_getCSC;
    unsigned int imc_ticks_adjust;
    unsigned int imc_ticks_total;
    int imc_latency_overhead_buffer_index IMC_KERNEL_NVM;

    unsigned int imc_last_tick;

    unsigned int imc_ticks_init_backup[2] IMC_KERNEL_NVM;
    unsigned int imc_ticks_getCSC_backup[2] IMC_KERNEL_NVM;
    unsigned int imc_ticks_adjust_backup[2] IMC_KERNEL_NVM;
    unsigned int imc_ticks_total_backup[2] IMC_KERNEL_NVM;

    void imc_recover_latency_overhead() {
        int index = imc_latency_overhead_buffer_index;
        imc_ticks_init = imc_ticks_init_backup[index];
        imc_ticks_getCSC = imc_ticks_getCSC_backup[index];
        imc_ticks_adjust = imc_ticks_adjust_backup[index];
        imc_ticks_total = imc_ticks_total_backup[index];
    }

    void imc_backup_latency_overhead() {
        unsigned int tick = DWT->CYCCNT;
        imc_ticks_total += tick - imc_last_tick;
        imc_last_tick = tick;

        int next_buffer_index = imc_latency_overhead_buffer_index == 0 ? 1 : 0;
        imc_ticks_init_backup[next_buffer_index] = imc_ticks_init;
        imc_ticks_getCSC_backup[next_buffer_index] = imc_ticks_getCSC;
        imc_ticks_adjust_backup[next_buffer_index] = imc_ticks_adjust;
        imc_ticks_total_backup[next_buffer_index] = imc_ticks_total;
        imc_latency_overhead_buffer_index = next_buffer_index;
    }
#endif

    void imc_recover_estimations()
    {
    #if(imcLOOP_OPT_DEBUG && imcLOOP_OPT_DEBUG_VERBOSE)
        portDISABLE_INTERRUPTS();
        printf("(DEBUG) recover estimation table from buffer %d\r\n", latest_buffer_index);
        portENABLE_INTERRUPTS();
    #endif
    memcpy(loop_pass_estimation, estimation_backup[latest_buffer_index], sizeof(float) * imcMAX_LOOP_IDS);
    __DSB();
}

void imc_backup_estimations() {
    int nextBufferIndex = latest_buffer_index == 0 ? 1 : 0;
    #if(imcLOOP_OPT_DEBUG && imcLOOP_OPT_DEBUG_VERBOSE)
        portDISABLE_INTERRUPTS();
        printf("(DEBUG) backup estimation table to buffer %d\r\n", nextBufferIndex);
        portENABLE_INTERRUPTS();
    #endif
    memcpy(estimation_backup[nextBufferIndex], loop_pass_estimation, sizeof(float) * imcMAX_LOOP_IDS);
    latest_buffer_index = nextBufferIndex;
    __DSB();
}

void imc_print_estimations() {
    portDISABLE_INTERRUPTS();
    for(int i=0; i<imcMAX_LOOP_IDS; i++) {
        printf("(DEBUG) loop_#%d: %f\r\n", i, loop_pass_estimation[i]);
    }
    portENABLE_INTERRUPTS();
}

void imc_init_energy_estimation() {
    #if (imcPRINT_LATENCY_OVERHEAD)
        unsigned int start_tick = DWT->CYCCNT;
    #endif

    if(!imc_energy_estimation_initialized) {
        #if(imcLOOP_OPT_DEBUG)
            portDISABLE_INTERRUPTS();
            printf("(DEBUG) initialize estimation table\r\n");
            portENABLE_INTERRUPTS();
        #endif
        for(int i=0; i<imcMAX_LOOP_IDS; i++) {
            loop_pass_estimation[i] = 3;
        }
        latest_buffer_index = 0;
        #if(imcPRINT_STATS)
            stat_buffer_index = 0;
        #endif
        imc_backup_estimations();
    } else {
        imc_recover_estimations();
    }

    #if (imcPRINT_LATENCY_OVERHEAD)
        unsigned int end_tick = DWT->CYCCNT;
        imc_ticks_init += end_tick - start_tick;
    #endif

    if(imc_energy_estimation_initialized && imc_is_passing_loops_backup != -1) {
        int loop_id = imc_is_passing_loops_backup;
        float new_estimation = loop_pass_estimation[loop_id] / 3;
        // int new_estimation = loop_pass_estimation[loop_id];
        // if(new_estimation >= 20) new_estimation /= 10;
        // else if(new_estimation < 20 && new_estimation > 2) new_estimation = 2;
        // else if(new_estimation == 2) new_estimation = 1;
        // else new_estimation = 0;

        loop_pass_estimation[loop_id] = new_estimation;
        estimation_backup[latest_buffer_index][loop_id] = new_estimation;
        #if (imcLOOP_OPT_DEBUG)
            portDISABLE_INTERRUPTS();
            printf("(DEBUG) (recovery) estimation for loop %d is updated to %f\r\n", loop_id, loop_pass_estimation[loop_id]);
            portENABLE_INTERRUPTS();
        #endif
    }
    imc_energy_estimation_initialized = 1;
    imc_is_passing_loops_backup = -1;
    imc_is_passing_loops = -1;
    __DSB();
}


int get_current_energy() {
    int volt = measure_voltage(ADC_HANDLER_SBC, EPS_CAP_ID_SBC);
    int energy = volt*volt - imcCAP_VOL_LOW*imcCAP_VOL_LOW;
    return energy;
}


void update_estimation(int loop_id, int counter, int current_energy, int energy_at_start) {
    float current_estimation = (float)loop_pass_estimation[loop_id];
    if(current_estimation > 1e6) return;
    const float error_margin = imcENERGY_TOTAL * 0.03;
    // const float error_margin = imcENERGY_TOTAL * 0;
    float max_consumed = energy_at_start + error_margin - (current_energy - error_margin);
    float min_consumed = energy_at_start - error_margin - (current_energy + error_margin);
    // if(max_consumed <= 0) max_consumed = 1e-3;
    // if(min_consumed <= 0) {
    //     min_consumed = 1e-3;
    // }
    if(max_consumed <= 0) max_consumed = 1e-3;
    if(min_consumed <= 0) {
        min_consumed = 1;
    }

    float measured_min = ((float)(imcENERGY_TOTAL)) / max_consumed * counter;
    // measured_min *= 2.5;
    float measured_max = ((float)(imcENERGY_TOTAL)) / min_consumed * counter;
    // measured_max *= 2;

    // float measured_iterations = ((float)(imcENERGY_TOTAL)) / energy_consumed * counter;
    // float variance_margin = 1;
    // float outlier_factor = 5;
    // int too_large_measured = measured_iterations > outlier_factor * current_estimation;
    // int too_small_measured = outlier_factor * measured_iterations < current_estimation;

    // printf("(DEBUG) lood_id:%d, counter:%d, consumed: %d, measured: %f, current: %f\r\n", loop_id, counter, energy_consumed, measured_iterations, current_estimation);

    // if(too_large_measured || too_small_measured) return;

    float new_estimation;

    if(measured_max < current_estimation) {
        float factor = measured_max / current_estimation;
        float factor_max = 0.25;
        factor = factor < factor_max ? factor_max : factor;
        new_estimation = current_estimation * factor;
        loop_pass_estimation[loop_id] = new_estimation;
    } else if (measured_min > current_estimation) {
        float factor = measured_min / current_estimation;
        float factor_max = 3;
        factor = factor > factor_max ? factor_max : factor;
        new_estimation = current_estimation * factor;
        loop_pass_estimation[loop_id] = new_estimation;
    }
    else {
        #if(imcLOOP_OPT_DEBUG)
            portDISABLE_INTERRUPTS();
            printf("(DEBUG) Estimation not updated; loop_id:%d, counter:%d, measured_min: %f, measured_max: %f, current_estimation: %f\r\n", 
                loop_id, counter, measured_min, measured_max, current_estimation);
            portENABLE_INTERRUPTS();
        #endif
        return;
    }

    const int max_estimation = 1000000;
    if(new_estimation > max_estimation) {
        new_estimation = max_estimation;
        loop_pass_estimation[loop_id] = max_estimation;
    }

    #if(imcLOOP_OPT_DEBUG)
    portDISABLE_INTERRUPTS();
    printf("(DEBUG) estimation for loop %d is updated to %f\r\n", loop_id, new_estimation);
    portENABLE_INTERRUPTS();

    #endif

    imc_backup_estimations();
}


void __imc_unflag_loop_skipping(int loop_id, int counter)
{
    if(loop_id != imc_is_passing_loops) return;

    #if (imcPRINT_LATENCY_OVERHEAD)
        unsigned int start_tick = DWT->CYCCNT;
    #endif

    int energy = get_current_energy();
    int energy_consumed = imc_energy_at_start - energy;
    // printf("(DEBUG) energy consumed: %d\r\n", energy_consumed);

    update_estimation(loop_id, counter, energy, imc_energy_at_start);
    // float measured_iterations = ((float)(imcENERGY_TOTAL)) / energy_consumed * counter;
    // float current_estimation = (float)loop_pass_estimation[loop_id];

    // printf("(DEBUG) loop_id:%d, counter:%d, measured: %f, estimation: %f\r\n", loop_id, counter, measured_iterations, current_estimation);

    imc_is_passing_loops_backup = -1;
    imc_is_passing_loops = -1;
    __DSB();

    #if (imcPRINT_LATENCY_OVERHEAD)
        unsigned int end_tick = DWT->CYCCNT;
        imc_ticks_adjust += end_tick - start_tick;
    #endif
}


void __imc_finish_loop_skipping() {
    int loop_id = imc_is_passing_loops;

    #if (imcPRINT_LATENCY_OVERHEAD)
        unsigned int start_tick = DWT->CYCCNT;
    #endif

    if(loop_id == 0) {
        printf("(DEBUG) imc_is_passing_loop is 0\r\n");
        while(1) { __ASM("  nop"); }
    }
    if(loop_id > 0) {
        // int volt = measure_voltage(ADC_HANDLER_SBC, EPS_CAP_ID_SBC);
        // int energy = volt*volt - imcCAP_VOL_LOW*imcCAP_VOL_LOW;
        #if(imcLOOP_OPT_DEBUG)
            portDISABLE_INTERRUPTS();
            // printf("(DEBUG) loop %d is not perfectly skipped\r\n", loop_id);
            portENABLE_INTERRUPTS();
        #endif
        int energy = get_current_energy();
        int energy_consumed = imc_energy_at_start - energy;

        int last_skip_is_zero = last_loop_skip_count == 0;
        if(!last_skip_is_zero) {
            update_estimation(loop_id, last_loop_skip_count, energy, imc_energy_at_start);
        }
    }
    imc_is_passing_loops_backup = -1;
    imc_is_passing_loops = -1;
    __DSB();

    #if (imcPRINT_LATENCY_OVERHEAD)
        unsigned int end_tick = DWT->CYCCNT;
        imc_ticks_adjust += end_tick - start_tick;
    #endif
}

int __imc_get_loop_pass_count(int loop_id) {
    #if(imcENABLE_STATIC_LOOP_PASS_COUNT)
        return imcLOOP_PASS_COUNT;
    #elif(imcENABLE_ADAPTIVE_LOOP_PASS_COUNT)
        if(imc_is_passing_loops != -1) return INT32_MAX;
        #if (imcPRINT_LATENCY_OVERHEAD)
            unsigned int start_tick = DWT->CYCCNT;
        #endif

        int volt = measure_voltage(ADC_HANDLER_SBC, EPS_CAP_ID_SBC);
        int energy = volt*volt - imcCAP_VOL_LOW*imcCAP_VOL_LOW;
        int ratio = energy * 100 / imcENERGY_TOTAL;

        imc_is_passing_loops = loop_id;
        imc_is_passing_loops_backup = loop_id;
        __DSB();
        imc_energy_at_start = energy;

        #if(imcLOOP_OPT_DEBUG && imcLOOP_OPT_DEBUG_VERBOSE)
            portDISABLE_INTERRUPTS();
            printf("(DEBUG) loopId: %d, ratio: %d\r\n", loop_id, ratio);
            portENABLE_INTERRUPTS();
        #endif

        float current_estimation = loop_pass_estimation[loop_id];
        if(ratio > 100) {
            last_loop_skip_count = current_estimation;
        }
        else if(ratio < 5) {
            last_loop_skip_count = 0;
            imc_is_passing_loops = -1;
            imc_is_passing_loops_backup = -1;
            // imc_sbc_power_off();
            #if(imcLOOP_OPT_DEBUG && imcLOOP_OPT_DEBUG_VERBOSE)
                portDISABLE_INTERRUPTS();
                printf("(DEBUG) ratio: %d\r\n", ratio);
                portENABLE_INTERRUPTS();
            #endif
            // while(1) { __ASM(" nop"); }
            // printf("__imc_get_loop_pass_count(%d) called; %d returned\r\n", loop_id, last_loop_skip_count);
        }
        else {
            double r = (double)(ratio) / 100.0;
            double target_skip_count = current_estimation * r;
            #if(imcLOOP_OPT_DEBUG && imcLOOP_OPT_DEBUG_VERBOSE)
                portDISABLE_INTERRUPTS();
                printf("(DEBUG) loop_id:%d, current estimation: %f, target_skip_count: %f, r: %f\r\n", loop_id, current_estimation, target_skip_count, r);
                portENABLE_INTERRUPTS();
            #endif
            const int max_skip_count = 10000;
            // if(target_skip_count > max_skip_count) {
            //     last_loop_skip_count = max_skip_count;
            // } else {
            //     last_loop_skip_count = (int)target_skip_count;
            // }
            if (isnan(target_skip_count) || target_skip_count < 1.0) { 
                last_loop_skip_count = 0;
            } else if (target_skip_count > max_skip_count) {
                last_loop_skip_count = max_skip_count;
            } else {
                last_loop_skip_count = (int)target_skip_count;
            }
            // last_loop_skip_count = loop_pass_estimation[loop_id] * ((double)(ratio) / 100.0);
        }
        // printf("loopId: %d, volt: %d, energy: %d, total: %d, ratio: %d, est:%f, count:%d\r\n", loop_id, volt, energy, imcENERGY_TOTAL, ratio, current_estimation, last_loop_skip_count);
        #if (imcPRINT_LATENCY_OVERHEAD)
            unsigned int end_tick = DWT->CYCCNT;
            imc_ticks_getCSC += end_tick - start_tick;
        #endif
        return last_loop_skip_count;
    #else
        return 0;
    #endif
}

void ckptTask() {
    while(1) {
    int volt = measure_voltage(ADC_HANDLER_SBC, 0);
        // __asm(" svc 7");
    if(volt < 3660) {
        HAL_GPIO_WritePin(GPIOF, GPIO_PIN_11, GPIO_PIN_SET);
        __asm(" svc 7");
        imc_sbc_power_off();
        while(1) __asm(" nop");
    }
    // printf("volt: %d\r\n", volt);
    const TickType_t xDelay = 25 / portTICK_PERIOD_MS;
    osDelay(xDelay);
    }
}