/*
 *    Copyright (C) 2001 Robert Kesterson <robertk@robertk.com>
 *
 *    This program is free software; you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation; either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with this program; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 *    This utility is based on "yuvmedianfilter.c" from mjpegtools.
 * 	 (basically used it as a framework and just replaced the filter part).
 *
 *    This filter pixelates an area of a video (like blanking out a 
 *    person's face or whatever).  Works by reading a list of commands from
 *    a text file (which must be specified on the command line) -- see the
 *    Usage() function for more details.
 *
 */
#include <config.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <string.h>

extern "C" {
#include "yuv4mpeg.h"
#include "mjpeg_logging.h"
}

#include <vector>
#include <fstream>
#include <string>
#include <algorithm>

struct pixelate
{
	int x;
	int y;
	int width;
	int height;
	int size;	// pixel range
	long startframe;
	long endframe;
};
				
int verbose = 1;

// just does a pixelate.size square block and fills it with the average pixel in that block
// and repeats across the desired x/y range
void filter(int width, int height, std::vector<pixelate>& commands, unsigned char *input[], unsigned char *output[])
{
	static long framenum = 1;
	int x, y, x1, y1, avgLuma, avgChroma1, avgChroma2;
	long lsum, csum1, csum2;
	int chromasize = (width / 2) * (height / 2);

	// first copy full output frame
	memcpy(output[0], input[0], width * height);
	memcpy(output[1], input[1], chromasize);
	memcpy(output[2], input[2], chromasize);
	// now do the pixelization
	for (std::vector<pixelate>::iterator it = commands.begin(); it != commands.end(); it++)
	{
		pixelate& pix = *it;
		if(pix.size == 0)
			pix.size = 6;
		for(y = pix.y; y < pix.y + pix.height; y += pix.size)
			for(x = pix.x; x < pix.x + pix.width; x += pix.size)
			{
				int pelcount = 0;
				lsum = csum1 = csum2 = 0;
				for(y1 = y; y1 < y + pix.size && y1 < height; y1++)
					for(x1 = x; x1 < x + pix.size && x1 < width; x1++)
					{
						++pelcount;
						lsum += input[0][y1 * width + x1];
						csum1 += input[1][(y1 / 2) * (width / 2) + x1 / 2];
						csum2 += input[2][(y1 / 2) * (width / 2) + x1 / 2];
					}
				avgLuma = lsum / pelcount;
				avgChroma1 = csum1 / pelcount;
				avgChroma2 = csum2 / pelcount;
				for(y1 = y; y1 < y + pix.size && y1 < height; y1++)
					for(x1 = x; x1 < x + pix.size && x1 < width; x1++)
					{
						output[0][y1 * width + x1] = avgLuma;
						output[1][(y1 / 2) * (width / 2) + x1 / 2] = avgChroma1;
						output[2][(y1 / 2) * (width / 2) + x1 / 2] = avgChroma2;
					}
			}
	}
	++framenum;
}

std::vector<pixelate> ReadCommandFile(const char *pFilename)
{
	char buffer[500];
	std::vector<pixelate> commands;
	
	// for command file format, see Usage()	
	std::ifstream in(pFilename, std::ios::in);
	if(in)
		while(!in.eof())
		{
			memset(buffer, 0, sizeof(buffer));
			in.getline(buffer, sizeof(buffer) - 1);
			char *pX = strtok(buffer, " \t");
			char *pY = strtok(NULL, " \t");
			char *pWidth = strtok(NULL, " \t");
			char *pHeight = strtok(NULL, " \t");
			char *pSize = strtok(NULL, " \t");
			char *pFrameStart = strtok(NULL, " \t");
			char *pFrameEnd = strtok(NULL, " \t");
			if(!pX || !pY || !pWidth || !pHeight || !pFrameStart | !pFrameEnd)
				continue;	// not a valid command line, skip it
			pixelate p;
			p.x = atoi(pX);
			p.y = atoi(pY);
			p.width = atoi(pWidth);
			p.height = atoi(pHeight);
			p.size = atoi(pSize);
			p.startframe = atol(pFrameStart);
			p.endframe = atol(pFrameEnd);
			commands.push_back(p);
		}
	return commands;
}

static void Usage(char *name )
{
	fprintf(stderr,
			"Usage: %s: [-h] [-b num] [-v num]\n"
			"-h   - Print out this help\n"
			"-l   - filename with list of pixelates, one per line, in the\n"
			"       form \"x y width height pixelsize framestart frameend\"\n"
			"       lines not in that format will be ignored.\n"
			"       pixelsize is how many pixels to put in a block. Default is 6.\n"
			"       Pixelation regions may overlap.\n"
			"-v   - Verbosity [0..2]\n" , name);
}

int main(int argc, char *argv[])
{
	unsigned char   *input_frame[3];
	unsigned char   *output_frame[3];
	y4m_stream_info_t istream, ostream;
	y4m_frame_info_t iframe;
	int output_fd = 1;
	int i;
	int input_fd = 0;
	int horz;
	int vert;
	int c;
	int chromasize;
	std::vector<pixelate> commands;
	std::vector<pixelate> active_commands;
	std::vector<long> framestarts;
	std::vector<long> frameends;

	while ((c = getopt(argc, argv, "l:v:h")) != EOF)
	{
		switch (c)
		{
			case 'l':
				commands = ReadCommandFile(optarg);
				break;
			case 'v':
				verbose = atoi (optarg);
				if ( verbose < 0 || verbose >2 )
				{
					Usage (argv[0]);
					exit (1);
				}
				break;        

			case 'h':
				Usage (argv[0]);
			default:
				exit(0);
		}
	}

	if(commands.size() < 1)
	{
		Usage(argv[0]);
		mjpeg_error_exit1("\n** ERROR ** :  No commands given\n");
	}
	for(std::vector<pixelate>::iterator it = commands.begin(); it != commands.end(); it++)
		if(find(framestarts.begin(), framestarts.end(), it->startframe) == framestarts.end())
			framestarts.push_back(it->startframe);
	std::sort(framestarts.begin(), framestarts.end());

	(void)mjpeg_default_handler_verbosity(verbose);

	y4m_init_stream_info(&istream);
	y4m_init_stream_info(&ostream);
	y4m_init_frame_info(&iframe);

	i = y4m_read_stream_header(input_fd, &istream);
	if (i != Y4M_OK)
		mjpeg_error_exit1("Input stream error: %s\n", y4m_strerr(i));

	horz = istream.width;
	vert = istream.height;
	mjpeg_debug( "width=%d height=%d\n", horz, vert);

	y4m_copy_stream_info(&ostream, &istream);

	chromasize = (horz / 2) * (vert / 2);
	input_frame[0] = (unsigned char *)alloca(horz * vert);
	input_frame[1] = (unsigned char *)alloca(chromasize);
	input_frame[2] = (unsigned char *)alloca(chromasize);

	output_frame[0] = (unsigned char *)alloca(horz * vert);
	output_frame[1] = (unsigned char *)alloca(chromasize);
	output_frame[2] = (unsigned char *)alloca(chromasize);

	y4m_write_stream_header(output_fd, &ostream);

	long frame_num = 1;
	// use this iterator to avoid having to walk the command array too much
	std::vector<long>::iterator it = framestarts.begin();
	while (y4m_read_frame(input_fd, &istream, &iframe, input_frame) == Y4M_OK)
	{
		// add any commands that start on this frame
		if(*it == frame_num)
		{
			std::vector<pixelate>::iterator cit = commands.begin();
			while (cit != commands.end())
			{
				if(cit->startframe == frame_num)
				{
					active_commands.push_back(*cit);
					if(find(frameends.begin(), frameends.end(), cit->endframe) == frameends.end())
						frameends.push_back(cit->endframe);
				}
				cit++;
			}
			sort(frameends.begin(), frameends.end());
			it++;	// move this to wait for next start frame
		}
		filter(horz, vert, active_commands, input_frame, output_frame);
		y4m_write_frame(output_fd, &ostream, &iframe, output_frame);
		// check for commands expiring on this frame
		if(frameends.size() && frameends[0] == frame_num)
		{
			std::vector<pixelate>::iterator endit = active_commands.begin();
			while (endit != active_commands.end())
			{
				if(endit->endframe == frame_num)
				{
					endit = active_commands.erase(endit);
				}
				else
					endit++;
			}
			// remove this end frame trigger
			frameends.erase(frameends.begin());
		}
		mjpeg_info("\rframes=%ld", frame_num++);
	}

	mjpeg_info("\nProcessed %ld frames\n", --frame_num);
	y4m_fini_stream_info(&istream);
	y4m_fini_stream_info(&ostream);
	y4m_fini_frame_info(&iframe);
	exit(0);
}