// Copyright 2023 wanderer <a_mirre at utb dot cz>
// SPDX-License-Identifier: GPL-3.0-or-later

package algo

import (
	"fmt"
	"log"
	"os"
	"sort"
	"time"

	"git.dotya.ml/wanderer/math-optim/bench"
	"git.dotya.ml/wanderer/math-optim/stats"
	"golang.org/x/exp/rand"
	"gonum.org/v1/gonum/stat/distuv"
)

func getRandomSearchLogPrefix() string {
	return " ***  random search:"
}

func fmtRandomSearchOut(input string) string {
	return getRandomSearchLogPrefix() + " " + input
}

func printRandomSearch(input string) {
	if _, err := fmt.Fprintln(os.Stderr, fmtRandomSearchOut(input)); err != nil {
		fmt.Fprintf(
			os.Stdout,
			getRandomSearchLogPrefix(),
			"error while printing to stderr: %q\n * original message was: %q",
			err, input,
		)
	}
}

func genValsRandomSearch(dimens uint, vals []float64, uniform *distuv.Uniform) {
	for i := uint(0); i < dimens; i++ {
		// using Uniform.Rand from gonum's stat/distuv package.
		// https://pkg.go.dev/gonum.org/v1/gonum/stat/distuv#Uniform.Rand
		// boundaries are already set at this point.
		vals[i] = uniform.Rand()
	}
}

// singleRandomSearch performs a single iteration of the 'RandomSearch'
// algorithm. it takes a couple of arguments:
// * dimens uint: number of dimensions of the objective function
// * f func([]float64) float64: bench func to execute (see Functions map in
// bench/functions.go)
// * uniformDist distuv.Uniform: uniform distribution representation with the
// min/max bounds already set to function-specific limits.
func singleRandomSearch(dimens uint, f func([]float64) float64, uniformDist distuv.Uniform) ([]float64, float64) {
	vals := make([]float64, dimens)

	genValsRandomSearch(dimens, vals, &uniformDist)

	// result of the bench function.
	res := f(vals)

	return vals, res
}

func RandomSearchNG(maxFES, benchMinIters int, theD []int, benchFunc string, ch chan []stats.Stats) {
	// perform basic sanity checks.
	if maxFES <= 0 {
		log.Fatalln(fmtRandomSearchOut("maxFES cannot be <= 0, bailing"))
	} else if benchMinIters <= 0 {
		log.Fatalln(fmtRandomSearchOut("benchMinIters cannot be <= 0, bailing"))
	} else if _, ok := bench.Functions[benchFunc]; !ok {
		log.Fatalln(fmtRandomSearchOut(
			"unknown benchFunc used: '" + benchFunc + "', bailing",
		))
	}

	for i := range theD {
		if theD[i] <= 0 {
			log.Fatalln(fmtRandomSearchOut(" no dimension in D can be <= 0, bailing"))
		}
	}

	// use func-local vars.
	var (
		fes               int
		localD            []int
		minIters          int
		randomSearchStats []stats.Stats
	)

	fes = maxFES
	localD = theD
	minIters = benchMinIters

	// create a continuous uniform distribution representation.
	uniformDist := &distuv.Uniform{
		Src: rand.NewSource(uint64(
			time.Now().
				UnixNano(),
		)),
	}

	rsMeans := &stats.AlgoBenchMean{
		Algo:       "Random Search",
		BenchMeans: make([]stats.BenchMean, 0, len(localD)),
	}

	// iterate over whatever was passed to us with theD - dimens slice.
	for _, dimens := range localD {
		randomSearchStatDimX := &stats.Stats{
			Algo:        "Random Search",
			Dimens:      dimens,
			Iterations:  minIters,
			Generations: fes,
		}
		funcStats := &stats.FuncStats{BenchName: benchFunc}
		benchFuncParams := bench.FunctionParams[benchFunc]
		dimXMean := &stats.BenchMean{
			Bench:       benchFunc,
			Dimens:      dimens,
			Iterations:  minIters,
			Generations: fes,
			// not applicable to Random Search...
			Neighbours: -1,
		}

		// set min/max bounds.
		uniformDist.Min = benchFuncParams.Min()
		uniformDist.Max = benchFuncParams.Max()

		printRandomSearch("running bench \"" + benchFunc + "\" for " +
			fmt.Sprint(randomSearchStatDimX.Dimens) + "D")

		funcStats.BenchResults = make([]stats.BenchRound, minIters)

		// perform the while dance 'minIters' times for "statistical relevance"
		for iter := 0; iter < minIters; iter++ {
			// have a fresh bestResult for each of the 'minIters' runs.
			var bestResult float64

			// set current iteration in funcStats.
			funcStats.BenchResults[iter].Iteration = iter

			// run the benchmarking function 'fes' times.
			for i := 0; i < fes; i++ {
				_, r := singleRandomSearch(
					uint(dimens),
					bench.Functions[benchFunc],
					*uniformDist,
				)

				// is a switch efficient, or should an if statement be used..?
				// TODO(me): perhaps benchmark this...
				switch i {
				case 0:
					// if it's our first, bestResult is the result we just got.
					bestResult = r
				default:
					// since we're minimising, the lower 'r' is the better.
					if r < bestResult {
						bestResult = r
					}
				}

				// save the 'best' result, since we only care about those.
				funcStats.BenchResults[iter].Results = append(
					funcStats.BenchResults[iter].Results,
					bestResult,
				)
			}
		}

		// get mean vals.
		dimXMean.MeanVals = stats.GetMeanVals(funcStats.BenchResults, fes)
		// save to funcStats, too.
		funcStats.MeanVals = dimXMean.MeanVals

		// save cumulative results of 'minIters' runs.
		randomSearchStatDimX.BenchFuncStats = append(
			randomSearchStatDimX.BenchFuncStats,
			*funcStats,
		)

		// save stats for each dimension to a stats slice.
		randomSearchStats = append(randomSearchStats, *randomSearchStatDimX)

		// save to AlgoMeans
		rsMeans.BenchMeans = append(rsMeans.BenchMeans, *dimXMean)
	}

	sort.Sort(rsMeans)

	// export AlgoMeans.
	mu.Lock()
	meanStats.AlgoMeans = append(meanStats.AlgoMeans, *rsMeans)
	mu.Unlock()

	ch <- randomSearchStats
}