#include <pico/stdlib.h>
#include <pico/stdio_usb.h>
#include <tusb.h>

#include <hardware/pwm.h>
#include <hardware/pio.h>
#include <hardware/dma.h>
#include <hardware/clocks.h>

#include <stdio.h>

#define SMPS_PIN 23

#define IRP_FB_PIN 2
#define IRP_RX_PIN 3
#define IRN_FB_PIN 0
#define IRN_RX_PIN 1

#define IR_PIO pio0
#define IRP_SM 0
#define IRN_SM 1

#define NUM_SAMPLES_BITS 10
#define NUM_SAMPLES (1 << NUM_SAMPLES_BITS)
static uint16_t irp_buffer[NUM_SAMPLES] __aligned(2 * NUM_SAMPLES);
static int dma_ch_irp_rx;

static uint16_t irn_buffer[NUM_SAMPLES] __aligned(2 * NUM_SAMPLES);
static int dma_ch_irn_rx;

#define BIT(n) (1 << (n))

static void init_pio_sm(pio_hw_t *pio, int sm, int rx_pin, int fb_pin)
{
	const uint16_t insn[] = {
		pio_encode_jmp_x_dec(4),
		pio_encode_in(pio_y, 16) | pio_encode_sideset_opt(1, 1),
		pio_encode_mov(pio_x, pio_osr),
		pio_encode_nop() | pio_encode_sideset_opt(1, 0),
		pio_encode_jmp_pin(6),
		pio_encode_jmp(0) | pio_encode_sideset_opt(1, 1),
		pio_encode_jmp_y_dec(0) | pio_encode_sideset_opt(1, 0),
	};

	pio_program_t prog = {
		.instructions = insn,
		.length = sizeof(insn) / sizeof(*insn),
		.origin = 0,
	};

	if (pio_can_add_program(pio, &prog))
		pio_add_program(pio, &prog);

	pio_sm_config pc = pio_get_default_sm_config();
	sm_config_set_clkdiv_int_frac8(&pc, 1, 0);
	sm_config_set_jmp_pin(&pc, rx_pin);
	sm_config_set_in_pins(&pc, rx_pin);
	sm_config_set_out_pins(&pc, fb_pin, 1);
	sm_config_set_set_pins(&pc, fb_pin, 1);
	sm_config_set_sideset_pin_base(&pc, fb_pin);
	sm_config_set_sideset(&pc, 2, true, false);
	sm_config_set_in_shift(&pc, false, true, 16);
	sm_config_set_out_shift(&pc, false, false, 32);

	pio_sm_init(pio, sm, prog.origin, &pc);
	pio_sm_set_wrap(pio, sm, prog.origin, prog.origin + prog.length - 1);

	pio_sm_restart(pio, sm);
	pio_sm_clear_fifos(pio, sm);

	pio->txf[sm] = 126;
	pio_sm_exec_wait_blocking(pio, sm, pio_encode_pull(false, false));
	pio_sm_exec_wait_blocking(pio, sm, pio_encode_mov(pio_x, pio_osr));

	pio_sm_set_consecutive_pindirs(pio, sm, rx_pin, 1, GPIO_IN);
	pio_sm_set_consecutive_pindirs(pio, sm, fb_pin, 1, GPIO_OUT);

	gpio_disable_pulls(rx_pin);
	gpio_disable_pulls(fb_pin);

	pio_gpio_init(pio, fb_pin);

	gpio_set_input_hysteresis_enabled(rx_pin, false);
	gpio_set_drive_strength(fb_pin, GPIO_DRIVE_STRENGTH_2MA);
	gpio_set_slew_rate(fb_pin, GPIO_SLEW_RATE_FAST);
}

static unsigned tail;

static bool read_sample(int *sample)
{
	unsigned head = (dma_hw->ch[dma_ch_irn_rx].write_addr >> 1) % NUM_SAMPLES;

	if (tail == head)
		return false;

	uint16_t next_pos = irp_buffer[tail];
	uint16_t next_neg = irn_buffer[tail];
	tail = (tail + 1) % NUM_SAMPLES;

	static uint16_t prev_pos;
	int pos = prev_pos - next_pos;
	prev_pos = next_pos;

	static uint16_t prev_neg;
	int neg = prev_neg - next_neg;
	prev_neg = next_neg;

	static int dc_pos;
	dc_pos += ((pos << 16) - dc_pos) >> 12;
	pos -= (dc_pos >> 16);

	static int dc_neg;
	dc_neg += ((neg << 16) - dc_neg) >> 12;
	neg -= (dc_neg >> 16);

	*sample = pos - neg;
	return true;
}

int main()
{
	stdio_usb_init();

	for (int i = 0; i < 30; i++) {
		if (stdio_usb_connected())
			break;

		sleep_ms(100);
	}

	printf("\nWelcome to PDM Test!\n");

	gpio_init(SMPS_PIN);
	gpio_disable_pulls(SMPS_PIN);
	gpio_put(SMPS_PIN, 1);
	gpio_set_dir(SMPS_PIN, GPIO_OUT);

	init_pio_sm(IR_PIO, IRP_SM, IRP_RX_PIN, IRP_FB_PIN);
	init_pio_sm(IR_PIO, IRN_SM, IRN_RX_PIN, IRN_FB_PIN);

	dma_ch_irp_rx = dma_claim_unused_channel(true);
	dma_ch_irn_rx = dma_claim_unused_channel(true);

	dma_channel_config dma_conf;

	dma_conf = dma_channel_get_default_config(dma_ch_irp_rx);
	channel_config_set_transfer_data_size(&dma_conf, DMA_SIZE_16);
	channel_config_set_read_increment(&dma_conf, false);
	channel_config_set_write_increment(&dma_conf, true);
	channel_config_set_ring(&dma_conf, true, NUM_SAMPLES_BITS + 1);
	channel_config_set_dreq(&dma_conf, pio_get_dreq(IR_PIO, IRP_SM, false));
	channel_config_set_chain_to(&dma_conf, dma_ch_irn_rx);
	dma_channel_configure(dma_ch_irp_rx, &dma_conf, irp_buffer, &IR_PIO->rxf[IRP_SM], 1, false);

	dma_conf = dma_channel_get_default_config(dma_ch_irn_rx);
	channel_config_set_transfer_data_size(&dma_conf, DMA_SIZE_16);
	channel_config_set_read_increment(&dma_conf, false);
	channel_config_set_write_increment(&dma_conf, true);
	channel_config_set_ring(&dma_conf, true, NUM_SAMPLES_BITS + 1);
	channel_config_set_dreq(&dma_conf, pio_get_dreq(IR_PIO, IRN_SM, false));
	channel_config_set_chain_to(&dma_conf, dma_ch_irp_rx);
	dma_channel_configure(dma_ch_irn_rx, &dma_conf, irn_buffer, &IR_PIO->rxf[IRN_SM], 1, false);

	dma_channel_start(dma_ch_irp_rx);

	pio_set_sm_mask_enabled(IR_PIO, BIT(IRP_SM) | BIT(IRN_SM), true);

	static int8_t batch[64];
	unsigned written = 0;

	stdio_usb_lock();

	while (true) {
		int sample;

		while (read_sample(&sample)) {
			batch[written++] = sample;

			if (written >= sizeof(batch) / sizeof(*batch)) {
				while (tud_cdc_n_write_available(1) < sizeof(batch)) {
					tud_cdc_n_write_flush(1);
					tud_task();
					sleep_us(10);
				}

				tud_cdc_n_write(1, batch, sizeof(batch));
				written = 0;
			}
		}

		tud_task();
		sleep_us(10);
	}

	stdio_usb_unlock();
}