/*
 ** 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.
 **
 ** 
 gcc -g libav-beatcolour.c -o beatcolour -I/sw/include -I/usr/local/include/ffmpeg -L/sw/lib -L/usr/local/lib -lavcodec -lavformat -lfftw3
 gcc -g libav-beatcolour.c -o beatcolour -I/opt/local/include -L/opt/local/lib -lavcodec -lavformat -lfftw3 -lpng
 
 */

#include	<stdio.h>
#include	<stdlib.h>
#include	<string.h>
#include	<ctype.h>
#include	<unistd.h>
#include	<math.h>

#include <png.h>
#include	<fftw3.h>
#include	<ffmpeg/avcodec.h>
#include	<ffmpeg/avformat.h>


#define	BUFFER_LEN      65536
#define HORZ 1280 // this should be dynamic
#define SPP 128
#define FFTSIZE 2048

#define INTENSITY 0.5

#if (defined (WIN32) || defined (_WIN32))
#define	snprintf	_snprintf
#endif

static uint8_t	*data;
static uint8_t	*imagedata;


static png_structp png_ptr;
static png_infop info_ptr;


static void
print_version (void)
{	char buffer [256] ;
	
	//	sf_command (NULL, SFC_GET_LIB_VERSION, buffer, sizeof (buffer)) ;
	printf ("\nVersion : %s\n\n", buffer) ;
} /* print_usage */


static void
print_usage (char *progname)
{	printf ("\nUsage : %s [-c 0|1] [-w <width>] [-f <sample size>] <file> <bpm>\n", progname) ;
	printf ("\nGenerates a beat graph in PPM format on stdout for <file> with <bpm>.\n\n") ;
#if (defined (_WIN32) || defined (WIN32))
	printf ("This is a Unix style command line application which\n"
			"should be run in a MSDOS box or Command Shell window.\n\n") ;
	printf ("Sleeping for 5 seconds before exiting.\n\n") ;
	
	/* This is the officially blessed by microsoft way but I can't get
	 ** it to link.
	 **     Sleep (15) ;
	 ** Instead, use this:
	 */
	_sleep (5 * 1000) ;
#endif
} /* print_usage */


#define HUE 1 	/* 1 for HSV colorspace fade, 0 for RGB colorspace fade */

#define NOHUE -1

void rgb2hsv(int r,int g,int b,double *hr, double *sr, double *vr)
{
	double rd, gd, bd, h, s, v, max, min, del, rc, gc, bc;
	
	/* convert RGB to HSV */
	rd = r / 255.0;            /* rd,gd,bd range 0-1 instead of 0-255 */
	gd = g / 255.0;
	bd = b / 255.0;
	
	/* compute maximum of rd,gd,bd */
	if (rd>=gd) { if (rd>=bd) max = rd;  else max = bd; }
	else { if (gd>=bd) max = gd;  else max = bd; }
	
	/* compute minimum of rd,gd,bd */
	if (rd<=gd) { if (rd<=bd) min = rd;  else min = bd; }
	else { if (gd<=bd) min = gd;  else min = bd; }
	
	del = max - min;
	v = max;
	if (max != 0.0) s = (del) / max;
	else s = 0.0;
	
	h = NOHUE;
	if (s != 0.0) {
		rc = (max - rd) / del;
		gc = (max - gd) / del;
		bc = (max - bd) / del;
		
		if      (rd==max) h = bc - gc;
		else if (gd==max) h = 2 + rc - bc;
		else if (bd==max) h = 4 + gc - rc;
		
		h = h * 60;
		if (h<0) h += 360;
	}
	
	*hr = h;  *sr = s;  *vr = v;
}


/* 0 <= h,s,v, <= 1.0 */
void hsv2rgb(double h, double s, double v, double *rd,double *gd,double *bd)
{
	double    j;
	double f, p, q, t;
	
	/* convert HSV back to RGB */
	if (s==0.0) { *rd = v;  *gd = v;  *bd = v; }
	else {
		while (h > 1.0) h -= 1.0;
		while (h < 0) h += 1.0;
		if (h==1.0) h = 0.0;
		
		j =  floor(h*6.0);
		//h = h * 6.0;
		// if (j<0) j=0;          /* either h or floor seem to go neg on some sys */
		
		p = v * (1-s);
		f = h * 6.0 - j;
		q = v * (1 - (s*f));
		t = v * (1 - (s*(1 - f)));
		
		if (j < 0.5) { *rd = v;  *gd = t;  *bd = p;  }
		else if (j < 1.5) {  *rd = q;  *gd = v;  *bd = p; }
		else if (j < 2.5) {  *rd = p;  *gd = v;  *bd = t; }
		else if (j < 3.5) {  *rd = p;  *gd = q;  *bd = v; }
		else if (j < 4.5) {  *rd = t;  *gd = p;  *bd = v; } 
		else if (j < 5.5) {  *rd = v;  *gd = p;  *bd = q; } 
		else { *rd = v;  *gd = t;  *bd = p; }  /* never happen */
	}
	
}

void calccolour (int p, double *in, double *r, double *g, double *b)
{
	
	int c; 
	double h,s,v,dc,tt,mx;
	double tr,tg,tb;
	double ttr,ttg,ttb,sc;
	
	ttr = 0; ttg=0; ttb=0;
	h=0; s=1;
	
	dc = in[0];
	sc = 1.0 / p;
	
	for (c = 1; c <= p; c++)
	{
		
		h =  c * sc;
		
		// want only the amplitude portion
		// v = (fabs(c_re(in[c])) + fabs(c_im(in[c])));
		//	fprintf(stderr,"spec: %lf %lf\n",in[c],in[c+p]);
		// v = (fabs(in[c]) + fabs(in[c+p-1]));
		v = fabs(in[c]);
		/* test the colour function */
		v = v / 128;
		
		hsv2rgb(h,1,v,&tr,&tg,&tb);
		//fprintf (stderr,"c: %d h: %lf v: %lf  r: %lf g: %lf b: %lf\n",c,h,v,tr,tg,tb);
		ttr += tr; ttg += tg; ttb += tb;
		//fprintf (stderr,"r: %lf g: %lf b: %lf\n",*r,*g,*b);
	}
	
	
	/* GG suggested a log of the intensity */
	/*  0.0 <= ttr,ttg,ttb <= 1000.0 (maybe) */
	/*
	 if (ttr > 1.0) ttr = 16* log2(ttr);
	 if (ttg > 1.0) ttg = 16* log2(ttg);
	 if (ttb > 1.0) ttb = 16* log2(ttb);
	 */
	*r = ttr; *g = ttg; *b = ttb;
	
	//fprintf (stderr,"r: %lf g: %lf b: %lf\n",*r,*g,*b);
	
}

// was hoping to make a read single bar, but overlaps between bar length and packet length caused too many issues.

// read the entire file into memory...

void * load_file (AVFormatContext *pfc, int s, unsigned int *total_data_size) {
	
	
	AVCodecContext  *pCodecCtx;
    AVCodec         *pCodec;
    AVPacket        packet;
    uint8_t         *buffer;
    static short *samples= NULL;
    uint8_t *ptr, samp;
	int data_size;
	int count;
	
	
	*total_data_size=0;
	
	samples = (short *) malloc(AVCODEC_MAX_AUDIO_FRAME_SIZE);
	data_size = AVCODEC_MAX_AUDIO_FRAME_SIZE;
	
	//fprintf (stderr,"size of: %d\n", data_size);
	
	pCodecCtx=pfc->streams[s]->codec;
	
	pCodec=avcodec_find_decoder(pCodecCtx->codec_id);
    if(pCodec==NULL)
        return NULL; // Codec not found
	
	// Open codec
    if(avcodec_open(pCodecCtx, pCodec)<0)
        return NULL; // Could not open codec
	
	buffer = (uint8_t *) malloc(64);
	
	while(av_read_frame(pfc, &packet)>=0)
    {
		/*
		 printf("total dur: %d ",pfc->streams[s]->duration);
		 printf("size: %d ",packet.size);
		 printf("dur: %d ",packet.duration);
		 */	
		
		avcodec_decode_audio2(pCodecCtx, samples, &data_size, packet.data, packet.size);
		
		//	printf ("samples decoded: %d\n",data_size);
		
		//printf ("realloc %d\n",(*total_data_size) + data_size);
		
		buffer =(uint8_t *) realloc(buffer,(*total_data_size) + data_size/2);
		
		data_size /= 2;
		for (count=0;count<data_size;count+=2)  {
			
			//fprintf (stderr,"read: %d, %d\n",samples[count],samples[count+1]);
			
			//	samp = (samples[count] + samples[count+1]) / 512 + 128;
			samp = samples[count] / 256;
			
			//	printf ("%d: (%hd + %hd) / 512 = %d\n", *total_data_size, samples[count], samples[count+1], samp);
			//  printf ("%d %d\n", *total_data_size, samples[count] + samples[count+1]));
			
			// This debug was to determine the file reading and conversion is correct.
			//	putchar(samp);
			
			buffer[((*total_data_size) + count)/2] = samp;
		}
		
		*total_data_size += data_size;
		data_size = AVCODEC_MAX_AUDIO_FRAME_SIZE;
	}
	
	*total_data_size /= 2;
	//printf ("total size: %d\n",*total_data_size);
	
	return buffer;
	
}

// open the audio file using an external library

void * open_file(char *fn, unsigned int *dl, unsigned int *sr) 
{
	unsigned char * da = NULL;
	AVFormatContext *pFormatCtx;
    int             audioStream,i;
	
	
	av_register_all();
	
	if(av_open_input_file(&pFormatCtx, fn, NULL, 0, NULL)!=0) {
		printf("Error: opening file %d.\n",fn);
		return NULL; 
	}
	
	if(av_find_stream_info(pFormatCtx)<0) {
		printf("Error: couldn't find stream information in file %d.\n",fn);
		return NULL; // Couldn't find stream information
	}
	
	audioStream = -1;
	
	for(i=0; i<pFormatCtx->nb_streams; i++) {
		//fprintf (stderr,"stream: %d = %d (%s)\n",i,pFormatCtx->streams[i]->codec->codec_type ,pFormatCtx->streams[i]->codec->codec_name);
		if(pFormatCtx->streams[i]->codec->codec_type==CODEC_TYPE_AUDIO)
		{
			audioStream=i;
			*sr = pFormatCtx->streams[i]->codec->sample_rate;
			//	*dl = pFormatCtx->streams[i]->nb_frames;
			
			//printf ("Duration: %d.\n",pFormatCtx->streams[i]->duration);
			//printf ("Frames: %u.\n",*dl);
			fprintf (stderr,"SampleRate: %u.\n",*sr);
			
			/*
			 da = (unsigned char *)malloc(*dl);
			 if (da) {
			 load_file(infile,da);
			 }
			 */
			// we can't seem to get duration until we read the first frame...
			fprintf (stderr,"Reading file.\n");
			da = load_file(pFormatCtx,audioStream,dl);
			return da;
		}
	}
	
	return NULL;
	
}

void 
write_close_png (uint8_t *ida, unsigned int horz, unsigned int vert) 
{
	
	int k;
	png_byte **row_pointers;
	
	row_pointers = (png_byte **) malloc(vert * sizeof (png_byte));
	
	for (k=0; k<vert; k++) {
		row_pointers[k] = ida + k * horz * 3;
	}
	
	png_write_image(png_ptr, row_pointers);
	png_write_end(png_ptr, info_ptr);			
	png_destroy_write_struct(&png_ptr, &info_ptr);
}


int 
init_png_writer (unsigned int width, unsigned int height, FILE *fp) {
	
		
	png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, (png_voidp)NULL, png_error_ptr_NULL, png_error_ptr_NULL);
	if (!png_ptr)
		return -1;
	
	info_ptr = png_create_info_struct(png_ptr);
	if (!info_ptr)
	{
		png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
		return -1;
	}
	png_set_compression_level(png_ptr, Z_BEST_COMPRESSION);
	// fprintf(stderr,"w: %d h: %d\n",width,height);
	png_set_IHDR(png_ptr, info_ptr, width, height,
				 8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
				 PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
	
	png_init_io(png_ptr, fp);
	
	png_write_info(png_ptr, info_ptr);
}					 

void
make_graph_bw (uint8_t *ida, uint8_t *da, unsigned int dl, double v, unsigned int width,int fftsize)
{
	double d,tot;
	int c,f,pixelcount=0;
	double e = 1.0 * v / width ;
	
	// fprintf (stderr,"samples per point %lf out of %d\n",e,dl);
	
	for (d=0; d < dl; d+=e) {
		tot = 0;
		for (c=0; c<fftsize; c++) {
			f = c + d + 0.5;
			tot += da[f];
		}
		tot /= fftsize;
		f = tot;
		
		//		printf ("%c%c%c",f,f,f);
		
		// fprintf (stderr,"pixelcount: %d\n",pixelcount);
		*(ida+pixelcount++) = f; 
		*(ida+pixelcount++) = f; 
		*(ida+pixelcount++) = f; 
		
	}
	
	
}


void
make_graph (uint8_t *ida, uint8_t *da, unsigned int dl, double v, unsigned int width, int fftsize)
{	
	
	/* REFACTOR make_graph 
	 ** read line, v = number of samples 
	 ** for x=0 to width
	 ** pixel colour = fft(x*(v/w),fftsize)
	 */
	
	
	int readcount,pixelcount=0;
	static double graph[HORZ][3];
	double *out;
	double *in;
	
	fftw_plan plan;
	int val=0; 
	unsigned int c,f;
	double d,e,rr,rg,rb,max,min=9999999;
	int ir, ig, ib;
	
	// v : samples per bar (4 beats)
	// e : pixels per bar/line
	
	e =  v / width ;
	out = fftw_malloc(fftsize * (sizeof(double)));
	in = fftw_malloc(fftsize * (sizeof(double)));
	
	//	fprintf(stderr,"memalloc: %x %x\n",out,in);
	
	if (out==NULL || in==NULL) {
		fprintf(stderr,"CANNOT ALLOCATE FFTW MEMORY\n");
		exit(-1);
	}
	
	// fprintf (stderr,"%lf steps between fft groups (%d samples)\n",e,fftsize);
	
	//	plan = fftw_create_plan(fftsize,FFTW_FORWARD,FFTW_MEASURE);
	plan = fftw_plan_r2r_1d(fftsize, in, out, FFTW_R2HC, FFTW_ESTIMATE); 
	
	/* d : horizontal pixel counter */
	for (d=0; d < dl; d+=e)
	{
		for (c=0; c<fftsize; c++)
		{
			if ( c+d < dl)
			{
				// scale sample data into fftsize
				// fprintf (stderr,"read c: [%d : %d] %d\n",c,d,da[c+d]);
				// it would be nice to put a windowing function here
				f = c + d;
				in[c]  = da[f];
				// printf ("read c: %d d: %d v: %lf = %d\n",c,d,in[c/2],samples[c+d]);
			} else {
				for (; c<fftsize; c++)
					in[c] = 0;
			}
		}	
		
		fftw_execute(plan);		
		calccolour(fftsize/2-1, out, &rr,&rg,&rb);
		
		if (rr > max) max=rr;
		if (rg > max) max=rg;
		if (rb > max) max=rb;
		
		if ((rr < min) && (rr>0)) min=rr;
		if ((rg < min) && (rg>0)) min=rg;
		if ((rb < min) && (rb>0)) min=rb;
		
		ir = rr * INTENSITY; ir = ir>255?255:ir;
		ig = rg * INTENSITY; ig = ig>255?255:ig;
		ib = rb * INTENSITY; ib = ib>255?255:ib;
		
		*(ida+pixelcount++) = ir; 
		*(ida+pixelcount++) = ig; 
		*(ida+pixelcount++) = ib; 
	}
	
	fftw_free(out);
	fftw_free(in);
	
	fprintf (stderr,"colour range (%lf - %lf)\n",min,max);
	
} 


int main (int argc, char *argv[])
{	
	static	char	strbuffer [BUFFER_LEN] ;
	char *progname, *infilename ;
	
	AVFormatContext *infile;
	unsigned int audioStream;
	
	int			k,horz ;
	unsigned int vert,sr,fr;
	double delay,sdelay;
	const static char *legal_flags = "c:w:f:h:o:";
	int algorithm=0,width=640,fftsize=256;
	FILE *outfile = NULL;
	
	progname = strrchr (argv [0], '/') ;
	progname = progname ? progname + 1 : argv [0] ;
	
	while ((k = getopt (argc, argv, legal_flags)) != -1) {
		switch (k) {
			case 'c':
				algorithm = atoi(optarg);
				break;
			case 'w':
				width = atoi(optarg);
				break;
			case 'f':
				fftsize = atoi(optarg);
				break;
			case 'o':
				outfile = fopen(optarg,"w");
				break;
			case 'h':
			default:
				print_usage (progname) ;
				exit (0);
		}
	}
	
	optind--;
	argc -= optind;
	argv += optind;
	
	if (argc < 3)
	{	print_usage (progname) ;
		return  1 ;
	} 
	
	//	for (k = 1 ; k < argc ; k+=2)
	//	{	
	
	// FIXME: should sanity check these variables.
	
	infilename = argv[1];
	delay = atof(argv[2]);
	if (!outfile) {
		fprintf(stderr,"No outfile selected writing to STDOUT\n");
		outfile = stdout;
	}
	
	data = open_file (infilename,&fr,&sr) ;
	
	if (! data)
	{	
		fprintf (stderr,"Error : Not able to open input file %s.\n", infilename);
		fflush (stdout);
		memset (data, 0, sizeof (data));	
	} else {	
		
		
		// get sample rate from infile
		sdelay = sr * 4.0 * 60 / delay;
		
		fprintf (stderr,"fr: %u\ndelay: %lf\n",fr,sdelay);
		
		//horz = 1 + (sdelay / SPP) ;
		
		horz = width;
		vert = (fr / sdelay) + 1; // 1 extra incase of rounding errors.
		
		init_png_writer(horz,vert,outfile);
		imagedata = (uint8_t * ) malloc ( horz * vert * 3);
		
		switch (algorithm)
		{
			case 1:
				make_graph (imagedata,data,fr,sdelay,horz,fftsize);
				break;
			default:
				make_graph_bw (imagedata, data,fr,sdelay,horz,fftsize);
				break;
		}
		
		write_close_png (imagedata,  horz,  vert); 
		
		
		//fprintf (stderr,"Signal Max  : %g\n\n", get_signal_max (infile)) ;
		
		fprintf (stderr,"data : %x\n", data);
		if (data) {
			free (data);
		}
		
		fclose (outfile);
	} 
	
	return 0 ;
} /* main */