#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 IR_BIAS_PIN 3
#define IR_RX_PIN 6
#define IR_FB_PIN 8

#define IR_PIO pio0
#define IR_SM 0

#define HP_GND_PIN 13
#define HP_LEFT_PIN 14
#define HP_RIGHT_PIN 15

#define HP_GND_SLICE 6
#define HP_GND_CHAN PWM_CHAN_B

#define HP_LR_SLICE 7

#define HP_PWM_RATE (16 * 48 * 1000)
#define HP_PWM_SCALE (SYS_CLK_HZ / HP_PWM_RATE)
#define HP_PWM_MID (HP_PWM_SCALE / 2)

typedef struct Sample {
	uint16_t x;
	uint16_t y;
} Sample;

#define NUM_SAMPLES_BITS 10
#define NUM_SAMPLES (1 << NUM_SAMPLES_BITS)
static Sample sample_buffer[NUM_SAMPLES] __aligned(4 * NUM_SAMPLES);
static int dma_ch_sample;
static int dma_ch_rx;

/* Jump to 0, side-set 1. */
static uint16_t sampling_script[] = { 0b0001100000000000 };

static void init_pio_sm(void)
{
	const uint16_t insn[] = {
		pio_encode_in(pio_y, 16),
		pio_encode_in(pio_x, 16) | pio_encode_sideset_opt(1, 0),
		pio_encode_jmp_pin(4),
		pio_encode_jmp_x_dec(2) | pio_encode_sideset_opt(1, 1),
		pio_encode_jmp_y_dec(2) | pio_encode_sideset_opt(1, 0),
	};

	pio_program_t prog = {
		.instructions = insn,
		.length = sizeof(insn) / sizeof(*insn),
		.origin = 0,
	};
	pio_add_program(IR_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, IR_RX_PIN);
	sm_config_set_in_pins(&pc, IR_RX_PIN);
	sm_config_set_out_pins(&pc, IR_FB_PIN, 1);
	sm_config_set_set_pins(&pc, IR_FB_PIN, 1);
	sm_config_set_sideset_pin_base(&pc, IR_FB_PIN);
	sm_config_set_sideset(&pc, 2, true, false);
	sm_config_set_in_shift(&pc, false, true, 32);
	sm_config_set_fifo_join(&pc, PIO_FIFO_JOIN_RX);

	pio_sm_init(IR_PIO, IR_SM, prog.origin, &pc);
	pio_sm_set_wrap(IR_PIO, IR_SM, prog.origin + 2, prog.origin + prog.length - 1);

	pio_sm_set_consecutive_pindirs(IR_PIO, IR_SM, IR_RX_PIN, 1, GPIO_IN);
	pio_sm_set_consecutive_pindirs(IR_PIO, IR_SM, IR_FB_PIN, 1, GPIO_OUT);

	pio_sm_restart(IR_PIO, IR_SM);
	pio_sm_clear_fifos(IR_PIO, IR_SM);
}

static unsigned tail;

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

	if (tail == head)
		return false;

	static Sample prev;
	Sample next = sample_buffer[tail];
	tail = (tail + 1) % NUM_SAMPLES;

	int s = (prev.y - next.y) - (prev.x - next.x);
	prev = next;

	static int dc;
	dc += ((s << 16) - dc) >> 10;
	s -= dc >> 16;

	*sample = s;
	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);

#if 0
	gpio_disable_pulls(HP_GND_PIN);
	gpio_disable_pulls(HP_LEFT_PIN);
	gpio_disable_pulls(HP_RIGHT_PIN);

	gpio_set_function(HP_GND_PIN, GPIO_FUNC_PWM);
	gpio_set_function(HP_LEFT_PIN, GPIO_FUNC_PWM);
	gpio_set_function(HP_RIGHT_PIN, GPIO_FUNC_PWM);

	gpio_set_drive_strength(HP_GND_PIN, GPIO_DRIVE_STRENGTH_2MA);
	gpio_set_drive_strength(HP_LEFT_PIN, GPIO_DRIVE_STRENGTH_2MA);
	gpio_set_drive_strength(HP_RIGHT_PIN, GPIO_DRIVE_STRENGTH_2MA);

	/* LR at 48 kHz */
	pwm_config pc = pwm_get_default_config();
	pwm_config_set_clkdiv_int(&pc, 1);
	pwm_config_set_phase_correct(&pc, false);
	pwm_config_set_wrap(&pc, HP_PWM_SCALE - 1);
	pwm_init(HP_LR_SLICE, &pc, false);

	/* GND at high rate, 50:50 for virtual ground */
	pwm_config_set_wrap(&pc, 1);
	pwm_init(HP_GND_SLICE, &pc, false);
	pwm_set_chan_level(HP_GND_SLICE, HP_GND_CHAN, 1);

	/* Start at 50:50. */
	pwm_set_both_levels(HP_LR_SLICE, HP_PWM_MID, HP_PWM_MID);

	/* Start both */
	pwm_set_enabled(HP_GND_SLICE, true);
	pwm_set_enabled(HP_LR_SLICE, true);

	while (true) {
		for (int i = 0; i < 50; i++) {
			sleep_us(2272);
			pwm_set_both_levels(HP_LR_SLICE, HP_PWM_MID + 1, HP_PWM_MID + 1);

			sleep_us(2272);
			pwm_set_both_levels(HP_LR_SLICE, HP_PWM_MID - 1, HP_PWM_MID - 1);
		}

		for (int i = 0; i < 55; i++) {
			sleep_us(2020);
			pwm_set_both_levels(HP_LR_SLICE, HP_PWM_MID + 1, HP_PWM_MID + 1);

			sleep_us(2020);
			pwm_set_both_levels(HP_LR_SLICE, HP_PWM_MID - 1, HP_PWM_MID - 1);
		}

		sleep_ms(500);
	}
#else
	// gpio_init(IR_BIAS_PIN);
	gpio_disable_pulls(IR_BIAS_PIN);
	// gpio_set_drive_strength(IR_BIAS_PIN, GPIO_DRIVE_STRENGTH_2MA);
	// gpio_put(IR_BIAS_PIN, 1);
	// gpio_set_dir(IR_BIAS_PIN, GPIO_OUT);

	gpio_disable_pulls(IR_RX_PIN);
	gpio_disable_pulls(IR_FB_PIN);

	pio_gpio_init(IR_PIO, IR_FB_PIN);

	gpio_set_input_hysteresis_enabled(IR_RX_PIN, false);
	gpio_set_drive_strength(IR_FB_PIN, GPIO_DRIVE_STRENGTH_2MA);
	gpio_set_slew_rate(IR_FB_PIN, GPIO_SLEW_RATE_FAST);

	init_pio_sm();

	dma_ch_sample = dma_claim_unused_channel(true);
	dma_ch_rx = dma_claim_unused_channel(true);
	int timer = dma_claim_unused_timer(true);

	dma_channel_config dma_conf;

	/* Copy sampling script into the PIO. */
	dma_conf = dma_channel_get_default_config(dma_ch_sample);
	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, false);
	channel_config_set_chain_to(&dma_conf, dma_ch_rx);
	channel_config_set_dreq(&dma_conf, dma_get_timer_dreq(timer));
	dma_channel_configure(dma_ch_sample, &dma_conf, &IR_PIO->sm[IR_SM].instr, sampling_script,
			      1, false);

	/* Read samples from PIO. */
	dma_conf = dma_channel_get_default_config(dma_ch_rx);
	channel_config_set_transfer_data_size(&dma_conf, DMA_SIZE_32);
	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 + 2);
	channel_config_set_chain_to(&dma_conf, dma_ch_sample);
	channel_config_set_dreq(&dma_conf, pio_get_dreq(IR_PIO, IR_SM, false));
	dma_channel_configure(dma_ch_rx, &dma_conf, sample_buffer, &IR_PIO->rxf[IR_SM], 1, false);

	/* 8 × 38 kHz */
	// dma_timer_set_fraction(timer, 8 * 16, 55579);
	dma_timer_set_fraction(timer, 1, 256);
	dma_channel_start(dma_ch_rx);
	pio_sm_set_enabled(IR_PIO, IR_SM, true);

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

#if 0
	static int sin_lut[16] = { 0, 98,  180,	 236,  255,  236,  180,	 98,
				   0, -98, -180, -236, -255, -236, -180, -98 };
	static int cos_lut[16] = { 255,	 236,  180,  98,  0, -98, -180, -236,
				   -255, -236, -180, -98, 0, 98,  180,	236 };

	int I = 0, Q = 0;
	int phase = 0;

	while (true) {
		/* Get last byte written. */
		unsigned head = (dma_hw->ch[dma_ch_rx].write_addr - 8u) & ((4 * NUM_SAMPLES) - 1);

		/* Convert to index into 4-byte array. */
		head = (head >> 2);

		while (tail != head) {
			int sample = read_sample();

			I += (cos_lut[phase] * sample) >> 8;
			Q += (sin_lut[phase] * sample) >> 8;

			if (++phase >= 16) {
				I >>= 3;
				Q >>= 3;

				batch[written++] = I;
				batch[written++] = Q;

				I = Q = 0;
				phase = 0;
			}

			if (written >= sizeof(batch) / sizeof(*batch)) {
				written = 0;
				if (tud_cdc_n_write_available(1) >= sizeof(batch)) {
					tud_cdc_n_write(1, batch, sizeof(batch));
					tud_cdc_n_write_flush(1);
					written = 0;
				}
			}
		}

		sleep_us(10);
	}
#else
	stdio_usb_lock();

	while (true) {
		int sample;

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

			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();
#endif
#endif
}