/*
 * MIT License
 *
 * Copyright (c) 2019 Alexey Edelev <semlanik@gmail.com>
 *
 * This file is part of NeuralNetwork project https://git.semlanik.org/semlanik/NeuralNetwork
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this
 * software and associated documentation files (the "Software"), to deal in the Software
 * without restriction, including without limitation the rights to use, copy, modify,
 * merge, publish, distribute, sublicense, and/or sell copies of the Software, and
 * to permit persons to whom the Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies
 * or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
 * PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
 * FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

package genetic

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

	"sort"

	neuralnetwork "git.semlanik.org/semlanik/NeuralNetwork/neuralnetwork"
)

// Genetic package implements basic proncipals of genetic and natural selection mechanisms for
// neuralnetwork.NeuralNetwork

// PopulationConfig is structure that is used to specify population parameters
// PopulationSize - size of population
// SelectionSize - percentage of best individuals used for next generation
// CrossbreedPart - percentage of weigts and biases used for crossbreed
type PopulationConfig struct {
	PopulationSize int
	SelectionSize  float64 // 0..1 percentage of success individuals to be used as parents for population
	CrossbreedPart float64 // 0..1 percentage of weights and biases to be exchanged beetween individuals while Crossbreed
}

// Population is main class for genetic and natural selection of neuralnetwork.NeuralNetwork
type Population struct {
	populationConfig *PopulationConfig
	Networks         []*neuralnetwork.NeuralNetwork
	verifier         PopulationVerifier
	mutagen          Mutagen
	etalonsCount     int
	bestFitness      *IndividalFitness
	bestNetwork      *neuralnetwork.NeuralNetwork
}

// NewPopulation is constructor of new Population with specified PopulationVerifier, Mutagen and PopulationConfig.
// sizes parameter also specified neuralnetwork.NeuralNetwork layers configuration
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),
		bestFitness:      nil,
		bestNetwork:      nil,
	}

	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
}

// NaturalSelection invokes natural selection process for specified number of generation
func (p *Population) NaturalSelection(generationCount int) {
	for g := 0; g < generationCount; g++ {
		p.crossbreedPopulation(p.verifier.Verify(p))
	}
}

// GetBestNetwork method returns best network in population according to it fitness
func (p *Population) GetBestNetwork() *neuralnetwork.NeuralNetwork {
	return p.bestNetwork
}

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
	})

	//Save best fitness individal
	if p.bestFitness == nil || p.bestFitness.Fitness < fitnesses[0].Fitness {
		p.bestFitness = fitnesses[0]
		p.bestNetwork = p.Networks[fitnesses[0].Index].Copy()
		fmt.Printf("New best Fitness %v \n", fitnesses[0].Fitness)
	}

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