// MDC 0.1 - (Motion Dependent Clip)  
// Copyright 2004 Carl Ritson <critson@perlfu.co.uk>

// 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., 675 Mass Ave, Cambridge, MA 02139, USA, or visit
// http://www.gnu.org/copyleft/gpl.html .

#include <stdio.h>
#include <windows.h>
#include <math.h>
#include "avisynth.h"
#include "info.h"

#include "scene-change.h"

#define MAX_CLIPS 3
#define MAX_SEARCH_FRAMES 20

enum {
	DF_MASK		= 0xff00,
	DF_CMASK	= 0x00ff,
	DF_INIT		= 0x0000,
	DF_DIFF_SET 	= 0x0100,
	DF_CLIP_SET	= 0x0200,
	DF_IS_SC	= 0x0400,
	DF_SC_SET	= 0x0800
};

// Structures

struct mclip {
	PClip 	clip;
	float 	t_lower;
	float 	t_upper;
};

struct mdecision {
	int	flags;
	float	diff;
	float	true_diff;
	float	avg_diff;
	int	map_to;
};

class MDC : public GenericVideoFilter {
	public: 
		MDC(PClip _child, 
			PClip c1, float t1_lower, float t1_upper,
			PClip c2, float t2_lower, float t2_upper,
			PClip c3, float t3_lower, float t3_upper,
			int _defclip,
			int _pre_frames, int _post_frames,
			bool _sc_combined_search, float _sc_threshold,
			bool _scf_fix, float _scf_threshold,
			int _sc_algorithm, int _sc_search,
			bool _debug,
			IScriptEnvironment *env); 
		PVideoFrame __stdcall GetFrame(int n, IScriptEnvironment *env); 
	private:
		void 		CalculateDiff(int n, IScriptEnvironment *env);
		void 		SetDiff(int n, float diff, float true_diff);
		void 		SetDiff(int n, float diff);
		float		GetDiff(int n, IScriptEnvironment *env);
		float		GetTrueDiff(int n, IScriptEnvironment *env);
		bool		IsSC(int n, IScriptEnvironment *env);
		PVideoFrame	OutputDebugInfo(PVideoFrame &frame,
					int n, IScriptEnvironment *env);
		void		MakeDecision(int n, IScriptEnvironment *env);
		
		VideoInfo	original_vi;
		mclip 		clips[3];
		mdecision 	*decision;
		int 		defclip, pre_frames, post_frames;
		int		sc_algorithm, sc_search;
		bool 		debug, sc_combined_search, scf_fix;
		float 		sc_threshold, scf_threshold;
};

// Methods

MDC::MDC(PClip _child,
		PClip c1, float t1_lower, float t1_upper,
		PClip c2, float t2_lower, float t2_upper,
		PClip c3, float t3_lower, float t3_upper,
		int _defclip,
		int _pre_frames, int _post_frames,
		bool _sc_combined_search, float _sc_threshold,
		bool _scf_fix, float _scf_threshold,
		int _sc_algorithm, int _sc_search,
		bool _debug,
		IScriptEnvironment *env) 
		: GenericVideoFilter(_child),
		defclip(_defclip),
		pre_frames(_pre_frames), post_frames(_post_frames),
		sc_combined_search(_sc_combined_search),
		sc_threshold(_sc_threshold), 
		scf_fix(_scf_fix), scf_threshold(_scf_threshold),
		sc_algorithm(_sc_algorithm), sc_search(_sc_search),
		debug(_debug) {
	VideoInfo target_vi;
	int i;
	
	for(i = 0; i < MAX_CLIPS; ++i)
		clips[i].clip = 0;
	
	// there must be at least one clip
	clips[0].clip = c1;
	clips[0].t_lower = t1_lower;
	clips[0].t_upper = t1_upper;
	if(c2 != 0) {
		clips[1].clip = c2;
		clips[1].t_lower = t2_lower;
		clips[1].t_upper = t2_upper;
	}
	if(c3 != 0) {
		clips[2].clip = c3;
		clips[2].t_lower = t3_lower;
		clips[2].t_upper = t3_upper;
	}
	
	target_vi = c1->GetVideoInfo();
	for(i = 0; i < MAX_CLIPS; ++i) {
		VideoInfo cvi = clips[i].clip->GetVideoInfo();
		if(cvi.height != target_vi.height)
			env->ThrowError("MDC: All output sources must have the same height!");
		if(cvi.width != target_vi.width)
			env->ThrowError("MDC: All output sources must have the same width!");
		if(!cvi.IsSameColorspace(target_vi))
			env->ThrowError("MDC: All output sources must have the same colourspace!");
		if(cvi.num_frames != target_vi.num_frames)
			env->ThrowError("MDC: All output sources must have the same number of frames!");
	}
	
	defclip -= 1;
	if(defclip < 0 || defclip >= MAX_CLIPS)
		env->ThrowError("MDC: Invalid default clip!");
	if(clips[defclip].clip == 0)
		env->ThrowError("MDC: Invalid default clip!");
	
	if(vi.num_frames != target_vi.num_frames)
		env->ThrowError("MDC: determination source must have the same number of frames other clips!");
	if(!vi.IsYV12())
		env->ThrowError("MDC: requires YV12 determination source");
	
	original_vi 	= vi;
	vi.height	= target_vi.height;	
	vi.width 	= target_vi.width;
	vi.pixel_type	= target_vi.pixel_type;
	vi.image_type	= target_vi.image_type;

	if(pre_frames < 0 || post_frames < 0)
		env->ThrowError("MDC: pre_frames and post_frames must be >= 0");
	if((pre_frames + post_frames) > MAX_SEARCH_FRAMES)
		env->ThrowError("MDC: pre_frames + post_frames must be <= 20");
	if(sc_search > MAX_SEARCH_FRAMES)
		env->ThrowError("MDC: sc_search, must be <= 20");
	if(sc_algorithm < 0 || sc_algorithm > 2)
		env->ThrowError("MDC: sc_algorithm must be >= 0 and <= 2");
	
	decision = new mdecision[vi.num_frames];
	for(i = 0; i < vi.num_frames; ++i) {
		decision[i].flags	= DF_INIT;
		decision[i].diff	= 0.0;
		decision[i].true_diff	= 0.0;
		decision[i].avg_diff	= 0.0;
		decision[i].map_to	= i;
	}
}

void MDC::CalculateDiff(int n, IScriptEnvironment *env) {
	if(n == 0) {
		SetDiff(n,0.0,0.0);
		return;
	}
	
	bool ISSE = !!(env->GetCPUFlags() & CPUF_INTEGER_SSE);
	
	PVideoFrame src1 = child->GetFrame(n-1,env);
	PVideoFrame src2 = child->GetFrame(n,env);
	
	float diff = 0.0;
	int plane = PLANAR_Y;
	
	for(;;) {
		const BYTE* srcp1 = src1->GetReadPtr(plane);
		const BYTE* srcp2 = src2->GetReadPtr(plane);

		int h = src1->GetHeight(plane);
		int w = src1->GetRowSize(plane);
		int pitch1 = src1->GetPitch(plane);
		int pitch2 = src2->GetPitch(plane);
		int b;
		
		w=(w/16)*16;

		if(ISSE)
			b = isse_scenechange_16(srcp1, srcp2, h, w, pitch1, pitch2);
		else
			b = C_scenechange_16(srcp1, srcp2, h, w, pitch1, pitch2);

		if(plane == PLANAR_Y)
			diff += (float)(b != 0 ? b : 1) / (float)(h * w) * 0.5f;
		else
			diff += (float)(b != 0 ? b : 1) / (float)(h * w) * 0.25f;
		if(plane == PLANAR_Y)
			plane = PLANAR_U;
		else if(plane == PLANAR_U)
			plane = PLANAR_V;
		else
			break; // leave loop
	}
	
	SetDiff(n,diff,diff);
}

void MDC::SetDiff(int n, float diff, float true_diff) {
	if(n < 0 || n >= vi.num_frames)
		return;
	decision[n].flags	|= DF_DIFF_SET;
	decision[n].diff 	= diff;
	decision[n].true_diff 	= diff;
}

void MDC::SetDiff(int n, float diff) {
	if(n < 0 || n >= vi.num_frames)
		return;
	decision[n].flags	|= DF_DIFF_SET;
	decision[n].diff	= diff;
}

float MDC::GetDiff(int n, IScriptEnvironment *env) {
	if(n < 0 || n >= vi.num_frames)
		return 0.0;
	if(!(decision[n].flags & DF_DIFF_SET))
		CalculateDiff(n,env);
	return decision[n].diff;
}

float MDC::GetTrueDiff(int n, IScriptEnvironment *env) {
	if(n < 0 || n >= vi.num_frames)
		return 0.0;
	if(!(decision[n].flags & DF_DIFF_SET))
		CalculateDiff(n,env);
	return decision[n].true_diff;
}

bool MDC::IsSC(int n, IScriptEnvironment *env) {
	if(n == 0)
		decision[n].flags |= DF_IS_SC | DF_SC_SET;
	if(n < 0 || n >= vi.num_frames)
		return false;
	if(!(decision[n].flags & DF_SC_SET)) {
		if(sc_algorithm == 0) {
			if(GetTrueDiff(n,env) > sc_threshold)
				decision[n].flags |= DF_IS_SC;
		} else if(sc_algorithm == 1 || sc_algorithm == 2) {
			double no_elem = (sc_search/2) + (sc_algorithm==1?1:0);
			double sum, diff = GetTrueDiff(n,env);
			int i;
			
			for(i = 1, sum = 0.0; i < sc_search/2; ++i) {
				sum += GetTrueDiff(n-i,env);
				sum += GetTrueDiff(n+i,env);
			}
			if(sc_algorithm == 1)
				sum += diff;
			
			double mean = sum / no_elem;
			for(i = 1, sum = 0.0; i < sc_search/2; ++i) {
				float diff1 = GetTrueDiff(n-i,env);
				float diff2 = GetTrueDiff(n+i,env);
				sum += (diff1 - mean) * (diff1 - mean);
				sum += (diff2 - mean) * (diff2 - mean);
			}
			if(sc_algorithm == 1)
				sum += (diff - mean) * (diff - mean);			
			double sd = sqrt(sum / no_elem);
			if(diff > (mean + (sd * (double)sc_threshold)))
				decision[n].flags |= DF_IS_SC;
		}
		decision[n].flags |= DF_SC_SET;
	}
	return (decision[n].flags & DF_IS_SC) == DF_IS_SC;
}

PVideoFrame MDC::OutputDebugInfo(PVideoFrame &frame, int n, IScriptEnvironment *env) {
	char buffer[80];
	
	env->MakeWritable(&frame);
	
	_snprintf(buffer,sizeof(buffer),
			"% 5d, diff: %.2f, true_diff: %.2f, avg_diff: %.2f",
			n,
			decision[n].diff,
			decision[n].true_diff,
			decision[n].avg_diff);
	DrawString(frame,0,0,buffer);
	
	_snprintf(buffer,sizeof(buffer),
			"% 5d, clip: %d, sc: %s,%s",
			decision[n].map_to,
			(decision[n].flags & DF_CMASK)+1,
			(decision[n].flags & DF_SC_SET) ? "true" : "false",
			(decision[n].flags & DF_IS_SC) ? "true" : "false");
	DrawString(frame,0,1,buffer);

	return frame;
}

void MDC::MakeDecision(int n, IScriptEnvironment *env) {
	float 	diff_sum = 0.0;
	int	backward = pre_frames, forward = post_frames;
	int	b_searched = 0, f_searched = 0;
	int 	i, j, frames = 0;
	
	if(IsSC(n,env)) {
		SetDiff(n,0.0);
		if(scf_fix) {
			if(GetDiff(n+1,env) < scf_threshold && 
					(n+1) < vi.num_frames) {
				decision[n].map_to = n+1;
				SetDiff(n+1,0.0);
			}
		}
		backward = 0;
	}

	diff_sum += GetDiff(n,env);
	frames++;
	
	for(j = 0; j < 2; ++j) {
		for(i = b_searched+1; i < backward; ++i) {
			if((n-i) < 0 || IsSC(n-i,env)) {
				b_searched++;
				break;
			}
			diff_sum += GetDiff(n-i,env);
			frames++;
		}

		for(i = f_searched+1; i < forward; ++i) {
			if((n+i) >= vi.num_frames || IsSC(n+i,env)) {
				f_searched++;
				break;
			}
			diff_sum += GetDiff(n+i,env);
			frames++;
		}

		if(!sc_combined_search)
			break;
		if(f_searched == post_frames && b_searched == pre_frames)
			break;
		
		backward += (f_searched - post_frames);
		forward += (b_searched - pre_frames);
	}	
	
	float avg_diff = diff_sum / ((float)frames);
	int clip = -1;

	for(i = 0; i < MAX_CLIPS && clip == -1; ++i) {
		if(clips[i].clip != 0 
				&& avg_diff >= clips[i].t_lower
				&& avg_diff < clips[i].t_upper)
			clip = i;
	}
	
	if(clip == -1)
		clip = defclip;

	decision[n].flags = ((decision[n].flags | DF_CLIP_SET) & DF_MASK) | (clip & DF_CMASK);
	decision[n].avg_diff = avg_diff;
}

PVideoFrame __stdcall MDC::GetFrame(int n, IScriptEnvironment *env) {
	if(!(decision[n].flags & DF_CLIP_SET))
		MakeDecision(n,env);
	
	PVideoFrame frame = clips[decision[n].flags & DF_CMASK].clip->GetFrame(decision[n].map_to,env);
	
	if(debug)
		frame = OutputDebugInfo(frame,n,env);
	
	return frame;
}

// Factory Methods

AVSValue __cdecl Create_MDC3(AVSValue args, void*, IScriptEnvironment *env) {
	return new MDC(
		args[0].AsClip(),
		args[1].AsClip(),args[2].AsFloat(0.0),args[3].AsFloat(0.0),
		args[4].AsClip(),args[5].AsFloat(0.0),args[6].AsFloat(0.0),
		args[7].AsClip(),args[8].AsFloat(0.0),args[9].AsFloat(0.0),
		args[10].AsInt(1),
		args[11].AsInt(2),args[12].AsInt(2),
		args[13].AsBool(true),
		args[14].AsFloat(args[17].AsInt(1) == 0 ? 20.0 : 1.5),
		args[15].AsBool(true),args[16].AsFloat(4.0),
		args[17].AsInt(1),args[18].AsInt(10),
		args[19].AsBool(false),
		env
	);
}

// Entry Point

extern "C" __declspec(dllexport) const char* __stdcall AvisynthPluginInit2(IScriptEnvironment* env) {
	env->AddFunction("MDC", 
		"c"
		"[c1]c[t1_lower]f[t1_upper]f"
		"[c2]c[t2_lower]f[t2_upper]f"
		"[c3]c[t3_lower]f[t3_upper]f"
		"[default]i"
		"[pre_frames]i[post_frames]i"
		"[sc_combined_search]b[sc_threshold]f"
		"[scf_fix]b[scf_threshold]f"
		"[sc_algorithm]i[sc_search]i"
		"[debug]b", 
		Create_MDC3, 0); 
	return "Motion Dependent Clip Filter";
}