package main

import (
	"fmt"
	"image"
	"math"

	"gocv.io/x/gocv"
)

// TURBIDITY SENSOR
type TurbiditySensor struct {
	template, standard gocv.Mat
	k1, k2             int //block size
	alpha, scale       float64
}

// Given a slice of test images of size n, return the sharpness and contrast scores
func (ts TurbiditySensor) Evaluate(n int, images []gocv.Mat) (Results, error) {

	var result Results
	result.New(n)

	for i := range images {

		// Transform image
		marker, err := ts.Transform(images[i])

		if err != nil {
			return result, fmt.Errorf("Image %v: %v", i, err)
		}

		// Apply sobel filter
		edge := ts.Sobel(marker, ts.scale)

		// Evaluate
		out, err := ts.EvaluateImage(marker, edge, ts.k1, ts.k2, ts.alpha)

		if err != nil {
			return result, err
		}

		result.Append(out[0], out[1], float64((i+1)*10), i)
	}

	return result, nil
}

// Evaluate image sharpness and contrast using blocks of size k1 by k2. Return a slice of the respective scores
func (ts TurbiditySensor) EvaluateImage(img, edge gocv.Mat, k1, k2 int, alpha float64) ([]float64, error) {

	// Slice to store results
	result := make([]float64, 2) // [0.0, 0.0]

	if img.Rows()%k1 != 0 || img.Cols()%k2 != 0 {
		return result, fmt.Errorf("Dimensions not compatible (%v, %v)", k1, k2)
	}

	l_step := int(img.Rows() / k1)
	k_step := int(img.Cols() / k2)

	for l := 0; l < img.Rows(); l = l + l_step {
		for k := 0; k < img.Cols(); k = k + k_step {

			//EME
			err := ts.EvaluateBlock(edge, l, k, l+l_step, k+k_step, result, "EME", alpha)

			if err != nil {
				return result, err
			}

			//AMEE
			err = ts.EvaluateBlock(img, l, k, l+l_step, k+k_step, result, "AMEE", alpha)

			if err != nil {
				return result, err
			}
		}
	}

	// EME
	result[0] = 2.0 / (float64(k1) * float64(k1)) * result[0]

	// AMEE
	result[1] = -1.0 / (float64(k1) * float64(k1)) * result[1]

	return result, nil
}

// Evaluate a block within and image and add to to the result slice
func (ts TurbiditySensor) EvaluateBlock(img gocv.Mat, l1, k1, l2, k2 int, result []float64, operation string, alpha float64) error {

	max := -1e10
	min := 1e10

	for i := l1; i < l2; i++ {
		for j := k1; j < k2; j++ {

			value := float64(img.GetUCharAt(i, j))

			// Check max/min conditions, zero values are ignored
			if value > max && value != 0.0 {
				max = value
			}
			if value < min && value != 0.0 {
				min = value
			}
		}
	}

	// Blocks which have no information are ignored
	if max != -1e10 && min != 1e10 && max != min {
		if operation == "EME" {

			result[0] += math.Log(max / min)

		} else if operation == "AMEE" {

			contrast := (max + min) / (max - min)
			result[1] += math.Pow(alpha*(contrast), alpha) * math.Log(contrast)

		} else {
			return fmt.Errorf("Invalid operation: %v", operation)
		}
	}
	return nil
}

// Search image for matching template. Returns the transformed image  which best match the template
func (ts TurbiditySensor) Transform(img gocv.Mat) (gocv.Mat, error) {

	out := gocv.NewMat()
	mask := gocv.NewMat()

	corners_img := gocv.NewMat()
	corners_template := gocv.NewMat()

	// Find corners in image
	if !gocv.FindChessboardCorners(img, image.Pt(3, 3), &corners_img, gocv.CalibCBNormalizeImage) {

		// Apply default if transformation fails
		// fmt.Println("Corner detection failed applying standard transformation")

		if !gocv.FindChessboardCorners(ts.standard, image.Pt(3, 3), &corners_img, gocv.CalibCBNormalizeImage) {
			return out, fmt.Errorf("Could not find corners in default image")
		}
	}

	// Find corners in template
	if !gocv.FindChessboardCorners(ts.template, image.Pt(3, 3), &corners_template, gocv.CalibCBNormalizeImage) {
		return out, fmt.Errorf("Could not find corners in template")
	}

	// Find and apply transformation
	H := gocv.FindHomography(corners_img, &corners_template, gocv.HomograpyMethodRANSAC, 3.0, &mask, 2000, 0.995)
	gocv.WarpPerspective(img, &out, H, image.Pt(ts.template.Rows(), ts.template.Cols()))

	return out, nil
}

// Apply sobel filter to an image with a given scale and return the result
func (ts TurbiditySensor) Sobel(img gocv.Mat, scale float64) gocv.Mat {

	dx := gocv.NewMat()
	dy := gocv.NewMat()
	sobel := gocv.NewMat()

	// Apply filter
	gocv.Sobel(img, &dx, gocv.MatTypeCV64F, 0, 1, 3, scale, 0.0, gocv.BorderConstant)
	gocv.Sobel(img, &dy, gocv.MatTypeCV64F, 1, 0, 3, scale, 0.0, gocv.BorderConstant)

	// Convert to unsigned
	gocv.ConvertScaleAbs(dx, &dx, 1.0, 0.0)
	gocv.ConvertScaleAbs(dy, &dy, 1.0, 0.0)

	// Add x and y components
	gocv.AddWeighted(dx, 0.5, dy, 0.5, 0, &sobel)

	return sobel
}