#include "step.hpp"
#include "pico/stdio.h"
#include "hardware/pio.h"
#include "hardware/irq.h"
#include <pico/time.h>

#include <math.h>

TMC2209_step_dual * g_step_driver_instance;

void step_irq0_handler() {
    g_step_driver_instance->irq_handler0();
};

void step_irq1_handler() {
    g_step_driver_instance->irq_handler1();
};

TMC2209_step_dual::TMC2209_step_dual(
    uint step0_pin, 
    uint step1_pin, 
    uint dir0_pin, 
    uint dir1_pin, 
    uint enable0_pin, 
    uint enable1_pin, 
    PIO pio_core
) {
    this->step0_pin = step0_pin;
    this->step1_pin = step1_pin;
    this->dir0_pin = dir0_pin;
    this->dir1_pin = dir1_pin;
    this->enable0_pin = enable0_pin;
    this->enable1_pin = enable1_pin;
    this->pio_core = pio_core;

    g_step_driver_instance = this;

    this->done_flag0 = false;
    this->done_flag1 = false;
    this->done_flag0_arm = false;
    this->done_flag1_arm = false;

    init_gpio();
    init_pio();
};

void TMC2209_step_dual::init_gpio() {
    gpio_init(this->enable0_pin);
    gpio_init(this->enable1_pin);
    gpio_set_dir(this->enable0_pin, GPIO_OUT);
    gpio_set_dir(this->enable1_pin, GPIO_OUT);
    gpio_put(this->enable0_pin, 0);
    gpio_put(this->enable1_pin, 0);
    sleep_ms(1);

    gpio_init(this->dir0_pin);
    gpio_init(this->dir1_pin);
    gpio_set_dir(this->dir0_pin, GPIO_OUT);
    gpio_set_dir(this->dir1_pin, GPIO_OUT);    
    gpio_put(this->dir0_pin, 1);
    gpio_put(this->dir1_pin, 1);
    sleep_ms(1);
};

void TMC2209_step_dual::init_pio() {
    pio_sm_set_enabled(this->pio_core, this->sm_num_divider0, false);
    pio_sm_set_enabled(this->pio_core, this->sm_num_divider1, false);
    pio_sm_set_enabled(this->pio_core, this->sm_num_counter0, false);
    pio_sm_set_enabled(this->pio_core, this->sm_num_counter1, false);

    this->pulser_program_offset = pio_add_program(this->pio_core, &pulser_program);

    this->sm_config_divider0 = pulser_program_get_default_config(this->pulser_program_offset);
    this->sm_config_divider1 = pulser_program_get_default_config(this->pulser_program_offset);
    this->sm_config_counter0 = pulser_program_get_default_config(this->pulser_program_offset);
    this->sm_config_counter1 = pulser_program_get_default_config(this->pulser_program_offset);

    pio_gpio_init(this->pio_core, this->step0_pin);
    pio_gpio_init(this->pio_core, this->step1_pin);

    pio_sm_set_consecutive_pindirs(this->pio_core, this->sm_num_counter0, this->step0_pin, 1, true);
    pio_sm_set_consecutive_pindirs(this->pio_core, this->sm_num_counter1, this->step1_pin, 1, true);

    sm_config_set_set_pins(&this->sm_config_counter0, this->step0_pin, 1);
    sm_config_set_set_pins(&this->sm_config_counter1, this->step1_pin, 1);

    pio_sm_init(
        this->pio_core, 
        this->sm_num_divider0, 
        this->pulser_program_offset + this->subprogram_offset_divider, 
        &this->sm_config_divider0
    );
    pio_sm_init(
        this->pio_core, 
        this->sm_num_divider1, 
        this->pulser_program_offset + this->subprogram_offset_divider, 
        &this->sm_config_divider1
    );
    pio_sm_init(
        this->pio_core, 
        this->sm_num_counter0, 
        this->pulser_program_offset + this->subprogram_offset_counter, 
        &this->sm_config_counter0
    );
    pio_sm_init(
        this->pio_core, 
        this->sm_num_counter1, 
        this->pulser_program_offset + this->subprogram_offset_counter, 
        &this->sm_config_counter1
    );

    float clock_div = ((float)(SYS_CLK_KHZ))/(base_clock_freq_khz);

    pio_sm_set_clkdiv(this->pio_core, this->sm_num_divider0, clock_div);
    pio_sm_set_clkdiv(this->pio_core, this->sm_num_divider1, clock_div);
    pio_sm_set_clkdiv(this->pio_core, this->sm_num_counter0, 1.0);
    pio_sm_set_clkdiv(this->pio_core, this->sm_num_counter1, 1.0);

    pio_sm_set_enabled(this->pio_core, this->sm_num_divider0, true);
    pio_sm_set_enabled(this->pio_core, this->sm_num_divider1, true);
    pio_sm_set_enabled(this->pio_core, this->sm_num_counter0, true);
    pio_sm_set_enabled(this->pio_core, this->sm_num_counter1, true);

    pio_set_irq0_source_enabled(this->pio_core, pis_interrupt2, true);
    pio_set_irq1_source_enabled(this->pio_core, pis_interrupt3, true);

    if(this->pio_core == pio0) {
        irq_set_exclusive_handler(PIO0_IRQ_0, step_irq0_handler);
        irq_set_exclusive_handler(PIO0_IRQ_1, step_irq1_handler);

        irq_set_enabled(PIO0_IRQ_0, true);
        irq_set_enabled(PIO0_IRQ_1, true);
    } else if(this->pio_core == pio1) {
        irq_set_exclusive_handler(PIO1_IRQ_0, step_irq0_handler);
        irq_set_exclusive_handler(PIO1_IRQ_1, step_irq1_handler);

        irq_set_enabled(PIO1_IRQ_0, true);
        irq_set_enabled(PIO1_IRQ_1, true);
    }

    set_conf0(1000, 0);
    set_conf1(1000, 0);
};

void TMC2209_step_dual::set_conf0(uint frequency, bool dir) {
    uint32_t cycles = (uint32_t)(base_clock_freq_khz / (((float)(frequency))/1000.0)) - 6;
    
    gpio_put(this->dir0_pin, dir);
    pio_sm_put_blocking(
        this->pio_core, 
        this->sm_num_divider0, 
        cycles
    );
};

void TMC2209_step_dual::set_conf1(uint frequency, bool dir) {
    uint32_t cycles = (uint32_t)(base_clock_freq_khz / (((float)(frequency))/1000.0)) - 6;

    gpio_put(this->dir1_pin, dir);
    pio_sm_put_blocking(
        this->pio_core, 
        this->sm_num_divider1, 
        cycles
    );
};

void TMC2209_step_dual::pulse0(uint n) {
    this->remaining_step_count0 = n;
    this->done_flag0_arm = true;
    this->done_flag0 = false;

    pio_sm_put_blocking(
        this->pio_core, this->sm_num_counter0, n
    );
};

void TMC2209_step_dual::pulse1(uint n) {
    this->remaining_step_count1 = n;
    this->done_flag1_arm = true;
    this->done_flag1 = false;

    pio_sm_put_blocking(
        this->pio_core, this->sm_num_counter1, n
    );
};

void TMC2209_step_dual::pulse0_int(int n) {
    this->remaining_step_count0 = n;
    this->done_flag0_arm = true;
    if(n < 0) {
        n = abs(n);
        gpio_put(this->dir1_pin, gpio_get(this->dir1_pin) ^ 1);
    }

    pio_sm_put_blocking(
        this->pio_core, this->sm_num_counter0, n
    );
};

void TMC2209_step_dual::pulse1_int(int n) {
    this->remaining_step_count1 = n;
    this->done_flag1_arm = true;
    if(n < 0) {
        n = abs(n);
        gpio_put(this->dir1_pin, gpio_get(this->dir1_pin) ^ 1);
    }

    pio_sm_put_blocking(
        this->pio_core, this->sm_num_counter1, n
    );
};

void TMC2209_step_dual::irq_handler0() {
    this->remaining_step_count0 = pio_sm_get(this->pio_core, this->sm_num_counter0);

    if(this->done_flag0_arm) {
        done_flag0 = true;
    } else {
        done_flag0 = false;
    }
    
    pio_interrupt_clear(this->pio_core, 2);
}

void TMC2209_step_dual::irq_handler1() {
    this->remaining_step_count1 = pio_sm_get(this->pio_core, this->sm_num_counter1);

    if(this->done_flag1_arm) {
        done_flag1 = true;
    } else {
        done_flag1 = false;
    }
    pio_interrupt_clear(this->pio_core, 3);
}

bool TMC2209_step_dual::wait0(int timeout_ms) {
    int t = timeout_ms*1000;
    bool continues = timeout_ms <= 0;

    if(!this->done_flag0_arm)
        return true;

    while(t > 0 || continues) {

        if(done_flag0) {
            this->done_flag0_arm = false;
            done_flag0 = false;
            return true;
        }

        if(!continues)
            t -= sleep_time_us;

        sleep_us(sleep_time_us);
    }

    return false;
};

bool TMC2209_step_dual::wait1(int timeout_ms) {
    int t = timeout_ms*1000;
    bool continues = timeout_ms <= 0;

    if(!this->done_flag1_arm)
        return true;

    while(t > 0 || continues) {

        if(done_flag1) {
            this->done_flag1_arm = false;
            done_flag1 = false;
            return true;
        }

        if(!continues)
            t -= sleep_time_us;

        sleep_us(sleep_time_us);
    }

    return false;
};

uint TMC2209_step_dual::get_remaining_steps0() {
    return this->remaining_step_count0;
};

uint TMC2209_step_dual::get_remaining_steps1() {
    return this->remaining_step_count1;
};