#include <stdio.h>
#include <string.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)
        printf("(DEBUG) recover estimation table from buffer %d\r\n", latest_buffer_index);
    #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)
        printf("(DEBUG) backup estimation table to buffer %d\r\n", nextBufferIndex);
    #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)
            printf("(DEBUG) initialize estimation table\r\n");
        #endif
        for(int i=0; i<imcMAX_LOOP_IDS; i++) {
            loop_pass_estimation[i] = 2;
        }
        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;
    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;
    }

    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 return;

    #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)
            printf("(DEBUG) loop %d is not perfectly skipped\r\n", loop_id);
        #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);
        }

        /*
        float possible_iterations_estimated = 0;
        int underestimated = 0;
        if(!last_skip_is_zero) {
            float energy_per_iteration = (float)energy_consumed / last_loop_skip_count;
            possible_iterations_estimated = (float)energy / energy_per_iteration;
            underestimated = possible_iterations_estimated > 3;
        }
        if(!last_skip_is_zero && underestimated) {
            int current_estimation = loop_pass_estimation[loop_id];
            // double inc_factor = 2.5;
            double inc_factor = possible_iterations_estimated > 3 ? 3 : possible_iterations_estimated;
            if (current_estimation * inc_factor > (double)INT32_MAX)
            {
                current_estimation = INT32_MAX/2-1;
            }
            else
            {
                current_estimation = (int)(current_estimation * inc_factor);
            }
            
            if (current_estimation < 1) {
                current_estimation = 1;
            }
            loop_pass_estimation[loop_id] = current_estimation;
            #if (imcLOOP_OPT_DEBUG)
                portDISABLE_INTERRUPTS();
                printf("(DEBUG) estimation for loop %d is updated to %d\r\n", loop_id, current_estimation);
                portENABLE_INTERRUPTS();
            #endif
        imc_backup_estimations();
        } else {
            #if (imcLOOP_OPT_DEBUG)
                portDISABLE_INTERRUPTS();
                printf("(DEBUG) estimation for loop %d is not updated (%d)\r\n", loop_id, loop_pass_estimation[loop_id]);
                portENABLE_INTERRUPTS();
            #endif
        }
        */
    }
    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)
            printf("(DEBUG) ratio: %d\r\n", ratio);
        #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)
                printf("(DEBUG) ratio: %d\r\n", ratio);
            #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;
            last_loop_skip_count = (int)((double)current_estimation * r);
            // 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:%d, 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);
    }
}