//go:build !ignore
// +build !ignore

/*
DESCRIPTION
  Holds the turbidity sensor struct. Can evaluate 4x4 chessboard markers
  in an image to measure the sharpness and contrast.

AUTHORS
  Russell Stanley <russell@ausocean.org>

LICENSE
  Copyright (C) 2020 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  in gpl.txt.  If not, see http://www.gnu.org/licenses.
*/

package main

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

	"gocv.io/x/gocv"
)

// TurbiditySensor is a software based turbidity sensor that uses CV to determine saturation and constrast level
// of a chessboard-like target submerged in water that can be correlated to turbidity/visibility values.
type TurbiditySensor struct {
	template, templateCorners gocv.Mat
	standard, standardCorners gocv.Mat
	k1, k2, sobelFilterSize   int
	scale, alpha              float64
}

// NewTurbiditySensor constructor for a turbidity sensor.
func NewTurbiditySensor(template, standard gocv.Mat, k1, k2, sobelFilterSize int, scale, alpha float64) (*TurbiditySensor, error) {
	ts := new(TurbiditySensor)
	templateCorners := gocv.NewMat()
	standardCorners := gocv.NewMat()

	// Validate template image is not empty and has valid corners.
	if template.Empty() {
		return nil, errors.New("template image is empty.")
	}
	if !gocv.FindChessboardCorners(template, image.Pt(3, 3), &templateCorners, gocv.CalibCBNormalizeImage) {
		return nil, errors.New("could not find corners in template image")
	}
	ts.template = template
	ts.templateCorners = templateCorners

	// Validate standard image is not empty and has valid corners.
	if standard.Empty() {
		return nil, errors.New("standard image is empty.")
	}
	if !gocv.FindChessboardCorners(standard, image.Pt(3, 3), &standardCorners, gocv.CalibCBNormalizeImage) {
		return nil, errors.New("could not find corners in standard image")
	}
	ts.standard = standard
	ts.standardCorners = standardCorners

	ts.k1, ts.k2, ts.sobelFilterSize = k1, k2, sobelFilterSize
	ts.alpha, ts.scale = alpha, scale

	return ts, nil
}

// Evaluate, given a slice of images, return the sharpness and contrast scores.
func (ts TurbiditySensor) Evaluate(imgs []gocv.Mat) (*Results, error) {
	result, err := NewResults(len(imgs))
	if err != nil {
		return result, err
	}

	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)

		gocv.IMWrite("marker.jpg", marker)
		gocv.IMWrite("edge.jpg", edge)

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

		result.Update(sharpScore, contScore, float64(i), i)
	}

	return result, nil
}

// EvaluateImage will evaluate image sharpness and contrast using blocks of size k1 by k2. Return the respective scores.
func (ts TurbiditySensor) EvaluateImage(img, edge gocv.Mat) (float64, float64, error) {
	var sharpness float64
	var contrast float64

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

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

	for l := 0; l < img.Rows(); l += lStep {
		for k := 0; k < img.Cols(); k += kStep {
			// Enhancement Measure Estimation (EME), provides a measure of the sharpness.
			sharpValue := ts.evaluateBlockEME(edge, l, k, l+lStep, k+kStep)
			sharpness += sharpValue

			// AMEE, provides a measure of the contrast.
			contValue := ts.evaluateBlockAMEE(img, l, k, l+lStep, k+kStep)
			contrast += contValue
		}
	}

	// Scale EME based on block size.
	sharpness = 2.0 / (float64(ts.k1 * ts.k2)) * sharpness

	// Scale and flip AMEE based on block size.
	contrast = -1.0 / (float64(ts.k1 * ts.k2)) * contrast

	return sharpness, contrast, nil
}

// minMax returns the max and min pixel values of an image block.
func (ts TurbiditySensor) minMax(img gocv.Mat, xStart, yStart, xEnd, yEnd int) (float64, float64) {
	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
			}
		}
	}
	return max, min
}

// evaluateBlockEME will evaluate an image block and return the value to be added to the sharpness result.
func (ts TurbiditySensor) evaluateBlockEME(img gocv.Mat, xStart, yStart, xEnd, yEnd int) float64 {
	max, min := ts.minMax(img, xStart, yStart, xEnd, yEnd)

	// Blocks which have no information are ignored.
	if max != -math.MaxFloat64 && min != math.MaxFloat64 && max != min {
		return math.Log(max / min)
	}
	return 0.0
}

// evaluateBlockAMEE will evaluate an image block and return the value to be added to the contrast result.
func (ts TurbiditySensor) evaluateBlockAMEE(img gocv.Mat, xStart, yStart, xEnd, yEnd int) float64 {
	max, min := ts.minMax(img, xStart, yStart, xEnd, yEnd)

	// Blocks which have no information are ignored.
	if max != -math.MaxFloat64 && min != math.MaxFloat64 && max != min {
		contrast := (max + min) / (max - min)
		return math.Pow(ts.alpha*(contrast), ts.alpha) * math.Log(contrast)
	}
	return 0.0
}

// transform will search img 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()
	imgCorners := gocv.NewMat()

	// Check image is valid.
	if img.Empty() {
		return out, errors.New("image is empty, cannot transform")
	}
	// Check image for corners, if non can be found corners will be set to default value.
	if !gocv.FindChessboardCorners(img, image.Pt(3, 3), &imgCorners, gocv.CalibCBFastCheck) {
		imgCorners = ts.standardCorners
	}

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

	return out, nil
}

// sobel will apply sobel filter to an image 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
}