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

package de

import (
	"git.dotya.ml/wanderer/math-optim/bench"
	"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 representats a single population individual.
type PopulationIndividual struct {
	CurX DecisionVector
	CurV DecisionVector
	CurC ConstraintVector
	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
	C ConstraintVector
	F FitnessVector
}

// Population groups population individuals (agents) with metadata about the population.
type Population struct {
	// Population is a slice of population individuals.
	Population []PopulationIndividual
	// Problem is the current benchmarking function this population is attempting to optimise.
	Problem string
	// Dimen is the dimensionality of the problem being optimised.
	Dimen int
	// Seed is the value used to (re)init population.
	Seed uint64
}

// GetIndividal 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 := bench.Functions[p.Problem]

	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 := bench.Functions[p.Problem]

	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
}

func (p *Population) SetX(n int, nuX DecisionVector) {}
func (p *Population) SetV(n int, nuV DecisionVector) {}

// Init initialises all individuals to random values.
func (p *Population) Init() {
	benchFuncParams := bench.FunctionParams[p.Problem]

	uniform := distuv.Uniform{
		Min: benchFuncParams.Min(),
		Max: benchFuncParams.Max(),
	}
	uniform.Src = rand.NewSource(p.Seed)

	// jDELogger.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
	}
}

// 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.CurC = make([]float64, p.Dimen)
			v.CurF = make([]float64, p.Dimen)
			v.CurV = make([]float64, p.Dimen)
			v.BestX = make([]float64, p.Dimen)
			v.BestC = make([]float64, p.Dimen)
			v.BestF = make([]float64, p.Dimen)
		}
	}
}

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

// newPopulation returns a pointer to a new, uninitialised population.
func newPopulation(benchProblem string, np, dimen int) *Population {
	p := &Population{}

	p.Problem = benchProblem

	p.Dimen = dimen

	// pre-alloc.
	p.Population = make([]PopulationIndividual, np)

	return p
}
