go-mxl-pattern-generator/internal/generator/cpu_patterns.go

package generator

import (
	"fmt"
)

var ebu75Colors = [...]YCbCr10{
	{Y: 721, Cb: 512, Cr: 512},
	{Y: 674, Cb: 176, Cr: 543},
	{Y: 581, Cb: 589, Cr: 176},
	{Y: 534, Cb: 253, Cr: 207},
	{Y: 251, Cb: 771, Cr: 817},
	{Y: 204, Cb: 435, Cr: 848},
	{Y: 111, Cb: 848, Cr: 481},
	{Y: 64, Cb: 512, Cr: 512},
}

var ebu100Colors = [...]YCbCr10{
	{Y: 940, Cb: 512, Cr: 512},
	{Y: 877, Cb: 64, Cr: 553},
	{Y: 754, Cb: 615, Cr: 64},
	{Y: 691, Cb: 167, Cr: 105},
	{Y: 313, Cb: 857, Cr: 919},
	{Y: 250, Cb: 409, Cr: 960},
	{Y: 127, Cb: 960, Cr: 471},
	{Y: 64, Cb: 512, Cr: 512},
}

var smpte75Colors = [...]YCbCr10{
	{Y: 721, Cb: 512, Cr: 512},
	{Y: 674, Cb: 176, Cr: 543},
	{Y: 581, Cb: 589, Cr: 176},
	{Y: 534, Cb: 253, Cr: 207},
	{Y: 251, Cb: 771, Cr: 817},
	{Y: 204, Cb: 435, Cr: 848},
	{Y: 111, Cb: 848, Cr: 481},
}

var grayBarsColors = [...]YCbCr10{
	{Y: 64, Cb: 512, Cr: 512},
	{Y: 137, Cb: 512, Cr: 512},
	{Y: 210, Cb: 512, Cr: 512},
	{Y: 283, Cb: 512, Cr: 512},
	{Y: 356, Cb: 512, Cr: 512},
	{Y: 429, Cb: 512, Cr: 512},
	{Y: 502, Cb: 512, Cr: 512},
	{Y: 575, Cb: 512, Cr: 512},
	{Y: 648, Cb: 512, Cr: 512},
	{Y: 721, Cb: 512, Cr: 512},
	{Y: 794, Cb: 512, Cr: 512},
	{Y: 867, Cb: 512, Cr: 512},
	{Y: 940, Cb: 512, Cr: 512},
}

var (
	ebu75BaseColor    = colorBars(ebu75Colors[:])
	ebu100BaseColor   = colorBars(ebu100Colors[:])
	grayBarsBaseColor = colorBars(grayBarsColors[:])
)

func NewCPUPatternGenerator(width, height uint, pattern string) (*CPUGenerator, error) {
	var baseColor baseColorFunc
	var dynamic bool
	switch pattern {
	case "ebu75":
		baseColor = ebu75BaseColor
	case "ebu75-move":
		baseColor, dynamic = ebu75BaseColor, true
	case "ebu100":
		baseColor = ebu100BaseColor
	case "ebu100-move":
		baseColor, dynamic = ebu100BaseColor, true
	case "smpte":
		if err := validateRP219Size(width, height); err != nil {
			return nil, err
		}
		baseColor = rp219BaseColor
	case "smpte-move":
		if err := validateRP219Size(width, height); err != nil {
			return nil, err
		}
		baseColor, dynamic = rp219BaseColor, true
	case "gray-bars":
		baseColor = grayBarsBaseColor
	case "gray-bars-move":
		baseColor, dynamic = grayBarsBaseColor, true
	case "gray-ramp":
		baseColor = grayRampBaseColor
	case "gray-ramp-move":
		baseColor, dynamic = grayRampBaseColor, true
	default:
		return nil, fmt.Errorf("cpu pattern %q is not implemented", pattern)
	}

	var patch FrameRenderer
	if dynamic {
		patch = movingSquarePatch(baseColor)
	}
	return NewCPUGenerator(width, height, baseRenderer(baseColor), patch)
}

func grayRampBaseColor(x, _, width, _ int) YCbCr10 {
	return YCbCr10{
		Y:  uint32(64 + (x*876)/width),
		Cb: 512,
		Cr: 512,
	}
}

func colorBars(colors []YCbCr10) baseColorFunc {
	return func(x, _, width, _ int) YCbCr10 {
		bar := min(x*len(colors)/width, len(colors)-1)
		return colors[bar]
	}
}

func baseRenderer(baseColor baseColorFunc) FrameRenderer {
	return func(dst []byte, width, height, _ int) error {
		return renderBasePattern(dst, width, height, baseColor)
	}
}

func movingSquarePatch(baseColor baseColorFunc) FrameRenderer {
	return func(dst []byte, width, height, frameIndex int) error {
		return patchMovingSquare(dst, width, height, frameIndex, baseColor)
	}
}

func renderBasePattern(dst []byte, width, height int, baseColor baseColorFunc) error {
	stride := V210LineSize(width)
	for y := 0; y < height; y++ {
		for x := 0; x < width; x += 6 {
			var pixels [6]YCbCr10
			for i := range pixels {
				pixels[i] = baseColor(x+i, y, width, height)
			}
			offset := y*stride + x/6*16
			packV210Block(dst[offset:offset+16], pixels)
		}
	}
	return nil
}

var (
	gray40 = YCbCr10{Y: 414, Cb: 512, Cr: 512}
	gray15 = YCbCr10{Y: 195, Cb: 512, Cr: 512}
	black  = YCbCr10{Y: 64, Cb: 512, Cr: 512}
	white  = YCbCr10{Y: 940, Cb: 512, Cr: 512}
)

func rp219BaseColor(x, y, width, height int) YCbCr10 {
	barsWidth := (height / 3) * 4
	barsStart := (width - barsWidth) / 2
	barsEnd := barsStart + barsWidth
	oneBarWidth := barsWidth / 7

	unitH := height / 12
	section1End := unitH * 7
	section2End := section1End + unitH
	section3End := section2End + unitH

	rampStart := barsStart + oneBarWidth
	rampEnd := barsEnd
	rampWidth := rampEnd - rampStart

	switch {
	case y < section1End:
		if x < barsStart || x >= barsEnd {
			return gray40
		}

		bar := min(
			(x-barsStart)*len(smpte75Colors)/barsWidth,
			len(smpte75Colors)-1,
		)
		return smpte75Colors[bar]

	case y < section2End:
		switch {
		case x < barsStart:
			return YCbCr10{Y: 754, Cb: 615, Cr: 64} // 100% cyan
		case x < rampStart:
			return YCbCr10{Y: 244, Cb: 612, Cr: 395} // -I
		case x >= barsEnd:
			return YCbCr10{Y: 127, Cb: 960, Cr: 471} // 100% blue
		default:
			return YCbCr10{Y: 721, Cb: 512, Cr: 512} // 75% white
		}

	case y < section3End:
		switch {
		case x < barsStart:
			return YCbCr10{Y: 877, Cb: 64, Cr: 553} // 100% yellow
		case x < rampStart:
			return YCbCr10{Y: 141, Cb: 697, Cr: 606} // +Q
		case x >= barsEnd:
			return YCbCr10{Y: 250, Cb: 409, Cr: 960} // 75% red
		default:
			rampY := 64 + ((x-rampStart)*876)/rampWidth
			return YCbCr10{Y: uint32(rampY), Cb: 512, Cr: 512}
		}

	default:
		if x < barsStart || x >= barsEnd {
			return gray15
		}

		bar := (x - barsStart) / oneBarWidth
		switch {
		case bar == 1:
			return white
		case bar == 3:
			offset := (x - barsStart) % oneBarWidth
			subBar := offset * 3 / oneBarWidth

			switch subBar {
			case 0:
				return YCbCr10{Y: 46, Cb: 512, Cr: 512}
			case 1:
				return black
			default:
				return YCbCr10{Y: 82, Cb: 512, Cr: 512}
			}

		default:
			return black
		}
	}
}

func validateRP219Size(width, height uint) error {
	if height < 12 {
		return fmt.Errorf("cpu RP 219 requires frame height of at least 12, got %d", height)
	}
	barsWidth := (height / 3) * 4
	if barsWidth < 7 || width < barsWidth {
		return fmt.Errorf(
			"cpu RP 219 requires frame width %d to fit a 4:3 pattern area for height %d, got %d",
			barsWidth,
			height,
			width,
		)
	}
	return nil
}