// ****************************************************************************
// This is a slightly modified version of the BitrateCalculator.cs 
// from MeGUI (http://sourceforge.net/projects/megui).
//
// Copyright (C) 2005  Doom9
// 
// 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; 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
// 
// ****************************************************************************

module BitrateCalculator;

import tango.math.Math;

///Arrays with the storage mediums, and there sizes.
char[][] storageMediums = 
	["1/4 CD", "1/2 CD", "1 CD", "2 CD", "3 CD",
	 "1/3 DVD-R", "1/4 DVD-R", "1/5 DVD-R", "DVD-5", "DVD-9"];
int[] storageMediumSizeKB =
	[179200, 358400, 716800, 1433600, 2150400, 1501184,
	 1126400, 901120, 4586496, 8333760,];

///Arrays with the framerates.
double[] FrameRate = [23.967, 24.0, 25.0, 29.97, 30.0, 50.0, 59.94, 60.0];
char[][] frameRates = ["23.967", "24.0", "25.0", "29.97", "30.0", "50.0", "59.94", "60.0"];

///The audio types used for the calculations.
enum AudioType { MP4AAC, RAWAAC, AC3, DTS, MP2, CBRMP3, VBRMP3, VORBIS }
char[][] audioTypes = ["MP4-AAC", "RAW-AAC", "AC3", "DTS", "MP2", "CBR-MP3", "VBR-MP3", "Ogg"];

///Videocodec used for calculations.
enum VideoCodec	{ Lavc, X264, Snow, XviD }
char[][] videoCodecs = ["Lavc", "X264", "Snow", "XviD"];

///Container used.
enum ContainerType { MKV, MP4, AVI }
char[][] containerTypes = ["MKV", "MP4", "AVI"];

///Info about the audio stream.
struct AudioStream
{
	long SizeBytes;
	AudioType Type;
}

///Bitrate Calculations
public class BitrateCalculator
{
	private const double mp4OverheadWithBframes = 10.4;
	private const double mp4OverheadWithoutBframes = 4.3;
	private const double aviVideoOverhead = 24;
	private const double cbrMP3Overhead = 23.75;
	private const double vbrMP3Overhead = 40;
	private const double ac3Overhead = 23.75;
	private const int AACBlockSize = 1024;
	private const int AC3BlockSize = 1536;
	private const int MP3BlockSize = 1152;
	private const int VorbisBlockSize = 1024;
	private const int mkvAudioTrackHeaderSize = 140;
	private const int mkvVorbisTrackHeaderSize = 4096;
	private const uint mkvIframeOverhead = 26;
	private const uint mkvPframeOverhead = 13;
	private const uint mkvBframeOverhead = 16;

	this(){}

	/// calculates the video bitrate for a video to be put into the MP4 container
	/// Params:
	/// audioStreams = the audio streams to be muxed
	/// desiredOutputSize = the desired size of the muxed output
	/// nbOfFrames = number of frames of the source
	/// useBframes = whether we have b-frames; the video
	/// framerate = framerate of the video
	/// videoSize = size of the raw video stream
	/// Returns: the calculated bitrate
	private int calculateMP4VideoBitrate(AudioStream[] audioStreams, long desiredOutputSizeBytes, ulong nbOfFrames,
		bool useBframes, double framerate, out int videoSizeKB)
	{
		double mp4Overhead = this.getMP4Overhead(useBframes);
		double totalOverhead = cast(double)nbOfFrames * mp4Overhead;
		double nbOfSeconds = cast(double)nbOfFrames / framerate;
		long audioSize = 0;
		foreach (AudioStream stream; audioStreams)
		{
			audioSize += stream.SizeBytes;
		}
		long videoTargetSize = desiredOutputSizeBytes - audioSize - cast(long)totalOverhead;
		videoSizeKB = rndint(videoTargetSize / 1024);
		long sizeInBits = videoTargetSize * 8;
		int bitrate = rndint(sizeInBits / (nbOfSeconds * 1000));
		return bitrate;
	}

	/// calculates the size of a muxed mp4 file given the desired video bitrate and the audio streams to be muxed
	/// Params:
	/// audioStreams = the audio streams to be muxed with the video
	/// desiredBitrate = the desired video bitrate
	/// nbOfFrames = the number of frames of the video source
	/// useBframes = whether the sources uses b-frames
	/// framerate = the framerate of the source
	/// rawVideoSize = the raw video size the stream will have; the container
	/// Returns: the size of the mp4 file; KB
	private long calculateMP4Size(AudioStream[] audioStreams, int desiredBitrate, ulong nbOfFrames, bool useBframes, double framerate, out int rawVideoSize)
	{
		double mp4Overhead = this.getMP4Overhead(useBframes);
		double totalOverhead = cast(double)nbOfFrames * mp4Overhead;
		double nbOfSeconds = cast(double)nbOfFrames / framerate;
		double bytesPerSecond = desiredBitrate * 1000 / 8;
		long videoSize = cast(long)(nbOfSeconds * bytesPerSecond);
		rawVideoSize = rndint(videoSize / 1024);
		long audioSize = 0;
		foreach (AudioStream stream; audioStreams)
		{
			audioSize += stream.SizeBytes;
		}
		long size = videoSize + audioSize + cast(long)totalOverhead;
		return size / 1024L;
	}

	/// calculates the bitrate a video with the given properties needs to have; order to be placed into a matroska file along with the
	/// given audio tracks and end up having the desired size
	/// Params:
	/// audioStreams = te audio streams to be muxed with this video
	/// desiredOutputSize = the final size of this file
	/// nbOfFrames = number of frames of the video
	/// framerate = framerate of the video
	/// useBframes = whether the video uses b-frames
	/// videoSize = size of the raw video stream; KB
	/// Returns: the video bitrate; kbit/s
	private int calculateMKVVideoBitrate(AudioStream[] audioStreams, long desiredOutputSize, ulong nbOfFrames, double framerate, bool useBframes, out int videoSize)
	{
		double totalOverhead = 0.0;
		ulong nbIframes = nbOfFrames / 10;
		ulong nbBframes = 0;
		if (useBframes)
			nbBframes = (nbOfFrames - nbIframes) / 2;
		ulong nbPframes = nbOfFrames - nbIframes - nbBframes;
		totalOverhead = cast(double)(4300 + 1400 + nbIframes * mkvIframeOverhead + nbPframes * mkvPframeOverhead +
			nbBframes * mkvBframeOverhead);
		double nbOfSeconds = cast(double)nbOfFrames / framerate;
		totalOverhead += nbOfSeconds / 2.0 * 12; // 12 bytes per cluster
		long audioSize = 0L;
		double audioOverhead = 0;
		foreach (AudioStream stream; audioStreams)
		{
			audioSize += stream.SizeBytes;
			audioOverhead += getMKVAudioOverhead(stream.Type, 48000, nbOfSeconds);
		}
		long videoTargetSize = desiredOutputSize - audioSize - cast(long)audioOverhead -
			cast(long)totalOverhead;
		videoSize = rndint(videoTargetSize / 1024);
		long sizeInBits = videoTargetSize * 8;
		int bitrate = rndint(sizeInBits / (nbOfSeconds * 1000));
		return bitrate;
	}

	/// calculates what size a given video and audio stream(s) will have at a given video bitrate
	/// Params:
	/// audioStreams = the audio streams to be considered
	/// desiredBitrate = the desired video bitrate
	/// nbOfFrames = number of frames of the video source
	/// framerate = framerate of the video source
	/// useBframes = whether we use b-frames for the video
	/// rawVideoSize = the raw size of the video stream; KB
	/// Returns: the size of the final file; KB
	private long calculateMKVSize(AudioStream[] audioStreams, int desiredBitrate, ulong nbOfFrames, double framerate, bool useBframes, out int rawVideoSize)
	{
		double totalOverhead = 0.0;
		ulong nbIframes = nbOfFrames / 10;
		ulong nbBframes = 0;
		if (useBframes)
			nbBframes = (nbOfFrames - nbIframes) / 2;
		ulong nbPframes = nbOfFrames - nbIframes - nbBframes;
		totalOverhead = cast(double)(4300 + 1400 + nbIframes * mkvIframeOverhead + nbPframes * mkvPframeOverhead +
			nbBframes * mkvBframeOverhead);
		double nbOfSeconds = cast(double)nbOfFrames / framerate;
		totalOverhead += nbOfSeconds / 2.0 * 12; // 12 bytes per cluster

		long audioSize = 0L;
		double audioOverhead = 0;
		foreach (AudioStream stream; audioStreams)
		{
			audioSize += stream.SizeBytes;
			audioOverhead += getMKVAudioOverhead(stream.Type, 48000, nbOfSeconds);
		}

		totalOverhead += audioOverhead;
		double bytesPerSecond = desiredBitrate * 1000 / 8;
		long videoSize = cast(long)(nbOfSeconds * bytesPerSecond);
		rawVideoSize = rndint(videoSize / 1024);
		long size = videoSize + audioSize + cast(long)totalOverhead;
		return size / 1024L;
	}

	/// gets the overhead a given audio type will incurr; the matroska container
	/// given its length and sampling rate
	/// Params:
	/// AudioType = type of the audio track
	/// samplingRate = sampling rate of the audio track
	/// length = length of the audio track
	/// Returns: overhead this audio track will incurr
	public int getMKVAudioOverhead(AudioType audioType, int samplingRate, double length)
	{
		if (audioType == 0)
			return 0;
		int nbSamples = rndint(cast(double)samplingRate * length);
		int headerSize = mkvAudioTrackHeaderSize;
		int samplesPerBlock = 0;
		if (audioType == AudioType.MP4AAC)
			samplesPerBlock = AACBlockSize;
		else if (audioType == AudioType.VBRMP3 || audioType == AudioType.CBRMP3)
			samplesPerBlock = MP3BlockSize;
		else if (audioType == AudioType.AC3)
			samplesPerBlock = AC3BlockSize;
		else if (audioType == AudioType.VORBIS)
		{
			samplesPerBlock = VorbisBlockSize;
			headerSize = mkvVorbisTrackHeaderSize;
		}
		else // unknown types.. we presume the same overhead as for DTS
		{
			samplesPerBlock = AC3BlockSize;
		}
		double blockOverhead = cast(double)nbSamples / cast(double)samplesPerBlock * 22.0 / 8.0;
		int overhead = rndint(headerSize + 5 * length + blockOverhead);
		return overhead;
	}

	/// calculates the AVI bitrate given the desired audio streams and video stream properties
	/// Params:
	/// audioStreams = the audio streams to be muxed to the video
	/// desiredOutputSize = the desired final filesize
	/// nbOfFrames = the number of frames of the source
	/// framerate = the framerate of the source
	/// videoSize = the size of the raw video stream
	/// Returns: the bitrate; kbit/s
	private int calculateAVIBitrate(AudioStream[] audioStreams, long desiredOutputSize, ulong nbOfFrames, double framerate, out int videoSize)
	{
		double videoOverhead = cast(double)nbOfFrames * aviVideoOverhead;
		double nbOfSeconds = cast(double)nbOfFrames / framerate;
		double totalOverhead = videoOverhead;
		long audioSize = 0;
		foreach (AudioStream stream; audioStreams)
		{
			audioSize += stream.SizeBytes;
			if (stream.SizeBytes > 0)
			{
				double audioOverhead = getAviAudioOverhead(stream.Type);
				totalOverhead += audioOverhead * nbOfFrames;
			}
		}
		long videoTargetSize = desiredOutputSize - audioSize - cast(long)totalOverhead;
		videoSize = rndint(videoTargetSize / 1024);
		long sizeInBits = videoTargetSize * 8;
		int bitrate = rndint(sizeInBits / (nbOfSeconds * 1000));
		return bitrate;
	}

	private long calculateAVISize(AudioStream[] audioStreams, int desiredBitrate, ulong nbOfFrames, double framerate, out int rawVideoSize)
	{
		double videoOverhead = this.aviVideoOverhead;
		double nbOfSeconds = cast(double)nbOfFrames / framerate;
		long audioSize = 0L;
		double audioOverhead = 0;
		foreach (AudioStream stream; audioStreams)
		{
			audioSize += stream.SizeBytes;
			if (stream.SizeBytes > 0)
				audioOverhead += getAviAudioOverhead(stream.Type) * nbOfFrames;
		}
		double totalOverhead = videoOverhead + audioOverhead;
		double bytesPerSecond = desiredBitrate * 1000 / 8;
		long videoSize = cast(long)(nbOfSeconds * bytesPerSecond);
		rawVideoSize = rndint(videoSize / 1024);
		long size = videoSize + cast(long)audioSize + cast(long)totalOverhead;
		return size / 1024L;
	}

	/// gets the avi container overhead given the type of audio file to be muxed
	/// Params:
	/// AudioType = the type of audio; question
	/// Returns: the overhead per video frame for the given audio type
	public double getAviAudioOverhead(AudioType audioType)
	{
		double audioOverhead = 0;
		if (audioType == AudioType.AC3)
			audioOverhead = ac3Overhead;
		else if (audioType == AudioType.CBRMP3)
			audioOverhead = cbrMP3Overhead;
		else if (audioType == AudioType.VBRMP3)
			audioOverhead = vbrMP3Overhead;
		else if (audioType == AudioType.MP4AAC)
			audioOverhead = 0;
		else
			audioOverhead = 0;
		return audioOverhead;
	}

	///
	public int CalculateBitrateKBits(VideoCodec codec, bool useBframes, ContainerType container, AudioStream[] audioStreams, long desiredOutputSizeBytes, ulong nbOfFrames, double framerate, out int videoSizeKB)
	{
		if (container == ContainerType.MP4)
			return calculateMP4VideoBitrate(audioStreams, desiredOutputSizeBytes, nbOfFrames, useBframes, framerate, videoSizeKB);
		if (container == ContainerType.AVI)
			return calculateAVIBitrate(audioStreams, desiredOutputSizeBytes, nbOfFrames, framerate, videoSizeKB);
		if (container == ContainerType.MKV)
			return calculateMKVVideoBitrate(audioStreams, desiredOutputSizeBytes, nbOfFrames, framerate, useBframes, videoSizeKB);
		
		videoSizeKB = 0;
		return 0;
	}

	///	
	public long CalculateFileSizeKB(VideoCodec codec, bool useBframes, ContainerType container, AudioStream[] audioStreams, int desiredBitrate, ulong nbOfFrames, double framerate, out int rawVideoSize)
	{
		if (container == ContainerType.MP4)
			return calculateMP4Size(audioStreams, desiredBitrate, nbOfFrames, useBframes, framerate, rawVideoSize);
		if (container == ContainerType.AVI)
			return calculateAVISize(audioStreams, desiredBitrate, nbOfFrames, framerate, rawVideoSize);
		if (container == ContainerType.MKV)
			return calculateMKVSize(audioStreams, desiredBitrate, nbOfFrames, framerate, useBframes, rawVideoSize);

		rawVideoSize = 0;
		return 0;
	}

	/// gets the video container overhead per frame given the b-frame choice
	/// Params:
	/// useBframes = whether we have b-frames; the source or not (causes a great increase; overhead)
	/// Returns: the overhead per video frame
	private double getMP4Overhead(bool useBframes)
	{
		if (useBframes)
			return cast(double)mp4OverheadWithBframes;
		else
			return cast(double)mp4OverheadWithoutBframes;
	}
}