package genetic

import (
	"log"

	"sort"

	neuralnetwork "../neuralnetwork"
)

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

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

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

	for i := 0; i < 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
	})

	for i := 1; i < p.populationSize; i += 2 {
		firstParent := results[i].index
		secondParent := results[i-1].index
		crossbreed(p.Networks[firstParent], p.Networks[secondParent])
		p.mutagen.Mutate(p.Networks[firstParent])
		p.mutagen.Mutate(p.Networks[secondParent])
	}
}

func crossbreed(firstParent, secondParent *neuralnetwork.NeuralNetwork) {
	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 < r/2; 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)
		}
	}
}