<commit_before>package diff

import (
	"context"
	"image"
	"image/color"
	"image/draw"
	"math"
	"unsafe"

	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/paramtools"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/golden/go/types"
)

var (
	pixelMatchColor = color.Transparent

	// Orange gradient.
	//
	// These are non-premultiplied RGBA values.
	pixelDiffColor = [][]uint8{
		{0xfd, 0xd0, 0xa2, 0xff},
		{0xfd, 0xae, 0x6b, 0xff},
		{0xfd, 0x8d, 0x3c, 0xff},
		{0xf1, 0x69, 0x13, 0xff},
		{0xd9, 0x48, 0x01, 0xff},
		{0xa6, 0x36, 0x03, 0xff},
		{0x7f, 0x27, 0x04, 0xff},
	}

	// Blue gradient.
	//
	// These are non-premultiplied RGBA values.
	pixelAlphaDiffColor = [][]uint8{
		{0xc6, 0xdb, 0xef, 0xff},
		{0x9e, 0xca, 0xe1, 0xff},
		{0x6b, 0xae, 0xd6, 0xff},
		{0x42, 0x92, 0xc6, 0xff},
		{0x21, 0x71, 0xb5, 0xff},
		{0x08, 0x51, 0x9c, 0xff},
		{0x08, 0x30, 0x6b, 0xff},
	}
)

// Returns the offset into the color slices (pixelDiffColor,
// or pixelAlphaDiffColor) based on the delta passed in.
//
// The number passed in is the difference between two colors,
// on a scale from 1 to 1024.
func deltaOffset(n int) int {
	ret := int(math.Ceil(math.Log(float64(n))/math.Log(3) + 0.5))
	if ret < 1 || ret > 7 {
		sklog.Fatalf("Input: %d", n)
	}
	return ret - 1
}

// DiffMetrics contains the diff information between two images.
type DiffMetrics struct {
	// NumDiffPixels is the absolute number of pixels that are different.
	NumDiffPixels int

	// CombinedMetric is a value in [0, 10] that represents how large the diff is between two
	// images. It is based off the MaxRGBADiffs and PixelDiffPercent.
	CombinedMetric float32

	// PixelDiffPercent is the percentage of pixels that are different. The denominator here is
	// (maximum width of the two images) * (maximum height of the two images).
	PixelDiffPercent float32

	// MaxRGBADiffs contains the maximum difference of each channel.
	MaxRGBADiffs [4]int

	// DimDiffer is true if the dimensions between the two images are different.
	DimDiffer bool
}

// ComputeDiffMetrics computes and returns the diff metrics between two given images.
func ComputeDiffMetrics(leftImg *image.NRGBA, rightImg *image.NRGBA) *DiffMetrics {
	defer metrics2.FuncTimer().Stop()
	ret, _ := PixelDiff(leftImg, rightImg)
	ret.CombinedMetric = CombinedDiffMetric(ret.MaxRGBADiffs, ret.PixelDiffPercent)
	return ret
}

// CombinedDiffMetric returns a value in [0, 10] that represents how large
// the diff is between two images. Implements the MetricFn signature.
func CombinedDiffMetric(channelDiffs [4]int, pixelDiffPercent float32) float32 {
	// Turn maxRGBA into a percent by taking the root mean square difference from
	// [0, 0, 0, 0].
	sum := 0.0
	for _, c := range channelDiffs {
		sum += float64(c) * float64(c)
	}
	normalizedRGBA := math.Sqrt(sum/float64(len(channelDiffs))) / 255.0
	// We take the sqrt of (pixelDiffPercent * normalizedRGBA) to straighten out
	// the curve, i.e. think about what a plot of x^2 would look like in the
	// range [0, 1].
	return float32(math.Sqrt(float64(pixelDiffPercent) * normalizedRGBA))
}

// getPixelDiffPercent returns the percentage of pixels that differ, as a float between 0 and 100
// (inclusive).
func getPixelDiffPercent(numDiffPixels, totalPixels int) float32 {
	return (float32(numDiffPixels) * 100) / float32(totalPixels)
}

func uint8ToColor(c []uint8) color.Color {
	return color.NRGBA{R: c[0], G: c[1], B: c[2], A: c[3]}
}

// diffColors compares two color values and returns a color to indicate the
// difference. If the colors differ it updates maxRGBADiffs to contain the
// maximum difference over multiple calls.
// If the RGB channels are identical, but the alpha differ then
// pixelAlphaDiffColor is returned. This allows to distinguish pixels that
// render the same, but have different alpha values.
func diffColors(color1, color2 color.Color, maxRGBADiffs *[4]int) color.Color {
	// We compare them before normalizing to non-premultiplied. If one of the
	// original images did not have an alpha channel (but the other did) the
	// equality will be false.
	if color1 == color2 {
		return pixelMatchColor
	}

	// Treat all colors as non-premultiplied.
	c1 := color.NRGBAModel.Convert(color1).(color.NRGBA)
	c2 := color.NRGBAModel.Convert(color2).(color.NRGBA)

	rDiff := util.AbsInt(int(c1.R) - int(c2.R))
	gDiff := util.AbsInt(int(c1.G) - int(c2.G))
	bDiff := util.AbsInt(int(c1.B) - int(c2.B))
	aDiff := util.AbsInt(int(c1.A) - int(c2.A))
	maxRGBADiffs[0] = util.MaxInt(maxRGBADiffs[0], rDiff)
	maxRGBADiffs[1] = util.MaxInt(maxRGBADiffs[1], gDiff)
	maxRGBADiffs[2] = util.MaxInt(maxRGBADiffs[2], bDiff)
	maxRGBADiffs[3] = util.MaxInt(maxRGBADiffs[3], aDiff)

	// If the color channels differ we mark with the diff color.
	if (c1.R != c2.R) || (c1.G != c2.G) || (c1.B != c2.B) {
		// We use the Manhattan metric for color difference.
		return uint8ToColor(pixelDiffColor[deltaOffset(rDiff+gDiff+bDiff+aDiff)])
	}

	// If only the alpha channel differs we mark it with the alpha diff color.
	//
	if aDiff > 0 {
		return uint8ToColor(pixelAlphaDiffColor[deltaOffset(aDiff)])
	}

	return pixelMatchColor
}

// recode creates a new NRGBA image from the given image.
func recode(img image.Image) *image.NRGBA {
	ret := image.NewNRGBA(img.Bounds())
	draw.Draw(ret, img.Bounds(), img, image.Pt(0, 0), draw.Src)
	return ret
}

// GetNRGBA converts the image to an *image.NRGBA in an efficient manner.
func GetNRGBA(img image.Image) *image.NRGBA {
	switch t := img.(type) {
	case *image.NRGBA:
		return t
	case *image.RGBA:
		for i := 0; i < len(t.Pix); i += 4 {
			if t.Pix[i+3] != 0xff {
				sklog.Warning("Unexpected premultiplied image!")
				return recode(img)
			}
		}
		// If every alpha is 0xff then t.Pix is already in NRGBA format, simply
		// share Pix between the RGBA and NRGBA structs.
		return &image.NRGBA{
			Pix:    t.Pix,
			Stride: t.Stride,
			Rect:   t.Rect,
		}
	default:
		// TODO(mtklein): does it make sense we're getting other types, or a DM bug?
		return recode(img)
	}
}

// PixelDiff is a utility function that calculates the DiffMetrics and the image of the
// difference for the provided images.
func PixelDiff(img1, img2 image.Image) (*DiffMetrics, *image.NRGBA) {
	defer metrics2.FuncTimer().Stop()
	img1Bounds := img1.Bounds()
	img2Bounds := img2.Bounds()

	// Get the bounds we want to compare.<commit_msg>Reland "[gold] Add logging to diff"<commit_after>package diff

import (
	"context"
	"image"
	"image/color"
	"image/draw"
	"math"
	"unsafe"

	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/paramtools"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/golden/go/types"
)

var (
	pixelMatchColor = color.Transparent

	// Orange gradient.
	//
	// These are non-premultiplied RGBA values.
	pixelDiffColor = [][]uint8{
		{0xfd, 0xd0, 0xa2, 0xff},
		{0xfd, 0xae, 0x6b, 0xff},
		{0xfd, 0x8d, 0x3c, 0xff},
		{0xf1, 0x69, 0x13, 0xff},
		{0xd9, 0x48, 0x01, 0xff},
		{0xa6, 0x36, 0x03, 0xff},
		{0x7f, 0x27, 0x04, 0xff},
	}

	// Blue gradient.
	//
	// These are non-premultiplied RGBA values.
	pixelAlphaDiffColor = [][]uint8{
		{0xc6, 0xdb, 0xef, 0xff},
		{0x9e, 0xca, 0xe1, 0xff},
		{0x6b, 0xae, 0xd6, 0xff},
		{0x42, 0x92, 0xc6, 0xff},
		{0x21, 0x71, 0xb5, 0xff},
		{0x08, 0x51, 0x9c, 0xff},
		{0x08, 0x30, 0x6b, 0xff},
	}
)

// Returns the offset into the color slices (pixelDiffColor,
// or pixelAlphaDiffColor) based on the delta passed in.
//
// The number passed in is the difference between two colors,
// on a scale from 1 to 1024.
func deltaOffset(n int) int {
	ret := int(math.Ceil(math.Log(float64(n))/math.Log(3) + 0.5))
	if ret < 1 || ret > 7 {
		sklog.Fatalf("Input: %d", n)
	}
	return ret - 1
}

// DiffMetrics contains the diff information between two images.
type DiffMetrics struct {
	// NumDiffPixels is the absolute number of pixels that are different.
	NumDiffPixels int

	// CombinedMetric is a value in [0, 10] that represents how large the diff is between two
	// images. It is based off the MaxRGBADiffs and PixelDiffPercent.
	CombinedMetric float32

	// PixelDiffPercent is the percentage of pixels that are different. The denominator here is
	// (maximum width of the two images) * (maximum height of the two images).
	PixelDiffPercent float32

	// MaxRGBADiffs contains the maximum difference of each channel.
	MaxRGBADiffs [4]int

	// DimDiffer is true if the dimensions between the two images are different.
	DimDiffer bool
}

// ComputeDiffMetrics computes and returns the diff metrics between two given images.
func ComputeDiffMetrics(leftImg *image.NRGBA, rightImg *image.NRGBA) *DiffMetrics {
	defer metrics2.FuncTimer().Stop()
	ret, _ := PixelDiff(leftImg, rightImg)
	ret.CombinedMetric = CombinedDiffMetric(ret.MaxRGBADiffs, ret.PixelDiffPercent)
	return ret
}

// CombinedDiffMetric returns a value in [0, 10] that represents how large
// the diff is between two images. Implements the MetricFn signature.
func CombinedDiffMetric(channelDiffs [4]int, pixelDiffPercent float32) float32 {
	// Turn maxRGBA into a percent by taking the root mean square difference from
	// [0, 0, 0, 0].
	sum := 0.0
	for _, c := range channelDiffs {
		sum += float64(c) * float64(c)
	}
	normalizedRGBA := math.Sqrt(sum/float64(len(channelDiffs))) / 255.0
	// We take the sqrt of (pixelDiffPercent * normalizedRGBA) to straighten out
	// the curve, i.e. think about what a plot of x^2 would look like in the
	// range [0, 1].
	return float32(math.Sqrt(float64(pixelDiffPercent) * normalizedRGBA))
}

// getPixelDiffPercent returns the percentage of pixels that differ, as a float between 0 and 100
// (inclusive).
func getPixelDiffPercent(numDiffPixels, totalPixels int) float32 {
	return (float32(numDiffPixels) * 100) / float32(totalPixels)
}

func uint8ToColor(c []uint8) color.Color {
	return color.NRGBA{R: c[0], G: c[1], B: c[2], A: c[3]}
}

// diffColors compares two color values and returns a color to indicate the
// difference. If the colors differ it updates maxRGBADiffs to contain the
// maximum difference over multiple calls.
// If the RGB channels are identical, but the alpha differ then
// pixelAlphaDiffColor is returned. This allows to distinguish pixels that
// render the same, but have different alpha values.
func diffColors(color1, color2 color.Color, maxRGBADiffs *[4]int) color.Color {
	// We compare them before normalizing to non-premultiplied. If one of the
	// original images did not have an alpha channel (but the other did) the
	// equality will be false.
	if color1 == color2 {
		return pixelMatchColor
	}

	// Treat all colors as non-premultiplied.
	c1 := color.NRGBAModel.Convert(color1).(color.NRGBA)
	c2 := color.NRGBAModel.Convert(color2).(color.NRGBA)

	rDiff := util.AbsInt(int(c1.R) - int(c2.R))
	gDiff := util.AbsInt(int(c1.G) - int(c2.G))
	bDiff := util.AbsInt(int(c1.B) - int(c2.B))
	aDiff := util.AbsInt(int(c1.A) - int(c2.A))
	maxRGBADiffs[0] = util.MaxInt(maxRGBADiffs[0], rDiff)
	maxRGBADiffs[1] = util.MaxInt(maxRGBADiffs[1], gDiff)
	maxRGBADiffs[2] = util.MaxInt(maxRGBADiffs[2], bDiff)
	maxRGBADiffs[3] = util.MaxInt(maxRGBADiffs[3], aDiff)

	// If the color channels differ we mark with the diff color.
	if (c1.R != c2.R) || (c1.G != c2.G) || (c1.B != c2.B) {
		// We use the Manhattan metric for color difference.
		return uint8ToColor(pixelDiffColor[deltaOffset(rDiff+gDiff+bDiff+aDiff)])
	}

	// If only the alpha channel differs we mark it with the alpha diff color.
	//
	if aDiff > 0 {
		return uint8ToColor(pixelAlphaDiffColor[deltaOffset(aDiff)])
	}

	return pixelMatchColor
}

// recode creates a new NRGBA image from the given image.
func recode(img image.Image) *image.NRGBA {
	ret := image.NewNRGBA(img.Bounds())
	draw.Draw(ret, img.Bounds(), img, image.Pt(0, 0), draw.Src)
	return ret
}

// GetNRGBA converts the image to an *image.NRGBA in an efficient manner.
func GetNRGBA(img image.Image) *image.NRGBA {
	switch t := img.(type) {
	case *image.NRGBA:
		return t
	case *image.RGBA:
		for i := 0; i < len(t.Pix); i += 4 {
			if t.Pix[i+3] != 0xff {
				sklog.Warning("Unexpected premultiplied image!")
				return recode(img)
			}
		}
		// If every alpha is 0xff then t.Pix is already in NRGBA format, simply
		// share Pix between the RGBA and NRGBA structs.
		return &image.NRGBA{
			Pix:    t.Pix,
			Stride: t.Stride,
			Rect:   t.Rect,
		}
	default:
		// TODO(kjlubick): does it make sense we're getting other types, or a DM bug?
		sklog.Errorf("Unexpected image type %T", img)
		return recode(img)
	}
}

// PixelDiff is a utility function that calculates the DiffMetrics and the image of the
// difference for the provided images.
func PixelDiff(img1, img2 image.Image) (*DiffMetrics, *image.NRGBA) {
	defer metrics2.FuncTimer().Stop()
	img1Bounds := img1.Bounds()
	img2Bounds := img2.Bounds()

	// Get the bounds we want to compare.
{'fraction_non_alphanumeric': 0.12200826381850784, 'fraction_numerical': 0.03570593279800421, 'mean_word_length': 3.81893313298272, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2061054', 'n_tokens_mistral': 5041, 'n_tokens_neox': 4624, 'n_words': 1447}