/*
 * avr.c
 *
 *  Created on: Oct 21, 2010
 *      Author: brainiaq
 */

#include <avr/io.h>
//#define F_CPU 20e6 //only if external
//#include <util/delay.h>
//#include <avr/interrupt.h>
#include <avr/pgmspace.h>
//#include <math.h>
#include <stdint.h>
#include <stdlib.h>

#define NUM_STATES 3
#define NUM_TRANS NUM_STATES*NUM_STATES
#define MAX_BUFFER 10
#define CMD_ON 1
#define CMD_OFF 2
#define MAX_BUTTON_OFF 200000
#define INT_CONV_CONST UINT32_MAX / UINT8_MAX

static uint32_t statesOn[NUM_STATES];
static uint32_t statesOff[NUM_STATES];
static uint32_t statesCount[NUM_STATES];
static uint8_t transMat[NUM_TRANS];
static uint8_t forBack[NUM_STATES*MAX_BUFFER];
static uint32_t onBuff[MAX_BUFFER];
static uint32_t offBuff[MAX_BUFFER];
static uint32_t probs[NUM_STATES];//convenience array
static uint8_t prevProbs[NUM_STATES];//for convenience
static uint32_t transProbs[NUM_STATES];

static uint8_t recording = 0;
static uint8_t cmd = 0;
static uint32_t count = 0;
static uint8_t bufferStart = 0;
static uint8_t bufferStop = 0;
static uint8_t currState = NUM_STATES;
static uint8_t j;
static uint8_t i;
static uint8_t k;



uint8_t getInd(uint8_t i, uint8_t j, uint8_t n){
	return i * n + j;
}

void turnON(){
	PORTB &= ~_BV(PB2);
	count = 0;
	cmd = CMD_ON;
}

void turnOFF(){
	PORTB |= _BV(PB2);
	count = 0;
	cmd = CMD_OFF;
}

//void showErr(){
//	while(1){
//		for(j = 0; j < 200; ++j){
//			for(i = 0; i < 200; ++i)
//				turnOFF();
//		}
//		for(j = 0; j < 200; ++j){
//			for(i = 0; i < 200; ++i)
//				turnON();
//		}
//	}
//}

void sampleState(){
	if(currState >= NUM_STATES){
		currState = rand() % NUM_STATES;
	} else {
		//sample a state
		uint8_t val = rand() % (UINT8_MAX);
		uint16_t sum = 0;
		for(j = 0; j < NUM_STATES; ++j){
			sum += transMat[getInd(currState, j, NUM_STATES)];
			if(sum >= val){
				currState = j;
				return;
			}
		}
		currState = NUM_STATES-1;
	}
}

void normalizeArr(uint32_t* arr, uint8_t length){
	uint32_t sum = 0;
	for(i = 0; i < length; ++i)
		sum += arr[i];
	if(sum == 0) sum = length;
	for(i = 0; i < length; ++i){
		arr[i] *= (UINT32_MAX / sum);
		if(arr[i] == 0) arr[i] = INT_CONV_CONST;
	}
}

void fillEmissionProbs(uint8_t buffInd){
	int32_t a;
	int32_t b;
	uint32_t max = 0;
	for(i = 0; i < NUM_STATES; ++i){
		a = onBuff[buffInd] - statesOn[i];
		b = offBuff[buffInd] - statesOff[i];
		a = a < 0 ? -a : a;
		b = b < 0 ? -b : b;
		probs[i] = a + b;
		//if(probs[i] > max) max = probs[i];
		max += probs[i];//better, i think
	}

	for(i = 0; i < NUM_STATES; ++i){
		probs[i] = (max - probs[i]);
	}

	normalizeArr(probs, NUM_STATES);

}

uint8_t nextBuffer(uint8_t curr){
	return (curr + 1) % MAX_BUFFER;
}

uint8_t lastBuffer(uint8_t curr){
	return curr == 0 ? MAX_BUFFER - 1 : curr - 1;
}

void normalizeForBackBuffer(uint8_t ind){
	uint32_t sum = 0;
	uint8_t tempInd;
	for(j = 0; j < NUM_STATES; ++j)
		sum += forBack[getInd(ind, j, MAX_BUFFER)];
	if(sum == 0) sum = NUM_STATES;
	for(j = 0; j < NUM_STATES; ++j){
		tempInd = getInd(ind, j, MAX_BUFFER);
		if(forBack[tempInd] <= 0) forBack[tempInd] = 1;
		forBack[tempInd] = (((uint32_t)forBack[tempInd]) * UINT8_MAX) / sum;
	}
}



void calcProbs(uint8_t ind){//for both forward and backward (if all vars are right!!)

	for(i = 0; i < NUM_STATES; ++i){
		transProbs[i] = 0;
		for(j = 0; j < NUM_STATES; ++j){
			transProbs[i] += ((uint32_t)prevProbs[j]) * transMat[getInd(j,i,NUM_STATES)];
		}
		//tricky, cause precisions don't match
		//need to divide by UINT8_MAX*UINT8_MAX and multiply by probs[i],
		//without overflow or underflow (a*%#$afa3#$%dk#%fs) :::punches face:::
		for(j = 0; j < 2; ++j){
			if(probs[i] > transProbs[i]){
				probs[i] /= UINT8_MAX;
			} else {
				transProbs[i] /= UINT8_MAX;
			}
		}
		probs[i] *= transProbs[i];
		if(probs[i] == 0) probs[i] = INT_CONV_CONST;
	}
}

void train(){
	uint8_t ind = bufferStart;
	uint8_t forBackInd = 0;
	uint8_t doLoop = 1;
	uint8_t numInBuffer = 0;

	uint8_t prevProbs[NUM_STATES];//for convenience

	while(doLoop){//forward loop
		fillEmissionProbs(ind);
		if(ind != bufferStart){
			calcProbs(ind);
		}

		for(j = 0; j < NUM_STATES; ++j){//convert to low precision etc.
			prevProbs[j] = probs[j] / (INT_CONV_CONST);
			forBack[getInd(forBackInd, j, MAX_BUFFER)] = prevProbs[j];
		}
		++forBackInd;
		++numInBuffer;
		if(ind == bufferStop || forBackInd == MAX_BUFFER) doLoop = 0;
		ind = nextBuffer(ind);
	}

	--forBackInd;
	ind = bufferStop;
	doLoop = 1;

	while(doLoop){//backward loop
		if(ind != bufferStop){
			calcProbs(ind);
		} else {//first time - even probs
			for(j = 0; j < NUM_STATES; ++j)
				probs[j] = UINT32_MAX / NUM_STATES;//just any constant
		}

		for(j = 0; j < NUM_STATES; ++j){
			prevProbs[j] = probs[j] / (INT_CONV_CONST);
			forBack[getInd(forBackInd, j, MAX_BUFFER)] = (((uint32_t)forBack[getInd(forBackInd, j, MAX_BUFFER)])
			                                                     * prevProbs[j]) / UINT8_MAX;
		}
		normalizeForBackBuffer(forBackInd);

		fillEmissionProbs(ind);//USE NEXT LOOP!

		if(ind == bufferStart) doLoop = 0;

		ind = lastBuffer(ind);
		--forBackInd;
	}
	//showErr();

	//update model parameters
	ind = bufferStart;//should be bufferStart
	forBackInd = 0; //should be zero!!
	doLoop = 1;
	//turnON();
	while(doLoop){//this way is simpler

		for(i = 0; i < NUM_STATES; ++i){
			//statesOn[i] /= UINT8_MAX;//pasodfasdf
			//statesOff[i] /= UINT8_MAX;
			statesOn[i] *= statesCount[i];
			statesOff[i] *= statesCount[i];
			//turnON();
			statesOn[i] += (forBack[getInd(forBackInd, i, MAX_BUFFER)] * (onBuff[ind] / UINT8_MAX));// / UINT8_MAX;
			statesOff[i] += (forBack[getInd(forBackInd, i, MAX_BUFFER)] * (offBuff[ind] / UINT8_MAX));// / UINT8_MAX;
			//turnOFF();
			statesCount[i] += forBack[getInd(forBackInd, i, MAX_BUFFER)];
			//turnON();
			statesOn[i] = (statesOn[i] / statesCount[i]);// * UINT8_MAX;
			statesOff[i] = (statesOff[i] / statesCount[i]);// * UINT8_MAX;
			//turnOFF();
		}
		++forBackInd;
		if(ind == bufferStop || forBackInd == MAX_BUFFER) doLoop = 0;
		ind = nextBuffer(ind);
	}

	//showErr();

	//turnOFF();

	//don't forget transition probabilities
	//use 'count' as a count and 'transProbs' to count transitions
	//turnON();
	for(i = 0; i < NUM_STATES; ++i){
		count = 0;
		for(j = 0; j < NUM_STATES; ++j)
			transProbs[j] = 0;
		for(j = 0; j < numInBuffer-1; ++j){
			for(k = 0; k < NUM_STATES; ++k){
				transProbs[k] += ((uint32_t)forBack[getInd(j+1,k,MAX_BUFFER)]) * forBack[getInd(j,i,MAX_BUFFER)];
				count += ((uint32_t)forBack[getInd(j+1,k,MAX_BUFFER)]) * forBack[getInd(j,i,MAX_BUFFER)];
			}
		}
		for(j = 0; j < NUM_STATES; ++j){
			transMat[getInd(i,j,NUM_STATES)] = (UINT8_MAX * transProbs[j]) / count;
		}
	}
	//showErr();
	//turnOFF();
}



void incBuffer(){
	bufferStop = nextBuffer(bufferStop);
	if(bufferStart == bufferStop){
		bufferStart = nextBuffer(bufferStart);
	}
}



void startRecording(){
	recording = 1;
	count = 0;
	bufferStart = 0;
	bufferStop = 0;
	turnON();
}

void startPlaying(){
	recording = 0;
	count = 0;
	currState = NUM_STATES;
	sampleState();
	turnON();
}

uint8_t buttonPressed(){
	return !(PINA & _BV(PA3));
}

uint8_t main(){
	CLKPR = (1 << CLKPCE);
	CLKPR = (0 << CLKPS3) | (0 << CLKPS2) | (0 << CLKPS1) | (0 << CLKPS0);
	//Button is PA3
	//LED = PB2
	//set
	PORTA |= _BV(PA3);//turn on pull-up resistor
	DDRB |= _BV(PB2);//make LED an output

//	while(1){
//		if(buttonPressed()){
//			turnON();
//		} else {
//			turnOFF();
//		}
//		//_delay_ms(200);
//	}

	//set up constants
	//init markov model (states on and off, trans mat)
	for(i = 0; i < NUM_TRANS; ++i)
		transMat[i] = (uint8_t)(UINT8_MAX / NUM_STATES);
	for(i = 0; i < NUM_STATES; ++i){
		statesOn[i] = rand();
		statesOff[i] = rand();
		statesCount[i] = 1;
		probs[i] = 1;
	}
	for(i = 0; i < NUM_STATES * MAX_BUFFER; ++i){
		forBack[i] = 1;
	}
	for(i = 0; i < MAX_BUFFER; ++i){
		onBuff[i] = 0;
		offBuff[i] = 0;
	}

	startPlaying();
	while(1){
		uint8_t button = buttonPressed();
		count = count + 1;
		if(recording){
			if(cmd == CMD_ON){
				if(!button){
					onBuff[bufferStop] = count;
					turnOFF();
				}
			} else if (cmd == CMD_OFF){
				if(button){
					offBuff[bufferStop] = count;
					incBuffer();
					turnON();
				} else if (count >= MAX_BUTTON_OFF){
					offBuff[bufferStop] = count;
					//incBuffer();
					train();
					startPlaying();
				}
			}
		} else {//playing
			if(cmd == CMD_ON){
				if(count >= statesOn[currState]){
					turnOFF();
				}
			} else if(cmd == CMD_OFF) {
				if(count  >= statesOff[currState]){
					sampleState();
					turnON();
				}
			}

			if(button){
				count = 0;
				startRecording();
			}
		}

	}
}