package main

import (
	"errors"
	"fmt"
	"image"
	"math"

	"gocv.io/x/gocv"
)

// Turbidity Sensor.
type TurbiditySensor struct {
	template, standard      gocv.Mat
	k1, k2, sobelFilterSize int
	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)

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

		result.update(scores[0], scores[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) ([]float64, error) {

	result := make([]float64, 2) // [0.0, 0.0]

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

	lStep := int(img.Rows() / ts.k1)
	kStep := int(img.Cols() / ts.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", ts.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", ts.alpha)
			if err != nil {
				return nil, err
			}
		}
	}

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

	// AMEE.
	result[1] = -1.0 / (float64(ts.k1) * float64(ts.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, xStart, yStart, xEnd, yEnd int, result []float64, operation string, alpha float64) error {

	max := -math.MaxFloat64
	min := math.MaxFloat64

	for i := xStart; i < xEnd; i++ {
		for j := yStart; j < yEnd; 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(ts.standard, 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) gocv.Mat {

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

	// Apply filter.
	gocv.Sobel(img, &dx, gocv.MatTypeCV64F, 0, 1, ts.sobelFilterSize, ts.scale, 0.0, gocv.BorderConstant)
	gocv.Sobel(img, &dy, gocv.MatTypeCV64F, 1, 0, ts.sobelFilterSize, ts.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
}