package net.bmahe.genetics4j.samples.neat;
   
   import java.time.Duration;
   import java.util.ArrayList;
   import java.util.List;
   import java.util.Map;
   import java.util.Objects;
   import java.util.Set;
   import java.util.random.RandomGenerator;
  
  import org.apache.logging.log4j.LogManager;
  import org.apache.logging.log4j.Logger;
  
  import net.bmahe.genetics4j.core.EASystem;
  import net.bmahe.genetics4j.core.EASystemFactory;
  import net.bmahe.genetics4j.core.Fitness;
  import net.bmahe.genetics4j.core.evolutionlisteners.EvolutionListeners;
  import net.bmahe.genetics4j.core.spec.EAConfiguration;
  import net.bmahe.genetics4j.core.spec.EAConfiguration.Builder;
  import net.bmahe.genetics4j.core.spec.EvolutionResult;
  import net.bmahe.genetics4j.core.spec.Optimization;
  import net.bmahe.genetics4j.core.spec.mutation.CreepMutation;
  import net.bmahe.genetics4j.core.spec.mutation.MultiMutations;
  import net.bmahe.genetics4j.core.spec.mutation.MutationPolicy;
  import net.bmahe.genetics4j.core.spec.mutation.RandomMutation;
  import net.bmahe.genetics4j.core.spec.selection.SelectionPolicy;
  import net.bmahe.genetics4j.core.spec.selection.Tournament;
  import net.bmahe.genetics4j.core.spec.statistics.distributions.NormalDistribution;
  import net.bmahe.genetics4j.core.termination.Terminations;
  import net.bmahe.genetics4j.neat.Activations;
  import net.bmahe.genetics4j.neat.Connection;
  import net.bmahe.genetics4j.neat.NeatEAExecutionContexts;
  import net.bmahe.genetics4j.neat.NeatUtils;
  import net.bmahe.genetics4j.neat.RecurrentNetwork;
  import net.bmahe.genetics4j.neat.chromosomes.NeatChromosome;
  import net.bmahe.genetics4j.neat.spec.NeatChromosomeSpec;
  import net.bmahe.genetics4j.neat.spec.combination.NeatCombination;
  import net.bmahe.genetics4j.neat.spec.mutation.AddConnection;
  import net.bmahe.genetics4j.neat.spec.mutation.AddNode;
  import net.bmahe.genetics4j.neat.spec.mutation.DeleteConnection;
  import net.bmahe.genetics4j.neat.spec.mutation.DeleteNode;
  import net.bmahe.genetics4j.neat.spec.mutation.NeatConnectionWeight;
  import net.bmahe.genetics4j.neat.spec.mutation.SwitchStateMutation;
  import net.bmahe.genetics4j.neat.spec.selection.NeatSelection;
  import net.bmahe.genetics4j.neat.util.GraphvizFormatter;
  
  public final class RecurrentParityExample {
  
  	private static final Logger logger = LogManager.getLogger(RecurrentParityExample.class);
  
  	public RecurrentParityExample() {
  	}
  
  	private static final RandomGenerator random = RandomGenerator.getDefault();
  
  	// tag::training_dataset[]
  	private static final List<List<Integer>> TRAINING_SEQUENCES = random.ints(100)
  			.mapToObj(i -> random.doubles(4).mapToInt(j -> j <= 0.5 ? 0 : 1).boxed().toList())
  			.toList();
  	// end::training_dataset[]
  
  	// tag::fitness_max_score[]
  	/**
  	 * 4 bits * range between odd-even (distance from -1 to 1) * number of sequences
  	 */
  	private static final float MAX_THEORETICAL_FITNESS = 4f * 2f * TRAINING_SEQUENCES.size();
  	private static final float FITNESS_TERMINATION_THRESHOLD = MAX_THEORETICAL_FITNESS * 0.99f;
  	// end::fitness_max_score[]
  
  	public static void main(final String[] args) {
  		final RecurrentParityExample recurrentParityExample = new RecurrentParityExample();
  
  		// tag::config_mutation_default_values[]
  		final double randomWeightsProb = 0.05;
  		final double switchMutationProb = 0.005;
  		final double addNodeProb = 0.015;
  		final double deleteNodeProb = 0.005;
  		final double addConnectionProb = 0.035;
  		final double deleteConnectionProb = 0.005;
  		// end::config_mutation_default_values[]
  
  		recurrentParityExample.run(
  				randomWeightsProb,
  				switchMutationProb,
  				addNodeProb,
  				deleteNodeProb,
  				addConnectionProb,
  				deleteConnectionProb);
  	}
  
  	public EvolutionResult<Float> run(final double randomWeightsProb, final double switchMutationProb,
  			final double addNodeProb, final double deleteNodeProb, final double addConnectionProb,
  			final double deleteConnectionProb) {
  
  		final List<MutationPolicy> mutations = buildMutations(
  				randomWeightsProb,
  				switchMutationProb,
  				addNodeProb,
  				deleteNodeProb,
 				addConnectionProb,
 				deleteConnectionProb);
 
 		final EAConfiguration<Float> configuration = buildParityConfiguration(mutations);
 		final var executionContext = NeatEAExecutionContexts.<Float>standard()
 				.populationSize(500)
 				.addEvolutionListeners(EvolutionListeners.ofLogTopN(logger, 3))
 				.build();
 
 		final EASystem<Float> eaSystem = EASystemFactory.from(configuration, executionContext);
 		final var result = eaSystem.evolve();
 
 		logger.info("Best fitness: {}", result.bestFitness());
 		describeBestChromosome(result.bestIndividual().genotype().getChromosome(0, NeatChromosome.class));
 
 		return result;
 	}
 
 	// tag::parity_configuration[]
 	private EAConfiguration<Float> buildParityConfiguration(final List<MutationPolicy> mutations) {
 		Objects.requireNonNull(mutations);
 
 		final Builder<Float> builder = new EAConfiguration.Builder<>();
 		builder.chromosomeSpecs(NeatChromosomeSpec.of(2, 1, -3, 3))
 				.parentSelectionPolicy(buildSelectionPolicy())
 				.combinationPolicy(NeatCombination.build())
 				.mutationPolicies(mutations)
 				.fitness(createParityFitness(false))
 				.optimization(Optimization.MAXIMIZE)
 				.termination(
 						Terminations.<Float>or(
 								Terminations.ofMaxTime(Duration.ofMinutes(5)),
 								Terminations.ofStableFitness(200),
 								Terminations.ofMaxGeneration(2000),
 								Terminations.ofFitnessAtLeast(FITNESS_TERMINATION_THRESHOLD)));
 		return builder.build();
 	}
 	// end::parity_configuration[]
 
 	// tag::config_selection[]
 	private SelectionPolicy buildSelectionPolicy() {
 		return NeatSelection.builder()
 				.minSpeciesSize(3)
 				.perSpeciesKeepRatio(0.20f)
 				.speciesSelection(Tournament.of(3))
 				.speciesPredicate(
 						(i1, i2) -> NeatUtils.compatibilityDistance(i1.genotype(), i2.genotype(), 0, 2, 2, 1f) < 3)
 				.build();
 	}
 	// end::config_selection[]
 
 	// tag::config_mutation[]
 	private List<MutationPolicy> buildMutations(final double randomWeightsProb, final double switchMutationProb,
 			final double addNodeProb, final double deleteNodeProb, final double addConnectionProb,
 			final double deleteConnectionProb) {
 		return List.of(
 				RandomMutation.of(randomWeightsProb),
 				MultiMutations.of(
 						CreepMutation.of(0.85f, NormalDistribution.of(0.0, 0.333)),
 						NeatConnectionWeight.builder().populationMutationProbability(0.85f).build()),
 				SwitchStateMutation.of(switchMutationProb),
 				AddNode.of(addNodeProb),
 				DeleteNode.of(deleteNodeProb),
 				AddConnection.builder()
 						.populationMutationProbability(addConnectionProb)
 						.allowOutputNodeAsSource(true)
 						.build(),
 				DeleteConnection.of(deleteConnectionProb));
 	}
 	// end::config_mutation[]
 
 	// tag::parity_fitness[]
 	private Fitness<Float> createParityFitness(final boolean verbose) {
 		return genotype -> {
 			final NeatChromosome chromosome = genotype.getChromosome(0, NeatChromosome.class);
 
 			if (chromosome.getConnections().isEmpty()) {
 				return 0f;
 			}
 
 			final RecurrentNetwork network = buildNetwork(chromosome);
 			final Integer outputNodeIndex = chromosome.getOutputNodeIds().getFirst();
 
 			float cumulativeError = 0F;
 			for (final List<Integer> sequence : TRAINING_SEQUENCES) {
 				network.resetState();
 
 				Map<Integer, Float> outputs = Map.of();
 				List<Integer> input = new ArrayList<>(4);
 				for (final int bit : sequence) {
 					input.add(bit);
 					final float expectedValue = expectedParity(input);
 
 					outputs = network.step(Map.of(0, encodeBit(bit), 1, 1.0f));
 
 					final float prediction = outputs.getOrDefault(outputNodeIndex, 0.0f);
 					final float error = Math.abs(expectedValue - prediction);
 					cumulativeError += error;
 
 					if (verbose) {
 						logger.info("Sequence {} -> expected {}, predicted {}", input, expectedValue, prediction);
 					}
 				}
 
 			}
 
 			return MAX_THEORETICAL_FITNESS - cumulativeError;
 		};
 	}
 	// end::parity_fitness[]
 
 	private RecurrentNetwork buildNetwork(final NeatChromosome chromosome) {
 		final Set<Integer> inputIndices = Set.copyOf(chromosome.getInputNodeIds());
 		final Set<Integer> outputIndices = Set.copyOf(chromosome.getOutputNodeIds());
 		final List<Connection> connections = chromosome.getConnections();
 
 		return new RecurrentNetwork(inputIndices, outputIndices, connections, Activations.tanhFloat, 12, 1e-4f, 0.0f);
 	}
 
 	private float encodeBit(final int bit) {
 		return bit == 0 ? -1.0f : 1.0f;
 	}
 
 	private float expectedParity(final List<Integer> bits) {
 		final long ones = bits.stream().filter(bit -> bit == 1).count();
 		return ones % 2 == 0 ? -1.0f : 1.0f;
 	}
 
 	private void describeBestChromosome(final NeatChromosome chromosome) {
 		final RecurrentNetwork network = buildNetwork(chromosome);
 		final Integer outputNodeIndex = chromosome.getOutputNodeIds().getFirst();
 
 		for (final List<Integer> sequence : TRAINING_SEQUENCES.subList(0, 10)) {
 			network.resetState();
 			Map<Integer, Float> outputs = Map.of();
 			for (final int bit : sequence) {
 				outputs = network.step(Map.of(0, encodeBit(bit), 1, 1.0f));
 			}
 
 			final float prediction = outputs.getOrDefault(outputNodeIndex, 0.0f);
 			final float expectedParity = expectedParity(sequence);
 			logger.info("Sequence {} -> Prediction: {} - Expected: {}", sequence, prediction, expectedParity);
 		}
 
 		final GraphvizFormatter graphvizFormatter = new GraphvizFormatter();
 		final String bestNetworkStr = graphvizFormatter.format(chromosome, Map.of(0, "Input", 1, "bias=1", 2, "Output"));
 		logger.info("{}", bestNetworkStr);
 
 	}
 }