// 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"
)

type (
	// DecisionVector is a []float64 abstraction representing the decision vector.
	DecisionVector []float64
	// FitnessVector is a []float64 abstraction representing the fitness vector.
	FitnessVector []float64
	// ConstraintVector is a []float64 abstraction representing the constraint vector.
	ConstraintVector []float64
)

// PopulationIndividual represents a single population individual.
type PopulationIndividual struct {
	CurX DecisionVector
	CurF FitnessVector

	BestX DecisionVector
	BestC ConstraintVector
	BestF FitnessVector
}

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

// Population groups population individuals (agents) with metadata about the population.
type Population struct {
	// Population is a slice of population individuals.
	Population []PopulationIndividual
	// Champion represents the best individual of the population.
	Champion ChampionIndividual
	// 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
	// Seed is the value used to (re)init population.
	Seed uint64

	// f is the differential weight (mutation/weighting factor) adapted over time.
	f []float64
	// cr is the crossover probability constant adapted over time.
	cr []float64
	// BestF is the best recorded value of the differential weight F.
	BestF float64
	// BestCR is the best recorded value of the differential weight CR.
	BestCR float64
	// CurF is the current value of F.
	CurF float64
	// CurCR is the current value of the differential weight CR.
	CurCR float64
}

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

// GetBestIdx returns the index of the best population individual.
func (p *Population) 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 *Population) 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 *Population) Init() {
	uniform := distuv.Uniform{
		Min: cec2020.SearchRange.Min(),
		Max: cec2020.SearchRange.Max(),
	}
	uniform.Src = rand.NewSource(p.Seed)

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

	for i, v := range p.Population {
		v.CurX = make([]float64, p.Dimen)

		for j := 0; j < p.Dimen; j++ {
			v.CurX[j] = uniform.Rand()
		}

		p.Population[i] = v
	}

	p.f = make([]float64, p.Size())
	p.cr = make([]float64, p.Size())

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

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

// ReinitN reinitialises the individual at position n.
func (p *Population) ReinitN(n uint) {}

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

func (p *Population) SelectDonors(currentIdx int) []PopulationIndividual {
	popCount := p.Size()

	idcs := make([]int, 0, popCount-1)

	// gather indices.
	for k := 0; k < popCount; k++ {
		if k != currentIdx {
			idcs = append(idcs, k)
		}
	}

	// randomly choose 3 of those idcs.
	selectedIdcs := make([]int, 0)

	selectedA := false
	selectedB := false
	selectedC := false

	for !selectedA {
		candidateA := rand.Intn(len(idcs)) % len(idcs)

		if candidateA != currentIdx {
			selectedIdcs = append(selectedIdcs, candidateA)
			selectedA = true
		}
	}

	for !selectedB {
		a := selectedIdcs[0]
		candidateB := rand.Intn(len(idcs)) % len(idcs)

		if candidateB != currentIdx && candidateB != a {
			selectedIdcs = append(selectedIdcs, candidateB)
			selectedB = true
		}
	}

	for !selectedC {
		a := selectedIdcs[0]
		b := selectedIdcs[1]
		candidateC := rand.Intn(len(idcs)) % len(idcs)

		if candidateC != currentIdx && candidateC != a && candidateC != b {
			selectedIdcs = append(selectedIdcs, candidateC)
			selectedC = true
		}
	}

	// selected contains the selected population individuals.
	selected := make([]PopulationIndividual, 0)

	// select individuals for donation.
	for _, idx := range selectedIdcs {
		for k := 0; k < popCount; k++ {
			if k == idx {
				selected = append(selected, p.Population[idx])
			}
		}
	}

	return selected
}

// MeanVelocity computes the mean current velocity of all individuals in the population.
func (p *Population) MeanVelocity() float64 { return 0.0 }

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

// newPopulation returns a pointer to a new, uninitialised population.
func newPopulation(benchProblem string, np, dimen int) *Population {
	p := &Population{
		Problem:     benchProblem,
		ProblemFunc: cec2020.Functions[benchProblem],
		Dimen:       dimen,
		Population:  make([]PopulationIndividual, np),
	}

	return p
}