package genetic

import (
	"log"

	"sort"

	neuralnetwork "../neuralnetwork"
)

type PopulationConfig struct {
	PopulationSize int
	SelectionSize  float64 // 0..1 persentage of success individuals to be used as parents for population
	CrossbreedPart float64 // 0..1 persentage of weights and biases to be exchanged beetween individuals while Crossbreed
}

type Population struct {
	populationConfig *PopulationConfig
	Networks         []*neuralnetwork.NeuralNetwork
	verifier         PopulationVerifier
	mutagen          Mutagen
}

func NewPopulation(verifier PopulationVerifier, mutagen Mutagen, populationConfig PopulationConfig, sizes []int) (p *Population) {
	if populationConfig.PopulationSize%2 != 0 {
		return nil
	}

	p = &Population{
		populationConfig: &populationConfig,
		Networks:         make([]*neuralnetwork.NeuralNetwork, populationConfig.PopulationSize),
		verifier:         verifier,
		mutagen:          mutagen,
	}

	for i := 0; i < populationConfig.PopulationSize; i++ {
		var err error
		p.Networks[i], err = neuralnetwork.NewNeuralNetwork(sizes, nil)
		if err != nil {
			log.Fatal("Could not initialize NeuralNetwork")
		}
	}

	return
}

func (p *Population) NaturalSelection(generationCount int) {
	for g := 0; g < generationCount; g++ {
		p.crossbreedPopulation(p.verifier.Verify(p))
	}
}

func (p *Population) crossbreedPopulation(results []*IndividalResult) {
	sort.Slice(results, func(i, j int) bool {
		return results[i].Result < results[j].Result
	})

	etalons := int(float64(p.populationConfig.PopulationSize) * p.populationConfig.SelectionSize)
	for i := 1; i < p.populationConfig.PopulationSize; i += 2 {
		firstParentBase := results[i%etalons].Index
		secondParentBase := results[(i-1)%etalons].Index
		firstParent := results[i].Index
		secondParent := results[i-1].Index
		p.Networks[firstParent] = p.Networks[firstParentBase].Copy()
		p.Networks[secondParent] = p.Networks[secondParentBase].Copy()
		crossbreed(p.Networks[firstParent], p.Networks[secondParent], p.populationConfig.CrossbreedPart)
		p.mutagen.Mutate(p.Networks[firstParent])
		p.mutagen.Mutate(p.Networks[secondParent])
	}
}

func crossbreed(firstParent, secondParent *neuralnetwork.NeuralNetwork, crossbreedPart float64) {
	for l := 1; l < firstParent.LayerCount; l++ {
		firstParentWeights := firstParent.Weights[l]
		secondParentWeights := secondParent.Weights[l]
		firstParentBiases := firstParent.Biases[l]
		secondParentBiases := secondParent.Biases[l]
		r, c := firstParentWeights.Dims()
		for i := 0; i < int(float64(r)*crossbreedPart); i++ {
			for j := 0; j < c; j++ {
				// Swap first half of weights
				w := firstParentWeights.At(i, j)
				firstParentWeights.Set(i, j, secondParentWeights.At(i, j))
				secondParentWeights.Set(i, j, w)
			}
			// Swap first half of biases
			b := firstParentBiases.At(i, 0)
			firstParentBiases.Set(i, 0, secondParentBiases.At(i, 0))
			secondParentBiases.Set(i, 0, b)
		}
	}
}