# AVSLib :: 'clip' package :: 'channels' module
# Copyright (c) 2005 George Zarkadas (gzarkadas@users.sourceforge.net)

# 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"

__lib__ = "avslib" #$$ [__lib__ = "avslib-debug"]

DefineModule(__lib__, "filters", "channels")

global __load__ = IsPackageLoaded(__lib__, "base") ? NOP : Eval("""
	LoadModule(__lib__, "base", "core")
	LoadModule(__lib__, "base", "conversion")
	""")
global __load__ = IsPackageLoaded(__lib__, "array") ? NOP : Eval("""
	LoadModule(__lib__, "array", "core")	
	LoadModule(__lib__, "array", "operators")
	""")
global __load__ = Eval("""
	LoadModule(__lib__, "clip", "core")
	""")

global __show_pxtype_msg = \
	" supports the following output pixel types: RGB, YUY2, or YV12"

# Provide RGB channels support for Avisynth versions < 2.5.6

global __load__ = VersionNumber() > 2.55 \
	? NOP \
	: Import("ver_2-5-5\rgb_channels.avsi")

# YUV Show... analogs of RGB Show... filters.
# Note that YUV channels do not upscale if an RGB[32/24] pixel type
# is requested. 
# To return a full [0.255] range from an ordinary CCIR YUV clip you 
# will have to manually perform ColorYUV(levels="tv->pc") *before*
# extracting the channels.

	# In proper use of the function 'c' is guaranteed to be YUV (and 
	# 'pixel_type' to be valid). 
	# Treat all values as pc range (ie do not upscale when converting to RGB)
	# for uniformity with the Show{R/G/B}, Merge(A)RGB filters' behavior.
	#
	Function __show_yuv_convert(clip c, string pixel_type)
	{
		pixel_type = pixel_type == "RGB" ? "RGB32" : pixel_type
		return pixel_type == c.ColorSpace ? c : ( \
			pixel_type == "YV12" ? \
				c.ConvertToYV12() : ( \
			pixel_type == "YUY2" ? \
				c.ConvertToYUY2() : ( \
			pixel_type == "RGB32" ? \
				c.ConvertToRGB32(matrix="PC.601") \
			: \
				c.ConvertToRGB24(matrix="PC.601") )))
	}

Function ShowY(clip c, string "pixel_type")
{
	Assert(c.IsYUV, "ShowY: YUV data only")
	pixel_type = UCase(Default(pixel_type, c.ColorSpace))
	Assert(IsPixelType(pixel_type, true), "ShowY" + __show_pxtype_msg)
	return __show_yuv_convert( c.Greyscale(), pixel_type )
}

Function ShowU(clip c, string "pixel_type")
{
	Assert(c.IsYUV, "ShowY: YUV data only")
	pixel_type = UCase(Default(pixel_type, c.ColorSpace))
	Assert(IsPixelType(pixel_type, true), "ShowU" + __show_pxtype_msg)
	return __show_yuv_convert( c.UToY(), pixel_type )	
}

Function ShowV(clip c, string "pixel_type")
{
	Assert(c.IsYUV, "ShowY: YUV data only")
	pixel_type = UCase(Default(pixel_type, c.ColorSpace))
	Assert(IsPixelType(pixel_type, true), "ShowV" + __show_pxtype_msg)
	return __show_yuv_convert( c.VToY(), pixel_type )	
}

# In order to simulate the Avisynth's 2.5.6 Merge(A)RGB filters' behavior
# we must convert RGB (or YUV) clips before passing them to YToUV.
# Rule set:
# 1. Treat all values as pc range (do not downscale when converting from RGB)
# 2. Greyscale RGB clips before converting them to YUV
# 3. Output colorspace is determined by ClipY (YUY2 if RGB/YUY2, YV12 if YV12)
# 4. Given (w,h) dimensions of ClipY, ClipU / ClipV dimensions must be
#    (w/2,h) if clipY == RGB / YUY2, (w/2,h/2) if clipY == YV12
	
	
	# returns a grayscale clip from c
	# if c is RGB converts to YUY2 without downnscaling
	#
	Function __merge_yuv_adjust(clip c)
	{
		c = c.Greyscale()
		return c.IsRGB ? c.ConvertToYUY2(matrix="PC.601") : c
	}	

# TODO in later versions:
# how do we handle the case where input clip(s) is interlaced?

	Function __merge_yuv_wok(clip y, clip u, clip v)
	{
		w = y.Width
		return u.Width == w/2 && v.Width == w/2
	}
	
	Function __merge_yuv_hok(clip y, clip u, clip v)
	{
		h = y.Height
		u_hcond = y.IsYV12 ? u.Height == h/2 : u.Height == h
		v_hcond = y.IsYV12 ? v.Height == h/2 : v.Height == h
		return u_hcond && v_hcond
	}

Function MergeYUV(clip clipY, clip clipU, clip clipV)
{
	output = clipY.IsRGB ? "YUY2" : clipY.ColorSpace
	clipY = __merge_yuv_adjust(clipY)
	clipU = __merge_yuv_adjust(clipU)
	clipV = __merge_yuv_adjust(clipV)
	# check input clips sizes
	Assert(__merge_yuv_wok(clipY, clipU, clipV), \
		"MergeYUV: 'clipU', 'clipV' width(s) must be half of 'clipY' width.")
	Assert(__merge_yuv_hok(clipY, clipU, clipV), \
		"MergeYUV: 'clipU', 'clipV' height(s) must be equal " + \
		"to 'clipY' height" + Chr(10) + "if the later is RGB or YUY2, " + \
		"or half of 'clipY' height if the later is YV12.")
	# convert clips to same (output) colorspace
	# (YToUV will throw an error if not all clips have same colorspace)
	clipU = clipU.ColorSpace != output \
		? Eval("clipU.ConvertTo" + output + "()") \
		: clipU
	clipV = clipV.ColorSpace != output \
		? Eval("clipV.ConvertTo" + output + "()") \
		: clipV
	return YToUV(clipU, clipV, clipY) 
}

# Splits clip's channels and places all inside an array.
# The very first element of the array (index 0) contains information about 
# the source clip's colorspace. Thus channels are placed in indexes:
#     1 (R/Y), 2 (G/U), 3 (B/V), 4 (A)
# If 'pixel_type' is provided converts channels to the requested colorspace;
# this is done through the Show... functions. Note that the number of channels
# does not changes; it remains that of the source clip's colorspace: 
# (R,G,B,A for RGB32, R,G,B for RGB24, Y,U,V for YUV).
#
Function SplitVideoChannels(clip c, string "pixel_type")
{
	ret = c.IsYUV ? \
			ArrayCreate( \
				c.ShowY(pixel_type), \
				c.ShowU(pixel_type), \
				c.ShowV(pixel_type) \
			) : ( \
		c.IsRGB32 ? \
			ArrayCreate( \
				c.ShowRed(pixel_type), \
				c.ShowGreen(pixel_type), \
				c.ShowBlue(pixel_type), \
				c.ShowAlpha(pixel_type) \
			) \
		: \
			ArrayCreate( \
				c.ShowRed(pixel_type), \
				c.ShowGreen(pixel_type), \
				c.ShowBlue(pixel_type) \
			) )
	return ret.ArrayInsert(c.ColorSpace, 0)
}

Function MergeVideoChannels(string channels)
{
	src_cs = UCase(channels.ArrayGet(0))
	return \
		src_cs == "RGB32" ? \
			MergeARGB( \
				channels.ArrayGet(4), \
				channels.ArrayGet(1), \
				channels.ArrayGet(2), \
				channels.ArrayGet(3) \
			) : ( \
		src_cs == "YUY2" || src_cs == "YV12" ? \
			MergeYUV( \
				channels.ArrayGet(1), \
				channels.ArrayGet(2), \
				channels.ArrayGet(3) \
			) : ( \
		src_cs == "RGB24" ? \
			MergeRGB( \
				channels.ArrayGet(1), \
				channels.ArrayGet(2), \
				channels.ArrayGet(3), \
				"RGB24" \
			) \
		: Throw("MergeVideoChannels: invalid channels' array format") ))
}

global __channel_id_msg = \
	"ChannelIndex: channel must be one of R, G, B, A, Y, U, V"

Function ChannelIndex(string ch)
{
	Assert(StrLen(ch) == 1, __channel_id_msg)
	i = FindStr("RGBAYUV", UCase(ch))
	return i == 0 \
		? Throw(__channel_id_msg) \
		: (i <= 4 ? i : i % 4)
}

# Functions to test the colorspace of channels.
# Return true only if ALL channels (array elements) share the same pixel type.

Function IsRGBChannels(string channels)
{
	return channels.ArrayLen == channels.ArrayOpFunc("IsRGB" \
		).ArrayOpFunc("CInt").ArraySum()
}

Function IsRGB32Channels(string channels)
{
	return channels.ArrayLen == channels.ArrayOpFunc("IsRGB32" \
		).ArrayOpFunc("CInt").ArraySum()
}

Function IsRGB24Channels(string channels)
{
	return channels.ArrayLen == channels.ArrayOpFunc("IsRGB24" \
		).ArrayOpFunc("CInt").ArraySum()
}

Function IsYUVChannels(string channels)
{
	return channels.ArrayLen == channels.ArrayOpFunc("IsYUV" \
		).ArrayOpFunc("CInt").ArraySum()
}

Function IsYUY2Channels(string channels)
{
	return channels.ArrayLen == channels.ArrayOpFunc("IsYUY2" \
		).ArrayOpFunc("CInt").ArraySum()
}

Function IsYV12Channels(string channels)
{
	return channels.ArrayLen == channels.ArrayOpFunc("IsYV12" \
		).ArrayOpFunc("CInt").ArraySum()
}
