import statistics 

# Fitness function that just optimizes for
# total percentage of assigned Highlighters
def fitness_1(buckets, assignment):
    '''
        buckets: []Bucket
        assignments: dict {Highlighter: Bucket_Name}

    '''

    assigned_buckets = list(assignment.values())

    fill_ratios = []
    for bucket in buckets:
        num_filled = sum(1 for assigned_bucket in assigned_buckets if assigned_bucket == bucket.name)
        fill_ratios.append(min(num_filled / bucket.fill_goal, 1.0))

    return statistics.mean(fill_ratios)

# Same as fitness_1 except it penalizes for overfilling buckets
def fitness_2(buckets, assignment):
    '''
        buckets: []Bucket
        assignments: dict {Highlighter: Bucket_Name}
    '''

    assigned_buckets = list(assignment.values())

    fill_ratios = []
    for bucket in buckets:
        num_filled = sum(1 for assigned_bucket in assigned_buckets if assigned_bucket == bucket.name)
        fill_ratio = num_filled / bucket.fill_goal
        if fill_ratio > 1.0:
            # Apply penalty for fill ratios greater than 1
            penalty = 1.0 - (fill_ratio - 1.0)
            fill_ratios.append(penalty)
        else:
            fill_ratios.append(fill_ratio)

    return statistics.mean(fill_ratios)



def fitness_3(buckets, assignments):
    """
    Calculate the fitness score based on the total number of assigned highlighters divided by the sum of bucket fill goals.
    
    :param buckets: List of Bucket objects
    :param assignments: Dictionary mapping Highlighter objects to bucket names
    :return: Fitness score
    """
    total_assigned_highlighters = len(assignments)
    total_fill_goals = sum(bucket.fill_goal for bucket in buckets)
    
    # Avoid division by zero
    if total_fill_goals == 0:
        return 0.0
    
    # Calculate fitness score
    fitness_score = total_assigned_highlighters / total_fill_goals
    
    return fitness_score
    

import random
def generate_random_assignment(buckets, highlighters):
    '''
        Each Highlighter can only be assigned to 1 Bucket.
    '''
    assignment = {}
    for h in highlighters:
        # Randomly choose one of the highlighter's assignable buckets
        chosen_bucket = random.choice(h.assignable_buckets)
        assignment[h] = chosen_bucket
    return assignment

assignment = generate_random_assignment(buckets, highlighters)

import random

def generate_weighted_assignment(buckets, highlighters):
    '''
        Generate assignments based on Bucket needs with some degree of randomness.
        Each Highlighter can only be assigned to one of its assignable buckets.
    '''
    assignment = {}
    
    # Dictionary to keep track of total filled per bucket
    filled_per_bucket = {bucket.name: 0 for bucket in buckets}
    
    # Calculate total filled per bucket
    for h in highlighters:
        if h in assignment:
            filled_per_bucket[assignment[h]] += 1

    random.shuffle(highlighters)
    for h in highlighters:
        # Filter buckets to only include those that are assignable to the highlighter
        assignable_buckets = [bucket for bucket in buckets if bucket.name in h.assignable_buckets]

        # If there are no assignable buckets, skip assignment for this highlighter
        if not assignable_buckets:
            continue
        
        # Calculate remaining needs ratios for each bucket
        remaining_ratios = [(bucket.fill_goal - filled_per_bucket[bucket.name]) / bucket.fill_goal for bucket in assignable_buckets]
        
        # Calculate weights based on remaining needs ratios
        weights = remaining_ratios
        
        # Calculate cumulative weights
        cumulative_weights = [sum(weights[:i+1]) for i in range(len(weights))]
        
        # Choose a random number within the total weight
        rand_num = random.uniform(0, sum(weights))
        
        # Find the corresponding bucket based on the random number
        chosen_bucket = None
        for i, cum_weight in enumerate(cumulative_weights):
            if rand_num <= cum_weight:
                chosen_bucket = assignable_buckets[i]
                break
        if chosen_bucket:
            assignment[h] = chosen_bucket.name
        
            # Update total filled per bucket
            filled_per_bucket[chosen_bucket.name] += 1
        else:
            assignment[h] = ''
    return assignment

assignment = generate_weighted_assignment(buckets, highlighters)

class GeneticAlgorithmRunner:
    def __init__(self, buckets, highlighters, mutation_rate=0.01, crossover_rate=0.1, cull_rate=0.1, max_pop=1000, max_epochs=50, fitness_function = fitness_1, generator_algorithm=generate_random_assignment):
        self.buckets = buckets
        self.highlighters = highlighters
        self.assignments = []
        self.cull_rate = cull_rate
        self.mutation_rate = mutation_rate
        self.crossover_rate = crossover_rate
        self.max_epochs = max_epochs
        self.max_population = max_pop
        self.fitness_function = fitness_function
        self.generator_algorithm = generator_algorithm

        self.population = []
                # Generate N random assignments
        for i in range(self.max_population):
            self.population.append(generator_algorithm(self.buckets, self.highlighters))

            # Select some % of fitnesses based on weights
            # selected_indices = random.choices(range(len(fitnesses)), weights=weights, k=int(self.cull_rate * len(fitnesses)))

    # Selects the top candidates 
    def cull(self, fitnesses, population):
        # Determine the number of individuals to select
        num_to_select = int(len(fitnesses) * self.cull_rate)
        
        # Create a list of indices and fitnesses
        indexed_fitnesses = list(enumerate(fitnesses))
        
        # Sort by fitness in descending order
        indexed_fitnesses.sort(key=lambda x: x[1], reverse=True)
        
        # Extract the top individuals based on sorted fitness scores
        selected_indices = [idx for idx, _ in indexed_fitnesses[:num_to_select]]
        
        # Select the fitnesses and population members by these indices
        selected_fitnesses = [fitnesses[i] for i in selected_indices]
        selected_population = [population[i] for i in selected_indices]
        
        return selected_fitnesses, selected_population

    def increase_population(self):
        """
        Increase the current population by breeding - apply mutation and crossover.
        """
        # Calculate the number of individuals needed to reach the maximum population size
        num_individuals_needed = self.max_population - len(self.population)

        original_population = self.population
        
        # Generate new individuals until the population reaches the maximum size
        while len(self.population) < self.max_population:
            # Select two parents randomly from the current population
            parent1, parent2 = random.sample(original_population, 2)
            
            offspring = self.breed(parent1, parent2)
            
            # Add the offspring to the population
            self.population.append(offspring)

    # Mutate a given individual's genes
    def mutate(self, assignment):
        for highlighter, bucket in assignment.items():
            if random.random() < self.mutation_rate:
                assignment[highlighter] = random.choice(highlighter.assignable_buckets)
        return assignment

    def crossover(self, assignment1, assignment2):
        """
        Perform crossover between two parents to produce a new individual.
        
        :param assignment1: Dictionary representing the assignment of highlighters to buckets for parent 1
        :param assignment2: Dictionary representing the assignment of highlighters to buckets for parent 2
        :return: Dictionary representing the assignment of highlighters to buckets for the offspring
        """
        # Initialize the offspring assignment
        offspring_assignment = {}

        # Randomly select genes (highlighter-bucket pairs) from each parent
        for highlighter in assignment1:
            if random.random() < 0.5:  # 50% chance to inherit from parent 1
                offspring_assignment[highlighter] = assignment1[highlighter]
            else:  # 50% chance to inherit from parent 2
                offspring_assignment[highlighter] = assignment2[highlighter]

        return offspring_assignment
    
    def breed(self, assignment1, assignment2):
        child_assignment = {}

        # Randomly make child a copy of either parent1 or parent2
        if random.random() < 0.5:
            child_assignment = assignment1
        else:
            child_assignment = assignment2

        # Probability of getting both parent's genes
        if random.random() < self.crossover_rate:
            child_assignment = self.crossover(assignment1, assignment2)

        # Mutate the genes
        child_assignment = self.mutate(child_assignment)        

        return child_assignment

    def get_best(self):
        """
        Return the highest score of the individuals in the current population.
        """
        best_score = float('-inf')  # Initialize with negative infinity to find the maximum score
        best_assignment = None
        
        # Iterate over each individual in the population
        for assignment in self.population:
            # Calculate the fitness score for the individual using the provided fitness function
            score = self.fitness_function(self.buckets, assignment)
            
            # Update the best score if the current score is higher
            if score > best_score:
                best_score = score
                best_assignment = assignment
        
        return best_assignment, best_score
        
    def run(self):
        for i in range(self.max_epochs):
            # Evaluate their fitnesses
            scores = []
            for assignment in self.population:
                scores.append(self.fitness_function(self.buckets, assignment))
    
            # Cull the herd (survival of the fittest!)
            best_scores, best_assignments = self.cull(scores, self.population)
            self.population = best_assignments

            # Breed to restore the population 
            self.increase_population()

            print(f"Epoch: {i}, Best Score: {max(best_scores)}")

            if max(best_scores) >= .9999:
                break
        

assignment, score = a.get_best()