package imageproc

import (
	"image"
	"math"

	"materialanalyzer/internal/model"
)

func ExtractTexture(img image.Image) model.TextureFeatures {
	small := ResizeToMax(img, 384)
	gray, w, h := grayscale(small)
	if w < 3 || h < 3 {
		return model.TextureFeatures{Algorithm: "entropy+sobel+lbp+glcm(local)"}
	}

	ent := entropy(gray)
	edgeDensity, frequency, orientationVar, primaryOrientation := sobelStats(gray, w, h)
	lbpUniformity := lbpUniformity(gray, w, h)
	glcm := glcmFeatures(gray, w, h, 8)

	return model.TextureFeatures{
		Entropy: round4(ent), TextureFrequency: round4(frequency), EdgeDensity: round4(edgeDensity),
		OrientationVariance: round4(orientationVar), PrimaryOrientationDeg: round4(primaryOrientation),
		LBPUniformity: round4(lbpUniformity), GLCM: glcm,
		Algorithm: "entropy+sobel+structure_tensor+lbp+glcm(local)",
	}
}

func grayscale(img image.Image) ([]uint8, int, int) {
	b := img.Bounds()
	w, h := b.Dx(), b.Dy()
	out := make([]uint8, w*h)
	i := 0
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			r, g, bb := rgba8(img.At(x, y))
			out[i] = uint8(math.Round(0.2126*float64(r) + 0.7152*float64(g) + 0.0722*float64(bb)))
			i++
		}
	}
	return out, w, h
}

func entropy(gray []uint8) float64 {
	if len(gray) == 0 {
		return 0
	}
	var hist [256]int
	for _, v := range gray {
		hist[v]++
	}
	var e float64
	denom := float64(len(gray))
	for _, c := range hist {
		if c == 0 {
			continue
		}
		p := float64(c) / denom
		e -= p * math.Log2(p)
	}
	return e
}

func sobelStats(gray []uint8, w, h int) (edgeDensity, frequency, orientationVariance, primaryOrientation float64) {
	var edges, n int
	var sumMag, sumWeight, sumCos, sumSin float64
	for y := 1; y < h-1; y++ {
		for x := 1; x < w-1; x++ {
			gx := -int(gray[(y-1)*w+x-1]) + int(gray[(y-1)*w+x+1]) -
				2*int(gray[y*w+x-1]) + 2*int(gray[y*w+x+1]) -
				int(gray[(y+1)*w+x-1]) + int(gray[(y+1)*w+x+1])
			gy := -int(gray[(y-1)*w+x-1]) - 2*int(gray[(y-1)*w+x]) - int(gray[(y-1)*w+x+1]) +
				int(gray[(y+1)*w+x-1]) + 2*int(gray[(y+1)*w+x]) + int(gray[(y+1)*w+x+1])
			mag := math.Hypot(float64(gx), float64(gy))
			if mag > 72 {
				edges++
			}
			if mag > 8 {
				angle := math.Atan2(float64(gy), float64(gx)) + math.Pi/2
				sumCos += mag * math.Cos(2*angle)
				sumSin += mag * math.Sin(2*angle)
				sumWeight += mag
			}
			sumMag += mag
			n++
		}
	}
	if n == 0 {
		return 0, 0, 0, 0
	}
	edgeDensity = float64(edges) / float64(n)
	frequency = math.Min(1, sumMag/(float64(n)*255.0))
	if sumWeight == 0 {
		return edgeDensity, frequency, 1, 0
	}
	coherence := math.Hypot(sumCos, sumSin) / sumWeight
	orientationVariance = 1 - coherence
	primary := 0.5 * math.Atan2(sumSin, sumCos) * 180 / math.Pi
	if primary < 0 {
		primary += 180
	}
	return edgeDensity, frequency, orientationVariance, primary
}

func lbpUniformity(gray []uint8, w, h int) float64 {
	if w < 3 || h < 3 {
		return 0
	}
	var hist [256]int
	var n int
	for y := 1; y < h-1; y++ {
		for x := 1; x < w-1; x++ {
			c := gray[y*w+x]
			code := 0
			if gray[(y-1)*w+x-1] >= c {
				code |= 1 << 7
			}
			if gray[(y-1)*w+x] >= c {
				code |= 1 << 6
			}
			if gray[(y-1)*w+x+1] >= c {
				code |= 1 << 5
			}
			if gray[y*w+x+1] >= c {
				code |= 1 << 4
			}
			if gray[(y+1)*w+x+1] >= c {
				code |= 1 << 3
			}
			if gray[(y+1)*w+x] >= c {
				code |= 1 << 2
			}
			if gray[(y+1)*w+x-1] >= c {
				code |= 1 << 1
			}
			if gray[y*w+x-1] >= c {
				code |= 1
			}
			hist[code]++
			n++
		}
	}
	if n == 0 {
		return 0
	}
	var uniformity float64
	for _, c := range hist {
		p := float64(c) / float64(n)
		uniformity += p * p
	}
	return uniformity
}

func glcmFeatures(gray []uint8, w, h, levels int) model.GLCMFeatures {
	if w < 2 || h < 2 || levels <= 1 {
		return model.GLCMFeatures{Levels: levels}
	}
	matrix := make([]float64, levels*levels)
	add := func(a, b uint8) {
		i := min(levels-1, int(a)*levels/256)
		j := min(levels-1, int(b)*levels/256)
		matrix[i*levels+j]++
		matrix[j*levels+i]++
	}
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			v := gray[y*w+x]
			if x+1 < w {
				add(v, gray[y*w+x+1])
			}
			if y+1 < h {
				add(v, gray[(y+1)*w+x])
			}
		}
	}
	var total float64
	for _, v := range matrix {
		total += v
	}
	if total == 0 {
		return model.GLCMFeatures{Levels: levels}
	}
	for i := range matrix {
		matrix[i] /= total
	}

	var meanI, meanJ float64
	for i := 0; i < levels; i++ {
		for j := 0; j < levels; j++ {
			p := matrix[i*levels+j]
			meanI += float64(i) * p
			meanJ += float64(j) * p
		}
	}
	var contrast, homogeneity, energy, varI, varJ, corr float64
	for i := 0; i < levels; i++ {
		for j := 0; j < levels; j++ {
			p := matrix[i*levels+j]
			d := float64(i - j)
			contrast += d * d * p
			homogeneity += p / (1 + math.Abs(d))
			energy += p * p
			varI += (float64(i) - meanI) * (float64(i) - meanI) * p
			varJ += (float64(j) - meanJ) * (float64(j) - meanJ) * p
			corr += (float64(i) - meanI) * (float64(j) - meanJ) * p
		}
	}
	if varI > 0 && varJ > 0 {
		corr /= math.Sqrt(varI * varJ)
	} else {
		corr = 0
	}
	return model.GLCMFeatures{
		Levels: levels, Contrast: round4(contrast), Homogeneity: round4(homogeneity),
		Energy: round4(energy), Correlation: round4(corr),
	}
}