#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
X**2 optimiziation
"""

# %%
import sys
import string
import random
from typing import TypedDict


class Individual(TypedDict):
    genome: str
    fitness: int


Population = list[Individual]


def initialize_individual(genome: str, fitness: int) -> Individual:
    return {"genome": genome, "fitness": fitness}


def initialize_pop(bitlen: int, pop_size: int) -> Population:
    population: Population = []
    for _ in range(pop_size):
        genome = "".join(random.choice("01") for _ in range(bitlen))
        individual = initialize_individual(genome=genome, fitness=0)
        population.append(individual)
    return population


def recombine_pair(parent1: Individual, parent2: Individual) -> Population:
    place = random.choice(range(len(parent1["genome"])))
    child1_genome = parent1["genome"][:place] + parent2["genome"][place:]
    child2_genome = parent2["genome"][:place] + parent1["genome"][place:]
    child1 = initialize_individual(genome=child1_genome, fitness=0)
    child2 = initialize_individual(genome=child2_genome, fitness=0)
    return [child1, child2]


def recombine_group(parents: Population, recombine_rate: float) -> Population:
    children: Population = []
    for ipair in range(0, len(parents) - 1, 2):
        if random.random() < recombine_rate:
            child1, child2 = recombine_pair(
                parent1=parents[ipair], parent2=parents[ipair + 1]
            )
        else:
            child1, child2 = parents[ipair], parents[ipair + 1]
        children.extend([child1, child2])
    return children


def mutate_individual(parent: Individual, mutate_rate: float) -> Individual:
    new_genome = "".join(
        [
            (random.choice("01") if random.random() < mutate_rate else ch)
            for ch in parent["genome"]
        ]
    )
    mutant = initialize_individual(genome=new_genome, fitness=0)
    return mutant


def mutate_group(children: Population, mutate_rate: float) -> Population:
    mutants: Population = []
    for child in children:
        mutants.append(mutate_individual(parent=child, mutate_rate=mutate_rate))
    return mutants


def evaluate_individual(bitlen: int, individual: Individual) -> None:
    individual["fitness"] = int(individual["genome"], 2) ** 2


def evaluate_group(bitlen: int, individuals: Population) -> None:
    for ind_index in range(len(individuals)):
        evaluate_individual(bitlen, individuals[ind_index])


def rank_group(individuals: Population) -> None:
    individuals.sort(key=lambda ind: ind["fitness"], reverse=True)


def parent_select(individuals: Population, number: int) -> Population:
    parents: Population = []
    fitnesses = [i["fitness"] for i in individuals]
    parents = random.choices(individuals, fitnesses, k=number)
    return parents


def survivor_select(individuals: Population, pop_size: int) -> Population:
    return individuals[:pop_size]


def evolve(bitlen: int, pop_size: int) -> Population:
    population = initialize_pop(bitlen=bitlen, pop_size=pop_size)
    evaluate_group(bitlen=bitlen, individuals=population)
    rank_group(individuals=population)
    best_fitness = population[0]["fitness"]
    perfect_fitness = ((2**bitlen) - 1) ** 2
    counter = 0
    while best_fitness < perfect_fitness:
        print(population)
        counter += 1
        parents = parent_select(individuals=population, number=pop_size)
        children = recombine_group(parents=parents, recombine_rate=0.8)
        mutate_rate = (1 - (best_fitness / perfect_fitness)) / 5
        mutants = mutate_group(children=children, mutate_rate=mutate_rate)
        evaluate_group(bitlen=bitlen, individuals=mutants)
        everyone = population + mutants
        rank_group(individuals=everyone)
        population = survivor_select(individuals=everyone, pop_size=pop_size)
        if best_fitness != population[0]["fitness"]:
            best_fitness = population[0]["fitness"]
            print("Iteration number", counter, "with best individual", population[0])
    return population


if __name__ == "__main__":
    import json

    if len(sys.argv) == 3:
        with open(file=sys.argv[1]) as finput:
            obj_name = finput.readlines()
            OBJECTIVE = int(obj_name[0].strip())
            POP_SIZE = int(obj_name[1])
        with open(file=sys.argv[2], mode="w") as foutput:
            population = evolve(bitlen=OBJECTIVE, pop_size=POP_SIZE)
            foutput.write(json.dumps(population) + "\n")
    else:
        BITLEN = int(input("What size bitstring you like to evolve?\n"))
        POP_SIZE = int(input("How many individuals would you like to evolve?\n"))
        population = evolve(bitlen=BITLEN, pop_size=POP_SIZE)