package main

import (
	"errors"
	"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, return the sharpness and contrast scores.
func (ts TurbiditySensor) Evaluate(imgs []gocv.Mat) (Results, error) {

	var result Results
	result.New(len(imgs))

	for i := range imgs {

		// Transform image.
		marker, err := ts.Transform(imgs[i])
		if err != nil {
			return result, fmt.Errorf("Image %v: %w", i, err)
		}

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

		// Evaluate image.
		out, err := ts.EvaluateImage(marker, edge, ts.k1, ts.k2, ts.alpha)
		if err != nil {
			return result, err
		}

		result.Update(out[0], out[1], float64(i*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 nil, fmt.Errorf("Dimensions not compatible (%v, %v)", k1, k2)
	}

	lStep := int(img.Rows() / k1)
	kStep := int(img.Cols() / k2)

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

			// Enhancement Measure Estimation (EME), provides a measure of the sharpness.
			err := ts.EvaluateBlock(edge, l, k, l+lStep, k+kStep, result, "EME", alpha)
			if err != nil {
				return nil, err
			}

			// AMEE, provides a measure of the contrast.
			err = ts.EvaluateBlock(img, l, k, l+lStep, k+kStep, result, "AMEE", alpha)
			if err != nil {
				return nil, err
			}
		}
	}

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

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

	return result, nil
}

// Evaluate a block within an 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 := -math.MaxFloat64
	min := math.MaxFloat64

	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 != -math.MaxFloat64 && min != math.MaxFloat64 && 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, errors.New("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, errors.New("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
}