package goimagehash

import (
	"image"
)

type WebPUpsampleLinePairFunc func(topY []uint8, bottomY []uint8,
	topU []uint8, topV []uint8,
	bottomU []uint8, bottomV []uint8,
	topDst []uint8, bottomDst []uint8, len int)

func FancyUpscale(yuv *image.YCbCr) *image.RGBA {
	rgb := image.NewRGBA(image.Rect(0, 0, yuv.Rect.Dx(), yuv.Rect.Dy()))
	var (
		upsample WebPUpsampleLinePairFunc = UpsampleRgbaLinePair

		lastRow = yuv.Rect.Dy()
		mbW     = yuv.Rect.Dx()
	)
	// First line is special cased. We mirror the u/v samples at boundary.
	upsample(yuv.Y[:yuv.YStride], nil,
		yuv.Cb[:yuv.CStride], yuv.Cr[:yuv.CStride],
		yuv.Cb[:yuv.CStride], yuv.Cr[:yuv.CStride],
		rgb.Pix[:rgb.Stride], nil,
		mbW,
	)

	var (
		topC      int
		bottomC   int
		topY      int
		bottomY   int
		topDst    int
		bottomDst int
	)
	// Loop over each output pairs of row.
	for row := 1; row+1 < lastRow; row += 2 {
		topC = yuv.COffset(0, row)
		bottomC = topC + yuv.CStride

		topY = yuv.YOffset(0, row)
		bottomY = topY + yuv.YStride

		topDst = rgb.PixOffset(0, row)
		bottomDst = topDst + rgb.Stride

		upsample(yuv.Y[topY:bottomY], yuv.Y[bottomY:bottomY+yuv.YStride],
			yuv.Cb[topC:bottomC], yuv.Cr[topC:bottomC],
			yuv.Cb[bottomC:bottomC+yuv.CStride], yuv.Cr[bottomC:bottomC+yuv.CStride],
			rgb.Pix[topDst:bottomDst], rgb.Pix[bottomDst:bottomDst+rgb.Stride],
			mbW,
		)

	}

	// Process the very last row of even-sized picture
	if lastRow%2 == 0 {
		upsample(yuv.Y[len(yuv.Y)-yuv.YStride:], nil,
			yuv.Cb[len(yuv.Cb)-yuv.CStride:], yuv.Cr[len(yuv.Cr)-yuv.CStride:],
			yuv.Cb[len(yuv.Cb)-yuv.CStride:], yuv.Cr[len(yuv.Cr)-yuv.CStride:],
			rgb.Pix[len(rgb.Pix)-rgb.Stride:], nil,
			mbW,
		)
	}
	return rgb
}

func UpsampleRgbaLinePair(topY []uint8, bottomY []uint8,
	topU []uint8, topV []uint8,
	bottomU []uint8, bottomV []uint8,
	topDst []uint8, bottomDst []uint8, rowLength int,
) {
	var (
		x             int
		lastPixelPair = (rowLength - 1) >> 1
		tlUV          = (uint32(topU[0]) | (uint32(topV[0]) << 16))       /* top-left sample */
		lUV           = (uint32(bottomU[0]) | (uint32(bottomV[0]) << 16)) /* left-sample */
	)
	{
		uv0 := (3*tlUV + lUV + 0x00020002) >> 2
		dst := VP8YuvToRgba(topY[0], uint8(uv0&0xff), uint8((uv0 >> 16)))
		copy(topDst[0:], dst[:])
	}
	if bottomY != nil {
		uv0 := (3*lUV + tlUV + 0x00020002) >> 2
		dst := VP8YuvToRgba(bottomY[0], uint8(uv0&0xff), uint8(uv0>>16))
		copy(bottomDst[0:], dst[:])
	}

	for x = 1; x <= lastPixelPair; x++ {
		var (
			tUV = (uint32(topU[x]) | (uint32(topV[x]) << 16))       /* top sample */
			uv  = (uint32(bottomU[x]) | (uint32(bottomV[x]) << 16)) /* sample */

			/* precompute invariant values associated with first and second diagonals*/
			avg    = tlUV + tUV + lUV + uv + 0x00080008
			diag12 = (avg + 2*(tUV+lUV)) >> 3
			diag03 = (avg + 2*(tlUV+uv)) >> 3
		)
		{
			var (
				uv0 = (diag12 + tlUV) >> 1
				uv1 = (diag03 + tUV) >> 1
			)
			dst := VP8YuvToRgba(topY[2*x-1], uint8(uv0&0xff), uint8(uv0>>16))
			copy(topDst[(2*x-1)*4:], dst[:])
			dst = VP8YuvToRgba(topY[2*x-0], uint8(uv1&0xff), uint8(uv1>>16))
			copy(topDst[(2*x-0)*4:], dst[:])
		}
		if bottomY != nil {
			var (
				uv0 = (diag03 + lUV) >> 1
				uv1 = (diag12 + uv) >> 1
			)
			dst := VP8YuvToRgba(bottomY[2*x-1], uint8(uv0&0xff), uint8(uv0>>16))
			copy(bottomDst[(2*x-1)*4:], dst[:])
			dst = VP8YuvToRgba(bottomY[2*x+0], uint8(uv1&0xff), uint8(uv1>>16))
			copy(bottomDst[(2*x+0)*4:], dst[:])
		}
		tlUV = tUV
		lUV = uv
	}
	if rowLength%2 == 0 {
		{
			uv0 := (3*tlUV + lUV + 0x00020002) >> 2
			dst := VP8YuvToRgba(topY[rowLength-1], uint8(uv0&0xff), uint8(uv0>>16))
			copy(topDst[(rowLength-1)*4:], dst[:])
		}
		if bottomY != nil {
			uv0 := (3*lUV + tlUV + 0x00020002) >> 2
			dst := VP8YuvToRgba(bottomY[rowLength-1], uint8(uv0&0xff), uint8(uv0>>16))
			copy(bottomDst[(rowLength-1)*4:], dst[:])
		}
	}
}
