package main

import (
	"fmt"

	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/plotutil"
	"gonum.org/v1/plot/vg"

	"gocv.io/x/gocv"
)

func main() {
	n_images := 6
	n_samples := 10

	// Load template
	template := gocv.IMRead("template.jpg", gocv.IMReadGrayScale)
	standard := gocv.IMRead("standard.jpg", gocv.IMReadGrayScale)

	images := make([][]gocv.Mat, n_images)

	//Load images
	for i := 0; i < n_images; i++ {
		images[i] = make([]gocv.Mat, n_samples)
		for j := 0; j < n_samples; j++ {
			images[i][j] = gocv.IMRead(fmt.Sprintf("images/t-%v/000%v.jpg", i, j), gocv.IMReadGrayScale)
		}
	}

	// Create turbidity sensor
	ts := TurbiditySensor{template: template, standard: standard, k1: 90, k2: 90, scale: 5.0, alpha: 1.0}

	var final_result Results
	final_result.New(n_images)

	// Score each image by calculating the average score from camera burst
	for i := range images {

		//Evaluate camera burst
		sample_result, err := ts.Evaluate(n_samples, images[i])

		if err != nil {
			fmt.Println(err)
		}

		// Append the average result from camera burst
		final_result.Append(Average(sample_result.saturation), Average(sample_result.contrast), float64((i+1)*10), i)
	}

	// Plot the final results
	err := PlotResults(final_result.turbidity, Normalize(final_result.saturation), Normalize(final_result.contrast))

	fmt.Println(final_result.saturation)
	fmt.Println(final_result.contrast)

	if err != nil {
		fmt.Println(err)
	}
}

// PLOTTING FUNCTIONS

// Normalize values in a slice between 0 and 1
func Normalize(slice []float64) []float64 {

	max := -1e10
	min := 1e10

	out := make([]float64, len(slice))

	if len(slice) <= 1 {
		return slice
	}

	for i := range slice {
		if slice[i] > max {
			max = slice[i]
		}
		if slice[i] < min {
			min = slice[i]
		}
	}

	for i := range slice {
		out[i] = (slice[i] - min) / (max - min)
	}

	return out
}

// Return the average of a slice
func Average(slice []float64) float64 {

	out := 0.0

	for i := range slice {
		out += slice[i]
	}

	return out / float64(len(slice))
}

func PlotResults(x, saturation, contrast []float64) error {

	err := plotToFile(
		"Results",
		"Turbidity (Almond Milk) (ml)",
		"Score",
		func(p *plot.Plot) error {
			return plotutil.AddLinePoints(p,
				"AMEE", plotterXY(x, contrast),
				"EME", plotterXY(x, saturation),
			)
		},
	)

	if err != nil {
		return fmt.Errorf("Could not plot results: %v", err)
	}

	return nil
}

// plotToFile creates a plot with a specified name and x&y titles using the
// provided draw function, and then saves to a PNG file with filename of name.
func plotToFile(name, xTitle, yTitle string, draw func(*plot.Plot) error) error {
	p := plot.New()

	p.Title.Text = name
	p.X.Label.Text = xTitle
	p.Y.Label.Text = yTitle

	err := draw(p)
	if err != nil {
		return fmt.Errorf("could not draw plot contents: %w", err)
	}

	if err := p.Save(15*vg.Centimeter, 15*vg.Centimeter, "plots/"+name+".png"); err != nil {
		return fmt.Errorf("could not save plot: %w", err)
	}
	return nil
}

// plotterXY provides a plotter.XYs type value based on the given x and y data.
func plotterXY(x, y []float64) plotter.XYs {

	xy := make(plotter.XYs, len(x))

	for i := range x {
		xy[i].X = x[i]
		xy[i].Y = y[i]
	}
	return xy
}