/* ------------------------------------------------------------------------
 * Microsoft Video 1 Encoder
 *
 * Dr. Tim Ferguson, 2001.
 * For more details on the algorithm:
 *         http://www.csse.monash.edu.au/~timf/videocodec.html
 *
 * You may freely use this source code.  I only ask that you reference its
 * source in your projects documentation:
 *       Tim Ferguson: http://www.csse.monash.edu.au/~timf/
 * ------------------------------------------------------------------------ */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <unistd.h>
#include <math.h>
#include "msvc.h"

#define DBUG	0

#define AVIIF_TWOCC         0x00000002L
#define AVIIF_KEYFRAME      0x00000010L

typedef struct {
	int width, height, key_rate, current_frame, last_key, mode_cnt[4];
	unsigned char *prev_img;
} msvc_enc_info;

/* ------------------------------------------------------------------------ */
#define MSVC_C1(ip, clr, rdec) { \
	*ip++ = clr; *ip++ = clr; *ip++ = clr; *ip = clr; ip -= rdec; \
	*ip++ = clr; *ip++ = clr; *ip++ = clr; *ip = clr; ip -= rdec; \
	*ip++ = clr; *ip++ = clr; *ip++ = clr; *ip = clr; ip -= rdec; \
	*ip++ = clr; *ip++ = clr; *ip++ = clr; *ip = clr; }

#define MSVC_C2(ip,flag,cA,cB,rdec) { \
	*ip++ = (flag & 0x01) ? (cB) : (cA); *ip++ = (flag & 0x02) ? (cB) : (cA); \
	*ip++ = (flag & 0x04) ? (cB) : (cA); *ip   = (flag & 0x08) ? (cB) : (cA); \
	ip-=rdec; \
	*ip++ = (flag & 0x10) ? (cB) : (cA); *ip++ = (flag & 0x20) ? (cB) : (cA); \
	*ip++ = (flag & 0x40) ? (cB) : (cA); *ip   = (flag & 0x80) ? (cB) : (cA); }

#define MSVC_C4(ip,flag,cA0,cA1,cB0,cB1,rdec) { \
	*ip++ = (flag & 0x01) ? (cB0) : (cA0); *ip++ =(flag & 0x02) ? (cB0) : (cA0); \
	*ip++ = (flag & 0x04) ? (cB1) : (cA1); *ip   =(flag & 0x08) ? (cB1) : (cA1); \
	ip-=rdec; \
	*ip++ = (flag & 0x10) ? (cB0) : (cA0); *ip++ =(flag & 0x20) ? (cB0) : (cA0); \
	*ip++ = (flag & 0x40) ? (cB1) : (cA1); *ip   =(flag & 0x80) ? (cB1) : (cA1); }

	/* convert a 15 bit colour value to 24 bits */
#define COL_15_TO_24(color) \
	((((color) << 9) & (0x1f << 19)) | (((color) << 6) & (0x1f << 11)) | \
		(((color) & 0x1f) << 3))

/* ------------------------------------------------------------------------ */
static void msvc_apply_block(unsigned char *img, int stride, unsigned long col[], unsigned long index)
{
int row_dec;
unsigned long *i_ptr = (unsigned long *)(img + stride * 3);

	row_dec = stride/4 + 3;

	if(index & 0x8000)		/* 1 color encoding (80-83) && (>=88)*/
		{
		unsigned long clr = COL_15_TO_24(index);
		MSVC_C1(i_ptr,clr,row_dec);
		}
	else		/* 2 or 8 color encoding */
		{
		unsigned long cA0,cB0;
		cB0 = COL_15_TO_24(col[0]);
		cA0 = COL_15_TO_24(col[1]);
		if(col[0] & 0x8000)   /* Eight Color Encoding */
			{
			unsigned long cA1, cB1, code = index >> 8;
			cB1 = COL_15_TO_24(col[2]);
			cA1 = COL_15_TO_24(col[3]);
			MSVC_C4(i_ptr,index,cA0,cA1,cB0,cB1,row_dec);
			i_ptr -= row_dec;
			cB0 = COL_15_TO_24(col[4]);
			cA0 = COL_15_TO_24(col[5]);
			cB1 = COL_15_TO_24(col[6]);
			cA1 = COL_15_TO_24(col[7]);
			MSVC_C4(i_ptr,code,cA0,cA1,cB0,cB1,row_dec);
			}
		else /* Two Color Encoding */
			{
			unsigned long code = index >> 8;
			MSVC_C2(i_ptr,index,cA0,cB0,row_dec);
			i_ptr -= row_dec;
			MSVC_C2(i_ptr,code,cA0,cB0,row_dec);
			}
		} /* end of 2 or 8 */
}


/* ------------------------------------------------------------------------ */
static double msvc_lbg_array(unsigned char pixels[][3], int len, unsigned long *col1,
	unsigned long *col2, unsigned long *index)
{
unsigned char bin_index[16], num[2];
unsigned long best[2][3], old_index = 0, new_index;
int i, j, iter, diff;
double sd = 0;

		/* init best 2 colours with first and last block colour */
	for(j = 0; j < 2; j++)
		for(i = 0; i < 3; i++) best[j][i] = pixels[j*(len-1)][i];

		/* LBG style VQ algorithm for finding best 2 colours for block */
	for(iter = 0; iter < 20; iter++)
		{
			/* sort into bins */
		num[0] = num[1] = 0;
		for(i = 0; i < len; i++)
			{
			bin_index[i] = (
			(abs(pixels[i][0] - best[0][0]) + abs(pixels[i][1] - best[0][1]) +
				abs(pixels[i][2] - best[0][2])) >
				(abs(pixels[i][0] - best[1][0]) + abs(pixels[i][1] - best[1][1]) +
				abs(pixels[i][2] - best[1][2]))  ?  1 : 0);
			num[bin_index[i]]++;
			}
		if(num[0] == 0) bin_index[0] = 0;		/* must always have stuff in bins */
		if(num[1] == 0) bin_index[1] = 1;

			/* generate new vectors */
		best[0][0] = best[0][1] = best[0][2] = best[1][0] = best[1][1] = best[1][2] = 0;
		num[0] = num[1] = 0;
		new_index = 0;
		for(i = 0; i < len; i++)
			{
			num[bin_index[i]]++;
			best[bin_index[i]][0] += pixels[i][0];
			best[bin_index[i]][1] += pixels[i][1];
			best[bin_index[i]][2] += pixels[i][2];
			new_index |= ((unsigned long)bin_index[i] << i);
			}

		for(j = 0; j < 2; j++)
			for(i = 0; i < 3; i++)
				{
				best[j][i] = best[j][i]/num[j];
				if(best[j][i] > 255) best[j][i] = 255;
				}

		if(old_index == new_index) break;	/* index never == 0 */
		old_index = new_index;
		}

		/* quantise colours */
	for(j = 0; j < 2; j++)
		for(i = 0; i < 3; i++) best[j][i] &= 0xf8;

		/* calculate error */
	for(i = 0; i < len; i++)
		{
		diff = (int)best[bin_index[i]][0] - pixels[i][0];  sd += (double)(diff * diff);
		diff = (int)best[bin_index[i]][1] - pixels[i][1];  sd += (double)(diff * diff);
		diff = (int)best[bin_index[i]][2] - pixels[i][2];  sd += (double)(diff * diff);
		}

	*col1 = (best[0][0] << 7) | (best[0][1] << 2) | (best[0][2] >> 3);
	*col2 = (best[1][0] << 7) | (best[1][1] << 2) | (best[1][2] >> 3);
	*index = new_index;

	return sd;
}


/* ------------------------------------------------------------------------
 *  1 colour for 4x4 block
 */
static double msvc_colour1(unsigned char *img, int stride, unsigned long *col)
{
double sd = 0;
int diff;
unsigned long avg_r, avg_g, avg_b;
unsigned char *lp, *cp;

	avg_r = avg_g = avg_b = 0;
	for(lp = img; lp < img + 4 * stride; lp += stride)
		for(cp = lp; cp < lp+12; cp += 3)
			{
			avg_r += (int)cp[2];
			avg_g += (int)cp[1];
			avg_b += (int)cp[0];
			}
	avg_r = (avg_r >> 4) & 0xf8;
	avg_g = (avg_g >> 4) & 0xf8;
	avg_b = (avg_b >> 4) & 0xf8;

	for(lp = img; lp < img + 4 * stride; lp += stride)
		for(cp = lp; cp < lp+12; cp += 3)
			{
			diff = (int)cp[2] - avg_r;  sd += (double)(diff * diff);
			diff = (int)cp[1] - avg_g;  sd += (double)(diff * diff);
			diff = (int)cp[0] - avg_b;  sd += (double)(diff * diff);
			}

	*col = (avg_r << 7) | (avg_g << 2) | (avg_b >> 3);
	return sd;
}


/* ------------------------------------------------------------------------ */
static double msvc_colour2(unsigned char *img, int stride, unsigned long *col1,
	unsigned long *col2, unsigned long *index)
{
unsigned char pixels[16][3];
unsigned long tmp;
int j;
double sd = 0;
unsigned char *lp, *cp;

	for(j = 0, lp = img + 3 * stride; lp >= img; lp -= stride)
		for(cp = lp; cp < lp+12; cp += 3, j++)
			{
			pixels[j][0] = (int)cp[2];
			pixels[j][1] = (int)cp[1];
			pixels[j][2] = (int)cp[0];
			}

	sd = msvc_lbg_array(pixels, 16, col2, col1, index);
	if(*index & 0x8000)		/* high index bit must be zero */
		{
		tmp = *col1; *col1 = *col2; *col2 = tmp;
		*index ^= 0xffff;
		}

	return sd;
}


/* ------------------------------------------------------------------------ */
static double msvc_colour8(unsigned char *img, int stride, unsigned long cols[8], unsigned long *index)
{
unsigned char pixels[4][3];
unsigned long tmp, new_index;
int j;
double sd;
unsigned char *lp, *cp;

	for(j = 0, lp = img + 3 * stride; lp >= img + 2 * stride; lp -= stride)
		for(cp = lp; cp < lp+6; cp += 3, j++)
			{
			pixels[j][0] = (int)cp[2];
			pixels[j][1] = (int)cp[1];
			pixels[j][2] = (int)cp[0];
			}
	sd = msvc_lbg_array(pixels, 4, cols+1, cols, &new_index);
	*index = (new_index & 0x3) | ((new_index & 0xC) << 2);

	for(j = 0, lp = img + 3 * stride; lp >= img + 2 * stride; lp -= stride)
		for(cp = lp+6; cp < lp+12; cp += 3, j++)
			{
			pixels[j][0] = (int)cp[2];
			pixels[j][1] = (int)cp[1];
			pixels[j][2] = (int)cp[0];
			}
	sd += msvc_lbg_array(pixels, 4, cols+3, cols+2, &new_index);
	*index |= (((new_index & 0x3) | ((new_index & 0xC) << 2)) << 2);

	for(j = 0, lp = img + stride; lp >= img; lp -= stride)
		for(cp = lp; cp < lp+6; cp += 3, j++)
			{
			pixels[j][0] = (int)cp[2];
			pixels[j][1] = (int)cp[1];
			pixels[j][2] = (int)cp[0];
			}
	sd += msvc_lbg_array(pixels, 4, cols+5, cols+4, &new_index);
	*index |= (((new_index & 0x3) | ((new_index & 0xC) << 2)) << 8);

	for(j = 0, lp = img + stride; lp >= img; lp -= stride)
		for(cp = lp+6; cp < lp+12; cp += 3, j++)
			{
			pixels[j][0] = (int)cp[2];
			pixels[j][1] = (int)cp[1];
			pixels[j][2] = (int)cp[0];
			}
	sd += msvc_lbg_array(pixels, 4, cols+7, cols+6, &new_index);
	if(new_index & 0x08)		/* high index bit must be zero */
		{
		tmp = cols[7]; cols[7] = cols[6]; cols[6] = tmp;
		new_index ^= 0xf;
		}
	*index |= (((new_index & 0x3) | ((new_index & 0xC) << 2)) << 10);

	*cols = *cols | 0x8000;
	return sd;
}


/* ------------------------------------------------------------------------ */
static double msvc_block_diff(unsigned char *cur_img, int cur_stride, unsigned char *prev_img, int prev_stride)
{
double sd = 0;
int diff;
unsigned char *c_lp, *c_cp, *p_lp, *p_cp;

	for(c_lp = cur_img, p_lp = prev_img; c_lp < cur_img + 4 * cur_stride; c_lp += cur_stride, p_lp += prev_stride)
		for(c_cp = c_lp, p_cp = p_lp; c_cp < c_lp+12; c_cp += 3, p_cp += 4)
			{
			diff = (int)c_cp[2] - (int)p_cp[2];  sd += (double)(diff * diff);
			diff = (int)c_cp[1] - (int)p_cp[1];  sd += (double)(diff * diff);
			diff = (int)c_cp[0] - (int)p_cp[0];  sd += (double)(diff * diff);
			}
	return sd;
}


/* ------------------------------------------------------------------------
 * Initalise a new encoding of a set of frames using MSVC.
 *
 * Input:
 *   width - the width of the input frames
 *   height - the height of the input frames
 *   key_rate - the rate at which key frames are to be inserted
 * Output:
 *   returns a context which must be passed to the encoder
 */
void *msvc_encode_init(int width, int height, int key_rate)
{
msvc_enc_info *mi;
int i;

	if((mi = calloc(sizeof(msvc_enc_info), 1)) == NULL) return NULL;
	if((mi->prev_img = calloc(width * height, 4)) == NULL)
		{
		free(mi);
		return NULL;
		}
	mi->width = width;
	mi->height = height;
	mi->key_rate = key_rate;
	mi->current_frame = mi->last_key = 0;
	for(i = 0; i < 4; i++) mi->mode_cnt[i] = 0;
	return (void *)mi;
}


/* ------------------------------------------------------------------------
 * This function encodes an RGB frame into a buffer using MSVC.
 *
 * Input:
 *   context - the context created using msvc_encode_init()
 *   frame - the input frame buffer (24-bit RGB format)
 *   width - the width of the input frame
 *   height - the height of the input frame
 *   quality - the quality this frame is to be encoded with (0=best to 100=bad)
 * Output:
 *   buf - the encoded output buffer to be saved in an AVI file
 *   size - the number of bytes put in the output buffer after encoding
 *   flags - the AVI flags the frame should be coded with (for the index)
 */
void msvc_encode(void *context, unsigned char *frame, int width,
	int height, int quality, unsigned char *buf, int *size, unsigned int *flags)
{
msvc_enc_info *mi = (msvc_enc_info *)context;
int i, x, y, block_cnt, done = 0, cmode, key_code = 0, num_skip;
unsigned long col1, col2a, col2b, col2_index, col8[8], col8_index;
double err[4], bits[4], lambda;
unsigned char *out_ptr, *img_ptr, *prev_img_ptr;
#define put_word(v) { *(out_ptr++) = (v) & 0xff; *(out_ptr++) = ((v) >> 8) & 0xff; }

	out_ptr = buf;

	lambda = (quality *  quality * quality)/10000.0;		/* a little more linear */

	if(width != mi->width || height != mi->height)
		{
		printf("MSVC ERROR: width and height have changed between init and encode!\n");
		return;
		}

	if(mi->current_frame == 0 || (mi->key_rate > 0 && mi->last_key >= mi->key_rate))
		{
		key_code = 1;
		mi->last_key = 0;
		}

	block_cnt = ((width * height) >> 4) + 1;
	done = 0;
	x = 0;
	y = height - 1;
	num_skip = 0;

	while(!done)
		{
		block_cnt--;
		if(!block_cnt || y < 0)
			{
			if(num_skip > 0) put_word(0x8400+num_skip);
			put_word(0);
			done = 1;
			continue;
			}

		img_ptr = frame + ((width * (y-3)) + x) * 3;
		prev_img_ptr = mi->prev_img + ((width * (y-3)) + x) * 4;

			/* Try coding using skip block */
		if(key_code) { err[0] = HUGE_VAL; bits[0] = 0; }
		else
			{
			err[0] = msvc_block_diff(img_ptr, width * 3, prev_img_ptr, width * 4);
			bits[0] = 16.0/(double)(num_skip + 1);
			}

			/* Try coding using 1 colour */
		err[1] = msvc_colour1(img_ptr, width * 3, &col1);
		bits[1] = 16;

			/* Try coding using 2 colours */
		err[2] = msvc_colour2(img_ptr, width * 3, &col2a, &col2b, &col2_index);
		bits[2] = 48;

			/* Try coding using 8 colours */
		err[3] = msvc_colour8(img_ptr, width * 3, col8, &col8_index);
		bits[3] = 144;

			/* select best coding type based on lambda = quality */
		cmode = 0;
		for(i = 0; i < 4; i++)
			{
			if(err[i] + lambda * bits[i] < err[cmode] + lambda * bits[cmode])
				cmode = i;
			}

			/* if not another skip block and skips pending: write skip */
		if((cmode != 0 && num_skip > 0) || num_skip >= 0x3FF)
			{
			put_word(0x8400+num_skip);
			num_skip = 0;
			}

			/* write encoded block and apply it to frame store */
		switch(cmode)
			{
			case 0: num_skip++; break;
			case 1:
				put_word(col1 | 0x8000);
				msvc_apply_block(prev_img_ptr, width * 4, col8, col1);
				break;
			case 2:
				put_word(col2_index);
				put_word(col2a);
				put_word(col2b);
				col8[0] = col2a; col8[1] = col2b;
				msvc_apply_block(prev_img_ptr, width * 4, col8, col2_index);
				break;
			case 3:
				put_word(col8_index);
				for(i = 0; i < 8; i++) put_word(col8[i]);
				msvc_apply_block(prev_img_ptr, width * 4, col8, col8_index);
				break;
			}

		mi->mode_cnt[cmode]++;
		x += 4;
		if(x >= width) { x=0; y -= 4; }
		}

	mi->current_frame++;
	mi->last_key++;

	*size = (out_ptr - buf);
	*flags = (key_code ? AVIIF_KEYFRAME : 0);
}


/* ------------------------------------------------------------------------
 * Free memory allocated when encoding a sequence.
 * Pass this function the context created by msvc_encode_init() and used
 * by msvc_encode().
 */
void msvc_encode_free(void *context)
{
msvc_enc_info *mi = (msvc_enc_info *)context;

	if(mi == NULL) return;

#if DBUG
	printf("MSVC modes - skip: %d  col1: %d  col2: %d  col8: %d\n", mi->mode_cnt[0], mi->mode_cnt[1],
		mi->mode_cnt[2], mi->mode_cnt[3]);
#endif

	if(mi->prev_img != NULL) free(mi->prev_img);
	free(mi);
}