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

package ga

import (
	"git.dotya.ml/wanderer/math-optim/bench/cec2020"
	"golang.org/x/exp/rand"

	"gonum.org/v1/gonum/stat/distuv"
)

// PopulationIndividual represents a single population individual.
type SOMAT3APopulationIndividual struct {
	CurX      DecisionVector
	PRTVector []float64
	Jumps     [][]float64
}

// ChampionIndividual is a representation of the best individual currently
// available in the population.
type SOMAT3AChampionIndividual struct {
	X DecisionVector
}

// Population groups population individuals (agents) with metadata about the population.
type SOMAT3APopulation struct {
	// Population is a slice of population individuals.
	Population []SOMAT3APopulationIndividual
	// Champion represents the best individual of the population.
	Champion SOMAT3AChampionIndividual
	// Problem is the current benchmarking function this population is attempting to optimise.
	Problem string
	// ProblemFunction is the actual function to optimise.
	ProblemFunc func([]float64) float64
	// Dimen is the dimensionality of the problem being optimised.
	Dimen int
	// fes is the current number of function evaluations.
	fes int
	// Seed is the value used to (re)init population.
	Seed         uint64
	zeroToOneRNG distuv.Uniform

	// PRT is the perturbation parameter.
	PRT float64
	// 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
}

// GetIndividual returns a reference to individual at position n.
func (p *SOMAT3APopulation) GetIndividual(n uint) *PopulationIndividual {
	return &PopulationIndividual{}
}

// GetBestIdx returns the index of the best population individual.
func (p *SOMAT3APopulation) GetBestIdx() int {
	f := p.ProblemFunc

	bestIndividual := 0
	// the first one is the best one.
	bestVal := f(p.Population[0].CurX)

	for i, v := range p.Population {
		current := f(v.CurX)

		if current < bestVal {
			bestIndividual = i
		}
	}

	return bestIndividual
}

// GetWorstIdx returns the index of the worst population individual.
func (p *SOMAT3APopulation) GetWorstIdx() int {
	f := p.ProblemFunc

	worstIndividual := 0
	// the first one is the worst one.
	worstVal := f(p.Population[0].CurX)

	for i, v := range p.Population {
		current := f(v.CurX)

		if current > worstVal {
			worstIndividual = i
		}
	}

	return worstIndividual
}

// Init initialises all individuals to random values.
func (p *SOMAT3APopulation) Init() {
	low := cec2020.SearchRange.Min()
	high := cec2020.SearchRange.Max()
	// zeroToOneRNG := distuv.Uniform{
	p.zeroToOneRNG = distuv.Uniform{
		Min: 0,
		Max: 1,
		Src: rand.NewSource(p.Seed),
	}

	// gaLogger.Printf("population initialisation - popCount: %d, seed: %d\n",
	// len(p.Population), p.Seed,
	// )

	p.PRT = p.getPrt()

	// for i,v  := range p.Population {
	for i := range p.Population {
		v := SOMAT3APopulationIndividual{}

		v.CurX = make([]float64, p.Dimen)

		for j := 0; j < p.Dimen; j++ {
			// with current settings, this will always be `low + 0` since
			// `high - low = 0`.
			v.CurX[j] = low + p.zeroToOneRNG.Rand()*(high-low)
		}

		v.PRTVector = p.getPrtVector()

		v.Jumps = make([][]float64, p.Njumps)

		p.Population[i] = v
	}

	p.Champion = SOMAT3AChampionIndividual{X: p.Population[p.GetBestIdx()].CurX}
}

// Reinit reinitialises all individuals.
func (p *SOMAT3APopulation) Reinit() {
	p.Init()
}

// Clear sets all vectors to 0.
func (p *SOMAT3APopulation) Clear() {
	if p.Population != nil {
		for _, v := range p.Population {
			v.CurX = make([]float64, p.Dimen)
		}
	}
}

// Size returns the number of population individuals.
func (p *SOMAT3APopulation) Size() int {
	return len(p.Population)
}

// newPopulation returns a pointer to a new, uninitialised population.
func newSOMAT3APopulation(benchProblem string, np, k, m, n, nj, dimen int) *SOMAT3APopulation {
	p := &SOMAT3APopulation{
		Problem:     benchProblem,
		ProblemFunc: cec2020.Functions[benchProblem],
		K:           k,
		M:           m,
		N:           n,
		Njumps:      nj,
		Dimen:       dimen,
		Population:  make([]SOMAT3APopulationIndividual, np),
	}

	return p
}