#include <avr/io.h>
#include <avr/interrupt.h>
#include <avr/pgmspace.h>

#ifdef ENABLE_OLED
# include "i2c_master.h"
#endif

// Port A0 - A3	-> LED Data
// Port A5 		-> Select Output
// Port A4, A6	-> I2C Display
// Port A7		-> Button Input (Pulled-Up)
// Port B2 		-> Button Input (Pulled-Up)

// Size of each strip
#define LED0_COUNT 12
#define LED1_COUNT 12
#define LED2_COUNT 12
#define LED3_COUNT 15

// ms per frame
#define LED0_MS 500
#define LED1_MS 500
#define LED2_MS 500
#define LED3_MS 500

// ms offset
#define LED0_MS_OFFSET 0
#define LED1_MS_OFFSET 0
#define LED2_MS_OFFSET 0
#define LED3_MS_OFFSET 0

#define LED0_PORT PORTA
#define LED1_PORT PORTA
#define LED2_PORT PORTA
#define LED3_PORT PORTA

#define LED0_BIT 0
#define LED1_BIT 1
#define LED2_BIT 2
#define LED3_BIT 3

// NeoPixel Defines
#define T1H  900    // Width of a 1 bit in ns
#define T1L  600    // Width of a 1 bit in ns
#define T0H  400    // Width of a 0 bit in ns
#define T0L  900    // Width of a 0 bit in ns
#define RES 6000
#define RES_HALF RES / 2
#define NS_PER_SEC (1000000000L)
#define CYCLES_PER_SEC (F_CPU)
#define NS_PER_CYCLE ( NS_PER_SEC / CYCLES_PER_SEC )
#define NS_TO_CYCLES(n) ( (n) / NS_PER_CYCLE )

// Simple RGB Colour Correction
const uint8_t PROGMEM gamma8[] =
{
  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   1,   1,   1,   1,
  1,   1,   1,   1,   1,   1,   1,   1,   1,   2,   2,   2,   2,   2,   2,   2,
  2,   3,   3,   3,   3,   3,   3,   3,   4,   4,   4,   4,   4,   5,   5,   5,
  5,   6,   6,   6,   6,   7,   7,   7,   7,   8,   8,   8,   9,   9,   9,  10,
 10,  10,  11,  11,  11,  12,  12,  13,  13,  13,  14,  14,  15,  15,  16,  16,
 17,  17,  18,  18,  19,  19,  20,  20,  21,  21,  22,  22,  23,  24,  24,  25,
 25,  26,  27,  27,  28,  29,  29,  30,  31,  32,  32,  33,  34,  35,  35,  36,
 37,  38,  39,  39,  40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  50,
 51,  52,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,  66,  67,  68,
 69,  70,  72,  73,  74,  75,  77,  78,  79,  81,  82,  83,  85,  86,  87,  89,
 90,  92,  93,  95,  96,  98,  99, 101, 102, 104, 105, 107, 109, 110, 112, 114,
115, 117, 119, 120, 122, 124, 126, 127, 129, 131, 133, 135, 137, 138, 140, 142,
144, 146, 148, 150, 152, 154, 156, 158, 160, 162, 164, 167, 169, 171, 173, 175,
177, 180, 182, 184, 186, 189, 191, 193, 196, 198, 200, 203, 205, 208, 210, 213,
215, 218, 220, 223, 225, 228, 231, 233, 236, 239, 241, 244, 247, 249, 252, 255
};

// ms counter
uint64_t ms = 0;

// Led update times
uint64_t led_ms_update[4] = {LED0_MS_OFFSET, LED1_MS_OFFSET, LED2_MS_OFFSET,
	LED3_MS_OFFSET};

#ifdef ENABLE_OLED
// I2C Display Address
uint8_t display_address = 0x78;

// Draw only changed digits
uint8_t old_time[7] = {0, 0, 0, 0, 0, 0, 0};

uint64_t old_time_ms = 0;
#endif

// ms Counter Func
ISR(TIM1_CAPT_vect)
{
	++ms;
}

// NeoPixel Funcs
static inline void send_bit_led0(uint8_t bitVal);
static inline void send_bit_led1(uint8_t bitVal);
static inline void send_bit_led2(uint8_t bitVal);
static inline void send_bit_led3(uint8_t bitVal);
static inline void send_byte(uint8_t byte, uint8_t led);
static inline void send_pixel(uint8_t r, uint8_t g, uint8_t b, uint8_t led);
void show_leds();

// Hardware Init
void init();

#ifdef ENABLE_OLED
// I2C Display Funcs
void init_display();
void send_command(uint8_t cmd);
void set_display_coords(uint8_t x_start, uint8_t x_end, uint8_t y_start, uint8_t y_end);
void show_time(uint64_t ms);
void draw_dots(uint8_t pos);
void draw_digit(uint8_t pos, uint8_t digit);
#endif

// Call function to show a frame and prepare idx for next frame
void sun_and_moon(uint8_t count, uint8_t* idx, uint8_t colour, uint8_t led);
void clock_bg(uint8_t count, uint8_t* idx, uint8_t colour, uint8_t led);
void clock_no_bg(uint8_t count, uint8_t* idx, uint8_t colour, uint8_t led);
void loading(uint8_t count, uint8_t* idx, uint8_t colour, uint8_t led);

void (*func_list[4])(uint8_t, uint8_t*, uint8_t, uint8_t) =
{
	sun_and_moon, clock_bg, clock_no_bg, loading
};

void (*send_ptr[4])(uint8_t) =
{
	send_bit_led0, send_bit_led1, send_bit_led2, send_bit_led3
};

uint8_t led0_idx = 0;
uint8_t led1_idx = 0;
uint8_t led2_idx = 0;
uint8_t led3_idx = 0;
uint8_t colour = 0;
uint8_t func_idx = 0;
uint8_t btn1_old_state = 1;
uint8_t btn2_old_state = 1;
uint8_t colour_state = 0;
uint8_t current_state;

int main()
{
	init();
// Main loop
while(1)
{
	if (ms > led_ms_update[0])
	{
		func_list[func_idx](LED0_COUNT, &led0_idx, colour, 0);
		led_ms_update[0] = ms + LED0_MS;
	}

	if (ms > led_ms_update[1])
	{
		func_list[func_idx](LED1_COUNT, &led1_idx, colour, 1);
		led_ms_update[1] = ms + LED1_MS;
	}

	if (ms > led_ms_update[2])
	{
		func_list[func_idx](LED2_COUNT, &led2_idx, colour, 2);
		led_ms_update[2] = ms + LED2_MS;
	}

	if (ms > led_ms_update[0])
	{
		func_list[func_idx](LED3_COUNT, &led3_idx, colour, 3);
		led_ms_update[3] = ms + LED3_MS;
	}

	show_leds();

#ifdef ENABLE_OLED
	if (ms - 1000 >= old_time_ms)
	{
		show_time(ms);
		old_time_ms = ms;
	}
#endif

	// Button 1 Press
	current_state = PINA & 0x80;
	if (btn1_old_state && !current_state)
	{
		led0_idx = 0;
		led1_idx = 0;
		led2_idx = 0;
		led3_idx = 0;

		func_idx++;
		func_idx %= sizeof(func_list) / sizeof(func_list[0]);
	}
	btn1_old_state = current_state;

	// Button 2 Press
	current_state = PINB & 0x04;
	if (btn2_old_state && !current_state)
	{
		switch (colour_state)
		{
		case 0:
			colour = 1;
			colour_state = 1;
			break;
		case 1:
			colour = 0;
			colour_state = 2;
			PORTA ^= 0x20;
			break;
		case 2:
			colour_state = 0;
			PORTA ^= 0x20;
			break;
		}
	}
	btn2_old_state = current_state;
}
}

inline void send_bit_led0(uint8_t value)
{
    if (value)
    {
		asm volatile
		(
			"sbi %[port], %[bit] \n\t"
			".rept %[onCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			"cbi %[port], %[bit] \n\t"
			".rept %[offCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			::
			[port]		"I" (_SFR_IO_ADDR(LED0_PORT)),
			[bit]		"I" (LED0_BIT),
			[onCycles]	"I" (NS_TO_CYCLES(T1H) - 2),
			[offCycles]	"I" (NS_TO_CYCLES(T1L) - 2)
		);
    }
    else
    {
		asm volatile
		(
			"sbi %[port], %[bit] \n\t"
			".rept %[onCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			"cbi %[port], %[bit] \n\t"
			".rept %[offCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			::
			[port]		"I" (_SFR_IO_ADDR(LED0_PORT)),
			[bit]		"I" (LED0_BIT),
			[onCycles]	"I" (NS_TO_CYCLES(T0H) - 2),
			[offCycles]	"I" (NS_TO_CYCLES(T0L) - 2)
		);
    }
}

inline void send_bit_led1(uint8_t value)
{
    if (value)
    {
		asm volatile
		(
			"sbi %[port], %[bit] \n\t"
			".rept %[onCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			"cbi %[port], %[bit] \n\t"
			".rept %[offCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			::
			[port]		"I" (_SFR_IO_ADDR(LED1_PORT)),
			[bit]		"I" (LED1_BIT),
			[onCycles]	"I" (NS_TO_CYCLES(T1H) - 2),
			[offCycles]	"I" (NS_TO_CYCLES(T1L) - 2)
		);
    }
    else
    {
		asm volatile
		(
			"sbi %[port], %[bit] \n\t"
			".rept %[onCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			"cbi %[port], %[bit] \n\t"
			".rept %[offCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			::
			[port]		"I" (_SFR_IO_ADDR(LED1_PORT)),
			[bit]		"I" (LED1_BIT),
			[onCycles]	"I" (NS_TO_CYCLES(T0H) - 2),
			[offCycles]	"I" (NS_TO_CYCLES(T0L) - 2)
		);
    }
}

inline void send_bit_led2(uint8_t value)
{
    if (value)
    {
		asm volatile
		(
			"sbi %[port], %[bit] \n\t"
			".rept %[onCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			"cbi %[port], %[bit] \n\t"
			".rept %[offCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			::
			[port]		"I" (_SFR_IO_ADDR(LED2_PORT)),
			[bit]		"I" (LED2_BIT),
			[onCycles]	"I" (NS_TO_CYCLES(T1H) - 2),
			[offCycles]	"I" (NS_TO_CYCLES(T1L) - 2)
		);
    }
    else
    {
		asm volatile
		(
			"sbi %[port], %[bit] \n\t"
			".rept %[onCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			"cbi %[port], %[bit] \n\t"
			".rept %[offCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			::
			[port]		"I" (_SFR_IO_ADDR(LED2_PORT)),
			[bit]		"I" (LED2_BIT),
			[onCycles]	"I" (NS_TO_CYCLES(T0H) - 2),
			[offCycles]	"I" (NS_TO_CYCLES(T0L) - 2)
		);
    }
}

inline void send_bit_led3(uint8_t value)
{
    if (value)
    {
		asm volatile
		(
			"sbi %[port], %[bit] \n\t"
			".rept %[onCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			"cbi %[port], %[bit] \n\t"
			".rept %[offCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			::
			[port]		"I" (_SFR_IO_ADDR(LED3_PORT)),
			[bit]		"I" (LED3_BIT),
			[onCycles]	"I" (NS_TO_CYCLES(T1H) - 2),
			[offCycles]	"I" (NS_TO_CYCLES(T1L) - 2)
		);
    }
    else
    {
		asm volatile
		(
			"sbi %[port], %[bit] \n\t"
			".rept %[onCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			"cbi %[port], %[bit] \n\t"
			".rept %[offCycles] \n\t"
			"nop \n\t"
			".endr \n\t"
			::
			[port]		"I" (_SFR_IO_ADDR(LED3_PORT)),
			[bit]		"I" (LED3_BIT),
			[onCycles]	"I" (NS_TO_CYCLES(T0H) - 2),
			[offCycles]	"I" (NS_TO_CYCLES(T0L) - 2)
		);
    }
}

inline void send_byte(uint8_t byte, uint8_t led)
{
	uint8_t i;

	for (i = 0; i < 8; i++)
	{
		send_ptr[led](byte & 0xF0);
		byte <<= 1;
	}
}
  
inline void send_pixel(uint8_t r, uint8_t g, uint8_t b, uint8_t led)
{
	send_byte(g, led);
	send_byte(r, led);
	send_byte(b, led);
}

void show_leds()
{
	asm volatile
	(
		".rept %[resCycles] \n\t"
		"nop \n\t"
		".endr \n\t"
		".rept %[resCycles] \n\t"
		"nop \n\t"
		".endr \n\t"
		::
		[resCycles]		"I" (NS_TO_CYCLES(RES_HALF) - 2)
	);
}

void init()
{
	DDRA = (1 << 0) | (1 << 1) | (1 << 2) | (1 << 3) | (1 << 5);
	PORTA = (1 << 7);
	PORTB = (1 << 2);

#ifdef ENABLE_OLED
	init_display();
#endif

	// TIM1 init
	TCCR1B = (1 << WGM13) | (1 << WGM12) | (1 << CS10);
	ICR1 = (F_CPU / 1000) - 1;
	TIMSK1 = (1 << ICIE1);
	sei();
}

void sun_and_moon(uint8_t count, uint8_t* idx, uint8_t colour, uint8_t led)
{
	uint8_t i;
	uint8_t colour1[3] = 
	{
		pgm_read_byte(&gamma8[0xFF]),
		pgm_read_byte(&gamma8[0x80]),
		pgm_read_byte(&gamma8[0x00])
	};
	uint8_t colour2[3] = 
	{
		pgm_read_byte(&gamma8[0x00]),
		pgm_read_byte(&gamma8[0x80]),
		pgm_read_byte(&gamma8[0xFF])
	};
	
	if (*idx <= count / 2)
	{
		for (i = 0; i < count; ++i)
		{
			if ((i < *idx) || (i >= *idx + count / 2))
			{
				send_pixel(colour1[0], colour1[1], colour1[2], led);
			}
			else
			{
				send_pixel(colour2[0], colour2[1], colour2[2], led);
			}
		}
	}
	else
	{
		for (i = 0; i < count; ++i)
		{
			if ((i < *idx) && (i >= *idx - count / 2))
			{
				send_pixel(colour1[0], colour1[1], colour1[2], led);
			}
			else
			{
				send_pixel(colour2[0], colour2[1], colour2[2], led);
			}
		}
	}
	*idx++;
	*idx %= count;
}

void clock_bg(uint8_t count, uint8_t* idx, uint8_t colour, uint8_t led)
{
	uint8_t i;
	uint8_t colour1[3] = 
	{
		pgm_read_byte(&gamma8[0xFF]),
		pgm_read_byte(&gamma8[0x80]),
		pgm_read_byte(&gamma8[0x00])
	};
	uint8_t colour2[3] = 
	{
		pgm_read_byte(&gamma8[0x00]),
		pgm_read_byte(&gamma8[0x80]),
		pgm_read_byte(&gamma8[0xFF])
	};
	uint8_t* main_colour;
	uint8_t* bg_colour;

	if (0 == colour)
	{
		main_colour = colour1;
		bg_colour = colour2;
	}
	else
	{
		main_colour = colour2;
		bg_colour = colour1;
	}

	for (i = 0; i < count; ++i)
	{
		if (i == *idx)
		{
			send_pixel(main_colour[0], main_colour[1], main_colour[2], led);
		}
		else
		{
			send_pixel(bg_colour[0], bg_colour[1], bg_colour[2], led);
		}
	}

	*idx++;
	*idx %= count;
}

void clock_no_bg(uint8_t count, uint8_t* idx, uint8_t colour, uint8_t led)
{
	uint8_t i;
	uint8_t colour1[3] = 
	{
		pgm_read_byte(&gamma8[0xFF]),
		pgm_read_byte(&gamma8[0x80]),
		pgm_read_byte(&gamma8[0x00])
	};
	uint8_t colour2[3] = 
	{
		pgm_read_byte(&gamma8[0x00]),
		pgm_read_byte(&gamma8[0x80]),
		pgm_read_byte(&gamma8[0xFF])
	};
	uint8_t black[3] =
	{
		pgm_read_byte(&gamma8[0x00]),
		pgm_read_byte(&gamma8[0x00]),
		pgm_read_byte(&gamma8[0x00])
	};
	uint8_t* main_colour;

	if (0 == colour)
	{
		main_colour = colour1;
	}
	else
	{
		main_colour = colour2;
	}

	for (i = 0; i < count; ++i)
	{
		if (i == *idx)
		{
			send_pixel(main_colour[0], main_colour[1], main_colour[2], led);
		}
		else
		{
			send_pixel(black[0], black[1], black[2], led);
		}
	}

	*idx++;
	*idx %= count;
}

void loading(uint8_t count, uint8_t* idx, uint8_t colour, uint8_t led)
{
	uint8_t i;
	uint8_t colour1[3] = 
	{
		pgm_read_byte(&gamma8[0xFF]),
		pgm_read_byte(&gamma8[0x80]),
		pgm_read_byte(&gamma8[0x00])
	};
	uint8_t colour2[3] = 
	{
		pgm_read_byte(&gamma8[0x00]),
		pgm_read_byte(&gamma8[0x80]),
		pgm_read_byte(&gamma8[0xFF])
	};

	if (*idx < count)
	{
		for (i = 0; i < count; ++i)
		{
			if (i <= *idx)
			{
				send_pixel(colour1[0], colour1[1], colour1[2], led);
			}
			else
			{
				send_pixel(colour2[0], colour2[1], colour2[2], led);
			}
		}
	}
	else
	{
		for (i = 0; i < count; ++i)
		{
			if (i <= *idx - count)
			{
				send_pixel(colour2[0], colour2[1], colour2[2], led);
			}
			else
			{
				send_pixel(colour1[0], colour1[1], colour1[2], led);
			}
		}
	}

	*idx++;
	*idx %= count * 2;
}

#ifdef ENABLE_OLED
void init_display()
{
	uint16_t i = 0;
	uint8_t cmdInit[] =
	{
		0xAE, 0xD5, 0x80, 0xA8,
		0x1F, 0xD3, 0x00, 0x40,
		0x8D, 0x14, 0x20, 0x00,
		0xA1, 0xC8, 0xDA, 0x02,
		0x81, 0x8F, 0xD9, 0xF1,
		0xDB, 0x40, 0xA4, 0xA6,
		0x2E, 0xAF
	};

	i2c_init();
    for (i = 0; i < sizeof(cmdInit); ++i)
    {
        send_command(cmdInit[i]);
    }

    set_display_coords(0x00, 0x7F, 0x00, 0x03);
    
    for(i = 0; i < 128*4; ++i)
    {
		i2c_write(0x00);
	}
	i2c_stop();
}

void send_command(uint8_t cmd)
{
    i2c_start(display_address);
    i2c_write(0x00);
    i2c_write(cmd);
    i2c_stop();
}

void set_display_coords(uint8_t x_start, uint8_t x_end, uint8_t y_start, uint8_t y_end)
{
	send_command(0x21);
	send_command(x_start);
	send_command(x_end);
	
	send_command(0x22);
	send_command(y_start);
	send_command(y_end);
	
	i2c_start(display_address);
	i2c_write(0x40);
}

void show_time(uint64_t ms)
{
	uint8_t new_time[7] = {0, 0, 0, 0, 0, 0, 0};
	uint8_t i;

	ms /= 1000;
	new_time[6] = ms % 10;
	ms /= 10;
	new_time[5] = ms % 6;
	ms /= 6;
	new_time[4] = ms % 10;
	ms /= 10;
	new_time[3] = ms % 6;
	ms /= 6;
	new_time[2] = ms % 10;
	ms /= 10;
	new_time[1] = ms % 10;
	ms /= 10;
	new_time[0] = ms % 10;

	for (i = 0; i < 7; ++i)
	{
		if (old_time[i] != new_time[i])
		{
			if (((2 == i) || (4 == i)) && (0 == old_time[i]))
			{
				draw_dots(i);
			}

			draw_digit(i, new_time[i]);
		}
		old_time[i] = new_time[i];
	}
}

void draw_dots(uint8_t pos)
{
	// Between hours and minutes
	if (2 == pos)
	{
		set_display_coords(50, 53, 1, 2);
	}

	// Between minutes and seconds
	if (4 == pos)
	{
		set_display_coords(90, 93, 1, 2);
	}

	i2c_write(0x0F);
	i2c_write(0x0F);
	i2c_write(0x0F);
	i2c_write(0x0F);
	i2c_write(0xF0);
	i2c_write(0xF0);
	i2c_write(0xF0);
	i2c_write(0xF0);
	i2c_stop();
}

void draw_digit(uint8_t pos, uint8_t digit)
{
	uint8_t i, j;
	uint8_t x_start;
	uint8_t x_end;
	uint8_t row[12];
	
	x_start = pos * 16;

	if (pos > 2)
	{
		x_start += 8;
	}

	if (pos > 4)
	{
		x_start += 8;
	}

	x_end = x_start + 15;
	set_display_coords(x_start + 2, x_end - 2, 0, 3);

	row[0] = 0xF0;
	row[1] = 0x00;
	row[2] = 0x00;
	row[3] = 0x03;
	row[4] = 0x00;
	row[5] = 0x00;
	row[6] = 0xC0;
	row[7] = 0xC0;
	row[8] = 0xC0;
	row[9] = 0x0F;
	row[10] = 0x03;
	row[11] = 0x0F;

	if (digit > 1)
	{
		row[1] = 0xF0;
	}

	if (digit > 2)
	{
		row[2] = 0xF0;
	}

	for (i = 3; i < 8; ++i)
	{
		if (digit > i)
		{
			row[i] = 0xFF;
		}
	}

	if (0 == digit)
	{
		row[2] = 0xF0;
		row[3] = 0xFF;
		row[5] = 0xFF;
		row[6] = 0xFF;
		row[8] = 0xFF;
	}

	for (i = 0; i < 12 * 4; ++i)
	{
		i2c_write(row[i / 4]);
	}

	i2c_stop();
}

#endif