package genetic

import (
	"fmt"
	"log"
	"math/rand"

	"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
	etalonsCount     int
}

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,
		etalonsCount:     int(float64(populationConfig.PopulationSize) * populationConfig.SelectionSize),
	}

	if p.etalonsCount%2 != 0 {
		p.etalonsCount -= 1
	}

	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(fitnesses []*IndividalFitness) {
	sort.Slice(fitnesses, func(i, j int) bool {
		return fitnesses[i].Fitness > fitnesses[j].Fitness //Descent order best will be on top, worst in the bottom
	})

	//Collect etalons from upper part of neural network list and crossbreed/mutate them
	etalonNetworks := make([]*neuralnetwork.NeuralNetwork, p.etalonsCount)
	for i := 1; i < p.etalonsCount; i += 2 {
		firstParent := fitnesses[i-1].Index
		secondParent := fitnesses[i].Index
		fmt.Printf("Result i %v firstParent %v secondParent %v firstFitness %v secondFitness %v\n", i, firstParent, secondParent, fitnesses[i-1].Fitness, fitnesses[i].Fitness)

		etalonNetworks[i-1] = p.Networks[firstParent].Copy()
		etalonNetworks[i] = p.Networks[secondParent].Copy()

		crossbreed(p.Networks[firstParent], p.Networks[secondParent], p.populationConfig.CrossbreedPart)
		p.mutagen.Mutate(p.Networks[firstParent])
		p.mutagen.Mutate(p.Networks[secondParent])
	}

	//Rest of networks are based on collected etalons but crossbreed/mutate own way
	for i := p.etalonsCount + 1; i < p.populationConfig.PopulationSize; i += 2 {
		firstParent := fitnesses[i-1].Index
		secondParent := fitnesses[i].Index
		fmt.Printf("Result i %v firstParent %v secondParent %v firstFitness %v secondFitness %v firstEtalon %v secondEtalon %v\n",
			i, firstParent, secondParent,
			fitnesses[i-1].Fitness, fitnesses[i].Fitness,
			fitnesses[(i-1)%p.etalonsCount].Index, fitnesses[(i)%p.etalonsCount].Index)

		firstParentEtalon := etalonNetworks[(i-1)%p.etalonsCount]
		secondParenEtalon := etalonNetworks[(i)%p.etalonsCount]

		p.Networks[firstParent] = firstParentEtalon.Copy()
		p.Networks[secondParent] = secondParenEtalon.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()

		//Minimal part of matrix will be chosen for crossbreed
		rWindow := int(float64(r) * crossbreedPart)
		cWindow := int(float64(c) * crossbreedPart)

		r = int(rand.Uint32())%(r-rWindow) + rWindow
		c = int(rand.Uint32())%(c-cWindow) + cWindow
		for i := 0; i < r; i++ {
			for j := 0; j < c; j++ {
				// Swap weights
				w := firstParentWeights.At(i, j)
				firstParentWeights.Set(i, j, secondParentWeights.At(i, j))
				secondParentWeights.Set(i, j, w)
			}
			// Swap biases
			b := firstParentBiases.At(i, 0)
			firstParentBiases.Set(i, 0, secondParentBiases.At(i, 0))
			secondParentBiases.Set(i, 0, b)
		}
	}
}