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

package algo

import (
	"fmt"
	"log"
	"math"
	"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"
)

// neighbour is an alias for a float64 slice used for more natural handling of
// the underlying data.
type neighbour []float64

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 []neighbour, 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]

	// fmt.Println("select best\nlen(neighbours)", len(neighbours))

	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 {
	// TODO(me): have this passed to HillClimb and from there into this func.
	// return (bench.MaxNeighbourVariancePercent * 0.5) * (searchSpaceSize * 0.01)
	return bench.MaxNeighbourVariancePercent * (searchSpaceSize * 0.01)
}

func genNeighbours(n, dimens int, benchName string, origin []float64, neighbVals []neighbour) {
	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{Src: time.Now().UnixMicro()}
	uniform := distuv.Uniform{}
	// uniform.Src.Seed(uint64(time.Now().UnixNano()))
	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 _, 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
			}
			// fmt.Println("newMin:", newMin, "newMax", newMax)

			// v = append(v, uniform.Rand())
			v[i] = uniform.Rand()
		}
	}
}

// singleHillClimb runs a single iteration of SHC.
func singleHillClimb(
	dimens uint,
	neighbCount int,
	benchName string,
	oldVals []float64,
) ([]float64, float64) {
	neighbours := make([]neighbour, 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(neighbCount, 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.
// nolint: gocognit
func HillClimb(
	maxFES, benchMinIters, neighbours 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 / neighbours))
	localD = theD
	minIters = benchMinIters

	shcMeans := &stats.AlgoBenchMean{
		Algo:       "Stochastic Hill Climbing",
		BenchMeans: make([]stats.BenchMean, 0, len(localD)),
	}

	for _, dimens := range localD {
		stochasticHCStatDimX := &stats.Stats{
			Algo:       "Stochastic Hill Climbing",
			Dimens:     dimens,
			Iterations: minIters,
			// this is subject to change, see note on Stats struct in pkg
			// stats.
			Generations: fes,
		}
		funcStats := &stats.FuncStats{BenchName: benchFunc}
		benchFuncParams := bench.FunctionParams[benchFunc]
		dimXMean := &stats.BenchMean{
			Bench:       benchFunc,
			Dimens:      dimens,
			Iterations:  minIters,
			Generations: fes,
			Neighbours:  neighbours,
		}
		uniDist := distuv.Uniform{
			Min: benchFuncParams.Min(),
			Max: benchFuncParams.Max(),
		}

		printSHC(fmt.Sprintf("running bench \"%s\" for %dD with %d Neighbours",
			benchFunc, stochasticHCStatDimX.Dimens, neighbours),
		)

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

		// create and seed a source of preudo-randomness
		// rand.Seed(uint64(time.Now().UnixNano()))
		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.BenchResults[iter].Iteration = iter

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

			// reseed using current time.
			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),
					neighbours,
					benchFunc,
					initVals,
				); bR != 0 {
					bestVals = bV
					bestResult = bR
					initialised = true
				}
			}

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

				// first or any of the subsequent iterations.
				switch i {
				case 0:
					bestResult = r
					bestVals = v
				default:
					if r < bestResult {
						bestResult = r
						bestVals = v
					}
				}

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

				// this block makes sure we properly count func evaluations for
				// the purpose of correctly comparable plot comparison. i.e.
				// append the winning (current best) value neighbours-1 (the
				// first best is already saved at this point) times to
				// represent the fact that while evaluating the neighbours to
				// the current best value is taking place in the background,
				// the current best value itself is kept around and
				// symbolically saved as the best of the Generation.
				for x := 0; x < neighbours-1; x++ {
					funcStats.BenchResults[iter].Results = append(
						funcStats.BenchResults[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),
		)

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

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

		stochasticHCStats = append(stochasticHCStats, *stochasticHCStatDimX)

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

	// log.Printf("%+v\n", shcMeans)
	sort.Sort(shcMeans)

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

	ch <- stochasticHCStats
}