// Copyright 2017 The goimagehash Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package goimagehash

import (
	"errors"
	"image"
	"image/draw"
	"sync"

	"gitea.narnian.us/lordwelch/goimagehash/etcs"
	"gitea.narnian.us/lordwelch/goimagehash/transforms"
	"github.com/disintegration/imaging"
)

var bufPool = sync.Pool{
	New: func() any {
		// The Pool's New function should generally only return pointer
		// types, since a pointer can be put into the return interface
		// value without an allocation:
		return &[]uint8{}
	},
}

func getBuf(size, capacity int) []uint8 {
	if capacity < size {
		capacity = size
	}

	buf := *bufPool.Get().(*[]uint8)
	if cap(buf) < capacity {
		bufPool.Put(&buf)
		buf = make([]uint8, capacity)
	}

	if len(buf) != size {
		buf = (buf)[:size]
	}
	return buf
}

func ToGray(img image.Image, pix []uint8) *image.Gray {
	c := img.Bounds().Dx() * img.Bounds().Dy()
	if cap(pix) < c {
		pix = make([]uint8, c)
	}
	pix = pix[:c]
	gray := &image.Gray{
		Pix:    transforms.Rgb2Gray(img, pix),
		Stride: img.Bounds().Dx(),
		Rect:   img.Bounds(),
	}
	return gray
}

func Resize(img image.Image, w, h int, gray *image.Gray) *image.Gray {
	resized := imaging.Resize(img, w, h, imaging.Lanczos)
	if gray == nil || len(gray.Pix) != w*h {
		gray = image.NewGray(image.Rect(0, 0, resized.Bounds().Dx(), resized.Bounds().Dy()))
	}
	draw.Draw(gray, resized.Bounds(), resized, resized.Bounds().Min, draw.Src)
	return gray
}

// AverageHash function returns a hash computation of average hash.
// Implementation follows
// http://www.hackerfactor.com/blog/index.php?/archives/432-Looks-Like-It.html
func AverageHash(img image.Image) (*ImageHash, error) {
	if img == nil {
		return nil, errors.New("image object can not be nil")
	}
	w, h := 8, 8
	c := w * h
	ahash := NewImageHash(0, AHash)

	pix := getBuf(img.Bounds().Dx()*img.Bounds().Dy(), 0)
	sizedPix := getBuf(c, 0)

	gray := Resize(ToGray(img, pix), w, h, &image.Gray{
		Pix:    sizedPix,
		Stride: w,
		Rect:   image.Rect(0, 0, w, h),
	})
	bufPool.Put(&pix)
	avg := etcs.MeanOfPixels(gray.Pix)

	for idx, p := range gray.Pix {
		if p > avg {
			ahash.leftShiftSet(len(gray.Pix) - idx - 1)
		}
	}
	bufPool.Put(&gray.Pix)
	return ahash, nil
}

// DifferenceHash function returns a hash computation of difference hash.
// Implementation follows
// http://www.hackerfactor.com/blog/?/archives/529-Kind-of-Like-That.html
func DifferenceHash(img image.Image) (*ImageHash, error) {
	if img == nil {
		return nil, errors.New("image object can not be nil")
	}
	w, h := 9, 8
	c := w * h
	dhash := NewImageHash(0, DHash)

	pix := getBuf(img.Bounds().Dx()*img.Bounds().Dy(), 0)
	sizedPix := getBuf(c, 0)

	gray := Resize(ToGray(img, pix), w, h, &image.Gray{
		Pix:    sizedPix,
		Stride: w,
		Rect:   image.Rect(0, 0, w, h),
	})
	bufPool.Put(&pix)
	idx := 0
	for y := 0; y < gray.Bounds().Dy(); y++ {
		for x := 0; x < gray.Bounds().Dx()-1; x++ {
			if gray.Pix[gray.PixOffset(x, y)] < gray.Pix[gray.PixOffset(x+1, y)] {
				dhash.leftShiftSet(64 - idx - 1)
			}
			idx++
		}
	}
	bufPool.Put(&gray.Pix)

	return dhash, nil
}

// PerceptionHash function returns a hash computation of phash.
// Implementation follows
// http://www.hackerfactor.com/blog/index.php?/archives/432-Looks-Like-It.html
func PerceptionHash(img image.Image) (*ImageHash, error) {
	if img == nil {
		return nil, errors.New("image object can not be nil")
	}

	w, h := 32, 32
	c := w * h

	phash := NewImageHash(0, PHash)

	pix := getBuf(img.Bounds().Dx()*img.Bounds().Dy(), 0)
	sizedPix := getBuf(c, 0)

	gray := Resize(ToGray(img, pix), w, h, &image.Gray{
		Pix:    sizedPix,
		Stride: w,
		Rect:   image.Rect(0, 0, w, h),
	})
	bufPool.Put(&pix)
	gray32 := make([]float64, len(gray.Pix))
	for i, p := range gray.Pix {
		gray32[i] = float64(p)
	}
	bufPool.Put(&gray.Pix)
	flattens := transforms.DCT2DFast32(&gray32)

	median := etcs.MedianOfPixelsFast64(flattens[:])

	for idx, p := range flattens {
		if p > median {
			phash.leftShiftSet(64 - idx - 1) // leftShiftSet
		}
	}

	return phash, nil
}

// FlattenPixels function flattens 2d array into 1d array.
func FlattenPixels(pixels [][]float64, x int, y int) []float64 {
	flattens := make([]float64, x*y)
	for i := 0; i < y; i++ {
		for j := 0; j < x; j++ {
			flattens[y*i+j] = pixels[i][j]
		}
	}
	return flattens
}

// ExtPerceptionHash function returns phash of which the size can be set larger than uint64
// Some variable name refer to https://github.com/JohannesBuchner/imagehash/blob/master/imagehash/__init__.py
// Support 64bits phash (width=8, height=8) and 256bits phash (width=16, height=16)
// Important: width * height should be a power of 2
func ExtPerceptionHash(img image.Image, hash_size, freq int) (*ExtImageHash, error) {
	imgSize := hash_size * freq
	if img == nil {
		return nil, errors.New("image object can not be nil")
	}
	if imgSize <= 0 || imgSize&(imgSize-1) != 0 {
		return nil, errors.New("width * height should be power of 2")
	}
	var phash []uint64
	// gray := resize(ToGray(img), imgSize, imgSize)
	// gray32 := make([][]float64, imgSize)
	// for i := range gray32 {
	// 	gray32[i] = make([]float64, imgSize)
	// 	for x := range gray32[i] {
	// 		gray32[i][x] = float64(gray.Pix[(gray.PixOffset(x, i))])
	// 	}
	// }

	// pixels := transforms.DCT2D(gray32, hash_size, hash_size)
	// flattens := FlattenPixels(pixels, hash_size, hash_size)
	// median := etcs.MedianOfPixels(flattens)
	// print(hash_size, "x", freq, " pix:", len(gray.Pix), " h:", imgSize, " flatten:", len(flattens), "\n")

	// lenOfUnit := 64
	// if (hash_size*hash_size)%lenOfUnit == 0 {
	// 	phash = make([]uint64, (hash_size*hash_size)/lenOfUnit)
	// } else {
	// 	phash = make([]uint64, (hash_size*hash_size)/lenOfUnit+1)
	// }
	// for idx, p := range flattens {
	// 	indexOfArray := idx / lenOfUnit
	// 	indexOfBit := lenOfUnit - idx%lenOfUnit - 1
	// 	if p > median {
	// 		phash[indexOfArray] |= 1 << uint(indexOfBit)
	// 	}
	// }
	return NewExtImageHash(phash, PHash, imgSize), nil
}

// ExtAverageHash function returns ahash of which the size can be set larger than uint64
// Support 64bits ahash (width=8, height=8) and 256bits ahash (width=16, height=16)
// func ExtAverageHash(img image.Image, width, height int) (*ExtImageHash, error) {
// 	if img == nil {
// 		return nil, errors.New("image object can not be nil")
// 	}
// 	var ahash []uint64
// 	imgSize := width * height

// 	gray := resize(ToGray(img), width, height)
// 	avg := etcs.MeanOfPixels(gray.Pix)

// 	lenOfUnit := 64
// 	if imgSize%lenOfUnit == 0 {
// 		ahash = make([]uint64, imgSize/lenOfUnit)
// 	} else {
// 		ahash = make([]uint64, imgSize/lenOfUnit+1)
// 	}
// 	for idx, p := range gray.Pix {
// 		indexOfArray := idx / lenOfUnit
// 		indexOfBit := lenOfUnit - idx%lenOfUnit - 1
// 		if p > avg {
// 			ahash[indexOfArray] |= 1 << uint(indexOfBit)
// 		}
// 	}
// 	return NewExtImageHash(ahash, AHash, imgSize), nil
// }

// ExtDifferenceHash function returns dhash of which the size can be set larger than uint64
// Support 64bits dhash (width=8, height=8) and 256bits dhash (width=16, height=16)
// func ExtDifferenceHash(img image.Image, width, height int) (*ExtImageHash, error) {
// 	if img == nil {
// 		return nil, errors.New("image object can not be nil")
// 	}

// 	var dhash []uint64
// 	imgSize := width * height

// 	gray := resize(ToGray(img), width+1, height)

// 	lenOfUnit := 64
// 	if imgSize%lenOfUnit == 0 {
// 		dhash = make([]uint64, imgSize/lenOfUnit)
// 	} else {
// 		dhash = make([]uint64, imgSize/lenOfUnit+1)
// 	}
// 	idx := 0
// 	for y := 0; y < gray.Bounds().Dy(); y++ {
// 		for x := 0; x < gray.Bounds().Dx()-1; x++ {
// 			if gray.Pix[gray.PixOffset(x, y)] < gray.Pix[gray.PixOffset(x+1, y)] {
// 				indexOfArray := idx / lenOfUnit
// 				indexOfBit := lenOfUnit - idx%lenOfUnit - 1
// 				dhash[indexOfArray] |= 1 << uint(indexOfBit)
// 			}
// 			idx++
// 		}
// 	}
// 	return NewExtImageHash(dhash, DHash, imgSize), nil
// }