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

package ga

import (
	"sort"
	"time"

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

// SOMAT3A holds the settings for an instance of SOMA T3A algorithm.
// nolint: unused
type SOMAT3A struct {
	// Generations is the number of generations to evolve for. Disable limit
	// with -1.
	Generations int
	// BenchMinIters is the number of iterations that the bench function will
	// be re-run.
	BenchMinIters int
	// Dimensions in which to look for a solution.
	Dimensions []int
	// NP is the initial population size.
	NP int
	// K is the number of individuals to choose the leader from.
	K int
	// M denotes how many individuals are picked from the population to form a
	// `team` during the organisation phase, can be thought of as "team size".
	M int
	// N is the number of best individuals in each team selected for actual
	// migration.
	N int
	// Njumps is the fixed number of jumps that each chosen migrating
	// individual performs on their way to the leader.
	Njumps int
	// BenchName is the human-friendly name of the benchmarking function.
	BenchName string
	// ch is a channel for writing back computed results.
	ch chan []stats.Stats
	// chAlgoMeans is a channel for writing back algo means.
	chAlgoMeans chan *stats.AlgoBenchMean

	// rng is a random number generator.
	rng distuv.Uniform

	// initialised denotes the initialisation state of the struct.
	initialised bool
}

var somat3aLogger = newLogger(" ***  SOMA T3A")

// Init performs checks on the input params and initialises an instance of
// SOMAT3A.
func (s *SOMAT3A) Init(
	generations, benchMinIters, np, k, m, n, njumps int,
	dimensions []int,
	bench string,
	ch chan []stats.Stats,
	chAlgoMeans chan *stats.AlgoBenchMean,
) error {
	switch {
	case generations == 0:
		return errGenerationsIsZero

	case generations == -1:
		somat3aLogger.Println("Generations is '-1', disabling generation limits")

	case benchMinIters < 1:
		return errBenchMinItersIsLTOne

	case np == 0:
		return errNPIsZero

	case k == 0:
		return errKIsZero

	case m == 0:
		return errMIsZero

	case n == 0:
		return errNIsZero

	case njumps == 0:
		return errNjumpsIsZero

	case len(dimensions) == 0:
		return errDimensionsUnset

	case bench == "":
		return errBenchUnset
	}

	s.Generations = generations
	s.BenchMinIters = benchMinIters
	s.NP = np
	s.K = k
	s.M = m
	s.N = n
	s.Njumps = njumps

	s.Dimensions = dimensions
	s.BenchName = bench
	s.ch = ch
	s.chAlgoMeans = chAlgoMeans

	s.initialised = true

	gaLogger.Printf("%s init done for bench %s", s.getAlgoName(), s.BenchName)

	return nil
}

// Run runs the SOMA T3A algo for each of `s.Dimensions` with benchmarking
// function `s.BenchName`, repeats the run `s.BenchMinIters` times and returns
// the stats via chans.
func (s *SOMAT3A) Run() {
	if !s.initialised {
		somat3aLogger.Fatalln(s.getAlgoName(), "needs to be initialised before calling Run(), exiting..")
	}

	var somat3aStats []stats.Stats
	// testing purposes.
	// s.BenchMinIters = 5

	somat3aMeans := &stats.AlgoBenchMean{
		Algo:       s.getAlgoName(),
		BenchMeans: make([]stats.BenchMean, 0, len(s.Dimensions)),
	}

	for _, dim := range s.Dimensions {
		maxFES := cec2020.GetMaxFES(dim)
		// TODO(me): check the slides for the formula to calculate this.
		// 4 fes per member per njump.
		// fesPerIter := int(float64(maxFES / (4 * s.NP * s.Njumps)))
		fesPerIter := int(float64(maxFES / (4 * s.NP)))
		somat3aStatDimX := &stats.Stats{
			Algo:        s.getAlgoName(),
			Dimens:      dim,
			Iterations:  s.BenchMinIters,
			Generations: maxFES,
			NP:          s.NP,
		}
		funcStats := &stats.FuncStats{BenchName: s.BenchName}
		dimXMean := &stats.BenchMean{
			Bench:       s.BenchName,
			Dimens:      dim,
			Iterations:  s.BenchMinIters,
			Generations: maxFES,
			Neighbours:  -1,
		}
		uniDist := distuv.Uniform{
			Min: cec2020.SearchRange.Min(),
			Max: cec2020.SearchRange.Max(),
		}

		if maxFES == -1 {
			somat3aLogger.Fatalf("could not get maxFES for current dim (%d), bailing", dim)
		}

		somat3aLogger.Printf("running bench \"%s\" for %dD, maxFES: %d\n",
			s.BenchName, dim, maxFES,
		)

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

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

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

		for iter := 0; iter < s.BenchMinIters; iter++ {
			somat3aLogger.Printf("run: %d, bench: %s, %dD, started at %s",
				iter, s.BenchName, dim, start,
			)

			funcStats.BenchResults[iter].Iteration = iter

			uniDist.Src = src
			// uniDist.Src = rand.NewSource(uint64(rand.Int63()))

			// create a population with known params.
			pop := newSOMAT3APopulation(s.BenchName, s.NP, s.K, s.M, s.N, s.Njumps, dim)

			// set population seed.
			pop.Seed = uint64(time.Now().UnixNano())
			// initialise the population.
			pop.Init()

			bestResult := s.evaluate(pop)

			// the core.
			for i := 0; i < fesPerIter; i++ {
				funcStats.BenchResults[iter].Results = append(
					funcStats.BenchResults[iter].Results,
					bestResult,
				)

				// call evolve where the inner workngs of SOMA T3A run.
				s.evolve(pop)

				// take measurements of current population fitness.
				r := s.evaluate(pop)

				// save if better.
				if r < bestResult {
					bestResult = r
				}

				// 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 4*NP-1 (the first
				// best is already saved at this point) times to represent the
				// fact that while evaluating (and comparing) other population
				// individuals 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 < 4*s.NP-1; x++ {
					funcStats.BenchResults[iter].Results = append(
						funcStats.BenchResults[iter].Results,
						bestResult,
					)
				}
			}
		}

		elapsed := time.Since(start)

		somat3aLogger.Printf("completed: bench: %s, %dD, computing took %s\n",
			s.BenchName, dim, elapsed,
		)

		// get mean vals.
		dimXMean.MeanVals = stats.GetMeanVals(funcStats.BenchResults, maxFES)
		funcStats.MeanVals = dimXMean.MeanVals

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

		somat3aStats = append(somat3aStats, *somat3aStatDimX)

		somat3aMeans.BenchMeans = append(somat3aMeans.BenchMeans, *dimXMean)
	}

	sort.Sort(somat3aMeans)

	s.chAlgoMeans <- somat3aMeans
	s.ch <- somat3aStats
}

func (s *SOMAT3A) evaluate(pop *SOMAT3APopulation) float64 {
	bestIndividual := pop.Population[pop.GetBestIdx()]

	return pop.ProblemFunc(bestIndividual.CurX)
}

func (s *SOMAT3A) evolve(pop *SOMAT3APopulation) {
	// organise.
	leader, migrants := pop.organise()
	// migrate.
	pop.migrate(leader, migrants)
	// actualise (position, best values..).
	pop.actualise(migrants)
}

func (s *SOMAT3A) getAlgoName() string {
	return "SOMA T3A"
}

// NewSOMAT3A returns a pointer to a new, uninitialised SOMAT3A instance.
func NewSOMAT3A() *SOMAT3A {
	return &SOMAT3A{}
}