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

package algo

import (
	"fmt"
	"log"
	"math"
	"os"
	"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 getSHCLogPrefix() string {
	return " ***  stochastic hill climbing:"
}

func fmtSHCOut(input string) string {
	return getSHCLogPrefix() + " " + input
}

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

func initValsSHC(dimens uint, vals []float64, uniform *distuv.Uniform) {
	for i := uint(0); i < dimens; i++ {
		vals[i] = uniform.Rand()
	}
}

// selectBestNeighbour evaluates the randomly generated neighbours passed to it
// as a parameter and points to the best one (the one yielding the greatest
// progress in a favourable way).
func selectBestNeighbour(neighbours [][]float64, benchFuncName string) ([]float64, float64) {
	// bestNeighbour is selected based on the value of the bench function it yields.
	var bestNeighbour []float64

	// bestSolution is used to store the best bench func name value.
	var bestSolution float64

	// f is the actual bench function, based on the name.
	f := bench.Functions[benchFuncName]

	for i, v := range neighbours {
		switch i {
		case 0:
			// the first one is automatically the best.
			bestSolution = f(v)
			bestNeighbour = v
		default:
			// calculate the value of the bench function for neighbour and
			// decide if this was the best we've seen so far or not.
			if solution := f(v); solution < bestSolution {
				bestSolution = solution
				bestNeighbour = v
			}
		}
	}

	return bestNeighbour, bestSolution
}

// getBenchSearchSpaceSize calculates and returns the absolute distance between
// min,max for a given bench func.
func getBenchSearchSpaceSize(benchName string) float64 {
	params := bench.FunctionParams[benchName]

	return math.Abs(params.Min() - params.Max())
}

// getAllowedTweak calculates the value of allowed movement correction based on
// the bench func's searchSpaceSize, also taking into account the globally
// applicable MaxNeighbourVariance (in %) value from the bench pkg.
// the allowedTweak value should therefore the amount to at most 10% of the
// searchSpaceSize. TODO(me): floor this down.
func getAllowedTweak(searchSpaceSize float64) float64 {
	return bench.MaxNeighbourVariancePercent * (searchSpaceSize * 0.01)
}

func genNeighbours(n, dimens int, benchName string, origin []float64, neighbVals [][]float64) {
	if dimens < 1 {
		log.Fatalf("error: neighbour count set to: %d; want len(neighbVals) > 0, bailing\n", n)
	}

	// create a new representation of the uniform distribution with the bounds
	// unset for the moment, since we need to find out what exactly those ought
	// to be (and we will - a couple of lines later).
	uniform := distuv.Uniform{}
	params := bench.FunctionParams[benchName]
	// get bench function bounds.
	benchMin := params.Min()
	benchMax := params.Max()
	searchSpaceSize := getBenchSearchSpaceSize(benchName)
	// allowedTweak tells how big of a tweak to each side is allowed.
	allowedTweak := getAllowedTweak(searchSpaceSize)

	// seed the src.
	uniform.Src = rand.NewSource(uint64(time.Now().UnixNano()))

	for _, v := range neighbVals {
		for i := 0; i < dimens; i++ {
			newMin := origin[i] - allowedTweak

			// if newMin is not within the bounds, fall back to benchMin.
			if newMin < benchMin {
				uniform.Min = newMin
			}

			// if newMax is not within the bounds, fall back to benchMax.
			newMax := origin[i] + allowedTweak

			if newMin > benchMax {
				uniform.Max = newMax
			}

			v[i] = uniform.Rand()
		}
	}
}

// singleHillClimb runs a single iteration of SHC.
func singleHillClimb(dimens uint, benchName string, oldVals []float64) ([]float64, float64) {
	neighbours := make([][]float64, dimens)
	// prealloc the slice the size of dimens (and oldVals, should be the same).
	vals := make([]float64, len(oldVals))

	for i := range neighbours {
		neighbours[i] = vals
	}

	genNeighbours(bench.Neighbourhood, int(dimens), benchName, oldVals, neighbours)

	newVals, newRes := selectBestNeighbour(neighbours, benchName)

	return newVals, newRes
}

// HillClimb performs 30 iterations of SHC (30 singleHillClimb func calls
// internally) to establish a semi-relevant statistical baseline, and reports
// the results of the computation.
func HillClimb(maxFES, benchMinIters int, theD []int, benchFunc string, ch chan []stats.Stats) {
	if maxFES <= 0 {
	} else if benchMinIters <= 0 {
		log.Fatalln(fmtSHCOut("benchMinIters cannot be <= 0, bailing"))
	} else if _, ok := bench.Functions[benchFunc]; !ok {
		log.Fatalln(fmtSHCOut(
			"unknown benchFunc used: '" + benchFunc + "', bailing",
		))
	}

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

	var (
		fes int
		// since Stochastic Hill Climbing performs multiple cost function
		// evaluations during each run (in search of a better neighbouring
		// value), fesPerIter stores how many evaluations are allowed during
		// each run based on the number of allowed evaluations in total (passed
		// into function in maxFES).
		fesPerIter        int
		localD            []int
		minIters          int
		stochasticHCStats []stats.Stats
	)

	fes = maxFES
	// no need to call math.Floor() apparently, as `int()` does that implicitly.
	fesPerIter = int(float64(fes / bench.Neighbourhood))
	localD = theD
	minIters = benchMinIters

	for _, dimens := range localD {
		stochasticHCStatDimX := &stats.Stats{
			Algo:       "Stochastic Hill Climbing",
			Dimens:     dimens,
			Iterations: minIters,
			// Generations for Stochastic Hill Climbing may vary based on both
			// the maxFES and bench.Neighbourhood, since that is how we arrive
			// at the value of fesPerIter.
			Generations: fesPerIter,
		}
		funcStats := &stats.FuncStats{BenchName: benchFunc}
		benchFuncParams := bench.FunctionParams[benchFunc]
		uniDist := distuv.Uniform{
			Min: benchFuncParams.Min(),
			Max: benchFuncParams.Max(),
		}

		printSHC("running bench \"" + benchFunc + "\" for " +
			fmt.Sprint(stochasticHCStatDimX.Dimens) + "D")

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

		// create a source of preudo-randomness.
		src := rand.NewSource(uint64(rand.Int63()))
		// src := rand.NewSource(uint64(time.Now().UnixNano()))

		// track execution time.
		start := time.Now()

		for iter := 0; iter < minIters; iter++ {
			funcStats.Solution[iter].Iteration = iter

			// create and stochastically populate the vals slice.
			initVals := make([]float64, dimens)

			uniDist.Src = src

			var bestResult float64

			var bestVals []float64

			initialised := false

			for !initialised {
				// generate initial vals, make sure they are not 0s.
				initValsSHC(uint(dimens), initVals, &uniDist)

				if bV, bR := singleHillClimb(
					uint(dimens),
					benchFunc,
					initVals,
				); bR != 0 {
					bestVals = bV
					bestResult = bR
					initialised = true
				}
			}

			for i := 0; i < fesPerIter; i++ {
				v, r := singleHillClimb(
					uint(dimens),
					benchFunc,
					bestVals,
				)

				// first or any other iteration.
				switch i {
				case 0:
					bestResult = r
					bestVals = v
				default:
					if r < bestResult {
						bestResult = r
						bestVals = v
					}
				}

				funcStats.Solution[iter].Results = append(
					funcStats.Solution[iter].Results,
					bestResult,
				)
			}
		}

		elapsed := time.Since(start)

		printSHC("computing " + fmt.Sprint(benchMinIters) + "I of bench \"" +
			benchFunc + "\" for " + fmt.Sprint(dimens) + "D took " + fmt.Sprint(elapsed),
		)

		stochasticHCStatDimX.BenchFuncStats = append(
			stochasticHCStatDimX.BenchFuncStats,
			*funcStats,
		)

		stochasticHCStats = append(stochasticHCStats, *stochasticHCStatDimX)
	}

	ch <- stochasticHCStats
}