population.py

from copy import deepcopy
import os
import random
import time

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

import util
from individual import Individual


class Population:
    def __init__(self, population_size, tournament_size, mutation_rate=0.01):
        """Initialize a population of individuals for the genetic algorithm.

        Args:
            population_size: Number of individuals in the population.
            tournament_size: Number of individuals drawn for tournament selection.
            mutation_rate: Probability of mutation (default 0.01).
        """
        self.population_size = population_size
        self.tournament_size = tournament_size
        self.mutation_rate = mutation_rate
        self.cities = {}
        self.individuals = []
        # Lazily built caches; invalidated whenever the cities change. They
        # depend only on city coordinates, so they are computed once and reused
        # across every generation (see _distance_matrix / _nearest_order).
        self._distance_matrix = None
        self._nearest_order = None

    def copy(self):
        """Create a deep copy of the population's cities and regenerate individuals.

        Returns:
            Population: A new instance with the same cities but fresh individuals.
        """
        clone = Population(self.population_size, self.tournament_size, mutation_rate=self.mutation_rate)
        clone.cities = deepcopy(self.cities)
        clone.individuals = []
        clone.generate_individuals()
        return clone

    def add_city(self, name, coord):
        """Add a city to the problem.

        Args:
            name: Unique identifier of the city.
            coord: Tuple or list [x, y] of coordinates.
        """
        self.cities[name] = coord
        self._invalidate_distance_cache()

    def _invalidate_distance_cache(self):
        """Drop the cached distance data after the set of cities changes."""
        self._distance_matrix = None
        self._nearest_order = None

    def distance_matrix(self):
        """Return the full city-to-city distance matrix, building it once.

        The matrix depends only on city coordinates, so it is cached and reused
        across generations instead of being recomputed on every crossover.

        Returns:
            dict: {city1: {city2: distance}} for every pair of cities.
        """
        if self._distance_matrix is None:
            cities = list(self.cities.keys())
            self._distance_matrix = {
                v1: {v2: self.distance_between(v1, v2) for v2 in cities}
                for v1 in cities
            }
        return self._distance_matrix

    def nearest_order(self):
        """Return, for each city, the other cities sorted by increasing distance.

        Cached alongside the distance matrix (same coordinate dependency).

        Returns:
            dict: {city: [other cities sorted by distance]}.
        """
        if self._nearest_order is None:
            dm = self.distance_matrix()
            cities = list(self.cities.keys())
            self._nearest_order = {
                c: [d for d in sorted(cities, key=lambda x: dm[c][x]) if d != c]
                for c in cities
            }
        return self._nearest_order

    def generate_individuals(self):
        """Generate an initial population of individuals with unique random routes.

        Fixes the first city as a common starting point, then generates random
        permutations of the remaining cities. Stops after population_size * 10
        attempts to avoid infinite loops.
        """
        city_names = list(self.cities.keys())
        attempts = 0
        while len(self.individuals) < self.population_size and attempts < self.population_size * 10:
            route = city_names.copy()
            first = route.pop(0)
            random.shuffle(route)
            route = [first] + route
            individual = Individual(self, route)
            if not any(individual.equals(other) for other in self.individuals):
                self.individuals.append(individual)
            attempts += 1

    def distance_between(self, city1, city2):
        """Compute the Euclidean distance between two cities.

        Args:
            city1: Name of the first city.
            city2: Name of the second city.

        Returns:
            float: Euclidean distance between the two cities.
        """
        x1, y1 = self.cities[city1]
        x2, y2 = self.cities[city2]
        return ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5

    def generate_random_cities(self, size):
        """Generate cities with random coordinates in [0, 1]^2.

        Args:
            size: Number of cities to generate.
        """
        self.cities = {}
        self.individuals = []
        city_names = util.gen_names(size)
        for name in city_names:
            coord = [np.random.rand(), np.random.rand()]
            self.add_city(name, coord)

    def select_tournament(self):
        """Select an individual by tournament.

        Draws tournament_size individuals at random and returns the fittest.

        Returns:
            Individual: The best individual among those drawn.
        """
        selected = random.sample(self.individuals, self.tournament_size)
        selected.sort(key=lambda individual: individual.fitness, reverse=True)
        return selected[0]

    def pmx(self, parent1, parent2, cut_points=None):
        """Partially Mapped Crossover (PMX).

        Swaps a segment between parents, then completes each child via positional
        mappings to keep every city present without duplication.

        Args:
            parent1, parent2: Routes (lists of cities) of the parents.
            cut_points: Tuple (c1, c2) of cut points, or None for random.

        Returns:
            tuple: Two child routes (child1, child2).
        """
        n = len(parent1)
        if cut_points is None:
            c1, c2 = sorted(random.sample(range(n), 2))
        else:
            c1, c2 = cut_points

        child1, child2 = [None] * n, [None] * n

        pos1 = {gene: i for i, gene in enumerate(parent1)}
        pos2 = {gene: i for i, gene in enumerate(parent2)}

        child1[c1:c2] = parent1[c1:c2]
        child2[c1:c2] = parent2[c1:c2]

        for i in range(c1, c2):
            gene = parent2[i]
            if gene not in child1:
                pos = i
                while child1[pos] is not None:
                    mapped_gene = parent1[pos]
                    pos = pos2[mapped_gene]
                child1[pos] = gene

        for i in range(c1, c2):
            gene = parent1[i]
            if gene not in child2:
                pos = i
                while child2[pos] is not None:
                    mapped_gene = parent2[pos]
                    pos = pos1[mapped_gene]
                child2[pos] = gene

        for i in range(n):
            if child1[i] is None:
                child1[i] = parent2[i]
            if child2[i] is None:
                child2[i] = parent1[i]

        return child1, child2

    def cx(self, parent1, parent2):
        """Cycle Crossover (CX).

        Builds children by following the mapping cycles between parents,
        alternating the source for each new cycle to preserve absolute positions.

        Args:
            parent1, parent2: Routes (lists of cities) of the parents.

        Returns:
            tuple: Two child routes (child1, child2).
        """
        n = len(parent1)
        child1, child2 = [None] * n, [None] * n
        visited = [False] * n
        pos1 = {gene: i for i, gene in enumerate(parent1)}

        start = 0
        use_parent1 = True
        while not all(visited):
            idx = start
            if use_parent1:
                while not visited[idx]:
                    child1[idx] = parent1[idx]
                    child2[idx] = parent2[idx]
                    visited[idx] = True
                    idx = pos1[parent2[idx]]
            else:
                while not visited[idx]:
                    child1[idx] = parent2[idx]
                    child2[idx] = parent1[idx]
                    visited[idx] = True
                    idx = pos1[parent2[idx]]

            if not all(visited):
                start = visited.index(False)
                use_parent1 = not use_parent1

        return child1, child2

    def erx(self, parent1, parent2):
        """Edge Recombination Crossover (ERX).

        Builds a map of each city's neighbors across both parents, then generates
        children by repeatedly choosing the city with the fewest remaining
        neighbors (edge-preservation heuristic).

        Args:
            parent1, parent2: Routes (lists of cities) of the parents.

        Returns:
            tuple: Two child routes (child1, child2).
        """
        def build_edge_map(p1, p2):
            edges = {c: set() for c in p1}
            for p in [p1, p2]:
                for i in range(len(p)):
                    left, right = p[i - 1], p[(i + 1) % len(p)]
                    edges[p[i]].update([left, right])
            return edges

        def make_child(edges):
            child = []
            current = random.choice(list(edges.keys()))
            while len(child) < len(edges):
                child.append(current)
                for e in edges.values():
                    e.discard(current)
                if edges[current]:
                    current = min(edges[current], key=lambda x: len(edges[x]))
                else:
                    remaining = [c for c in edges if c not in child]
                    if remaining:
                        current = random.choice(remaining)
            return child

        edges = build_edge_map(parent1, parent2)
        edges_copy = {k: set(v) for k, v in edges.items()}
        return make_child(edges), make_child(edges_copy)

    def hx(self, parent1, parent2, distance_matrix):
        """Heuristic Crossover (HX).

        Builds children by choosing, at each step, the successor (in either
        parent) closest to the current city. If no successor is available,
        picks a random remaining city.

        Args:
            parent1, parent2: Routes (lists of cities) of the parents.
            distance_matrix: Precomputed dict {city1: {city2: distance}}.

        Returns:
            tuple: Two child routes (child1, child2).
        """
        n = len(parent1)

        pos1 = {city: i for i, city in enumerate(parent1)}
        pos2 = {city: i for i, city in enumerate(parent2)}

        def next_in_parent(p, pos_map, city):
            i = pos_map[city]
            return p[(i + 1) % n]

        def choose_next(current, used):
            c1 = next_in_parent(parent1, pos1, current)
            c2 = next_in_parent(parent2, pos2, current)
            candidates = []
            if c1 not in used:
                candidates.append(c1)
            if c2 not in used and c2 != c1:
                candidates.append(c2)
            if candidates:
                return min(candidates, key=lambda c: distance_matrix[current][c])
            remaining = [c for c in parent1 if c not in used]
            return random.choice(remaining) if remaining else None

        def make_child():
            child = []
            used = set()
            current = random.choice(parent1)
            while len(child) < n:
                child.append(current)
                used.add(current)
                nxt = choose_next(current, used)
                if nxt is None:
                    break
                current = nxt
            return child

        return make_child(), make_child()

    def hx_extended(self, parent1, parent2, distance_matrix, nearest_order):
        """Heuristic Crossover Extended - an improved version of HX.

        Considers the 4 neighbors (predecessor and successor in each parent) of
        the current city and picks the closest unvisited one. On failure, falls
        back to a list of cities pre-sorted by increasing distance.

        Args:
            parent1, parent2: Routes (lists of cities) of the parents.
            distance_matrix: Precomputed dict {city1: {city2: distance}}.
            nearest_order: Precomputed dict {city: [cities sorted by distance]}.

        Returns:
            tuple: Two child routes (child1, child2).
        """
        n = len(parent1)

        pos1 = {city: i for i, city in enumerate(parent1)}
        pos2 = {city: i for i, city in enumerate(parent2)}
        neighbor4 = {}
        for city in parent1:
            i1 = pos1[city]
            i2 = pos2[city]
            neighbor4[city] = (
                parent1[(i1 - 1) % n], parent1[(i1 + 1) % n],
                parent2[(i2 - 1) % n], parent2[(i2 + 1) % n],
            )

        def fallback_nearest(current, used):
            for cand in nearest_order[current]:
                if cand not in used:
                    return cand
            return None

        def make_child():
            child = []
            used = set()
            current = random.choice(parent1)
            while len(child) < n:
                child.append(current)
                used.add(current)

                cand = []
                seen = set()
                for v in neighbor4[current]:
                    if v not in used and v not in seen:
                        cand.append(v)
                        seen.add(v)

                if cand:
                    current = min(cand, key=lambda c: distance_matrix[current][c])
                else:
                    nxt = fallback_nearest(current, used)
                    if nxt is None:
                        break
                    current = nxt
            return child

        return make_child(), make_child()

    def ox(self, parent1, parent2, cut_points=None):
        """Order Crossover (OX).

        Swaps a segment between parents, then fills each child with the missing
        cities while preserving the relative order of the opposite parent
        (with wrap-around).

        Args:
            parent1, parent2: Routes (lists of cities) of the parents.
            cut_points: Tuple (c1, c2) of cut points, or None for random.

        Returns:
            tuple: Two child routes (child1, child2).
        """
        n = len(parent1)
        if cut_points is None:
            cut1, cut2 = sorted(random.sample(range(n), 2))
        else:
            cut1, cut2 = cut_points

        child1 = [None] * n
        child2 = [None] * n

        child1[cut1:cut2 + 1] = parent2[cut1:cut2 + 1]
        child2[cut1:cut2 + 1] = parent1[cut1:cut2 + 1]

        def fill_order(child, donor, c1, c2):
            seq = []
            i = (c2 + 1) % n
            while len(seq) < n:
                seq.append(donor[i])
                i = (i + 1) % n
            seq = [g for g in seq if g not in child]
            pos = (c2 + 1) % n
            for g in seq:
                while c1 <= pos <= c2:
                    pos = (pos + 1) % n
                child[pos] = g
                pos = (pos + 1) % n
            return child

        child1 = fill_order(child1, parent1, cut1, cut2)
        child2 = fill_order(child2, parent2, cut1, cut2)

        return child1, child2

    def crossover(self, method='ox', mutation='2-opt'):
        """Run a full reproduction step: selection, crossover, mutation.

        Selects two distinct parents by tournament, applies the chosen crossover
        operator, then mutates the children with probability mutation_rate. Adds
        the unique children to the population and removes the worst if the size
        is exceeded.

        Args:
            method: Crossover type ('pmx', 'cx', 'erx', 'ox', 'hx', 'hx extended').
            mutation: Mutation type ('2-opt', '2-opt alt', '2-opt extended', '3-opt').

        Returns:
            tuple: The two generated children (child1, child2).
        """
        parent1 = self.select_tournament()
        parent2 = self.select_tournament()
        while parent1.equals(parent2):
            parent2 = self.select_tournament()

        match method:
            case 'pmx':
                child1, child2 = self.pmx(parent1.route, parent2.route)
            case 'cx':
                child1, child2 = self.cx(parent1.route, parent2.route)
            case 'erx':
                child1, child2 = self.erx(parent1.route, parent2.route)
            case 'ox':
                child1, child2 = self.ox(parent1.route, parent2.route)
            case 'hx':
                child1, child2 = self.hx(parent1.route, parent2.route, self.distance_matrix())
            case 'hx extended':
                child1, child2 = self.hx_extended(parent1.route, parent2.route,
                                                  self.distance_matrix(), self.nearest_order())

        child1 = Individual(self, child1)
        child2 = Individual(self, child2)

        if np.random.rand() < self.mutation_rate:
            match mutation:
                case '2-opt':
                    child1.mutation_2opt()
                case '2-opt alt':
                    child1.mutation_2opt_alternative()
                case '2-opt extended':
                    child1.mutation_2opt_extended()
                case '3-opt':
                    child1.mutation_3opt()
        if np.random.rand() < self.mutation_rate:
            match mutation:
                case '2-opt':
                    child2.mutation_2opt()
                case '2-opt alt':
                    child2.mutation_2opt_alternative()
                case '2-opt extended':
                    child2.mutation_2opt_extended()
                case '3-opt':
                    child2.mutation_3opt()

        if not any(child1.equals(other) for other in self.individuals):
            self.individuals.append(child1)

        if not any(child2.equals(other) for other in self.individuals):
            self.individuals.append(child2)

        # Trim back to the target size by dropping the least-fit individuals
        # (steady-state, elitist replacement). Each pass locates the worst by
        # index in a single scan, avoiding the redundant index() lookup and the
        # wrong-individual risk when fitnesses tie.
        while len(self.individuals) > self.population_size:
            worst_idx = min(
                range(len(self.individuals)),
                key=lambda i: self.individuals[i].fitness,
            )
            self.individuals.pop(worst_idx)

        return child1, child2

    def generate_circle_cities(self, size):
        """Generate cities laid out on a circle (for tests / benchmarks).

        Places the cities uniformly on a circle of radius 0.4 centered at
        (0.5, 0.5). Useful for testing because the optimal solution is known.

        Args:
            size: Number of cities to generate.
        """
        self.cities = {}
        self.individuals = []
        city_names = util.gen_names(size)
        angle_step = 2 * np.pi / size
        for i, name in enumerate(city_names):
            angle = i * angle_step
            coord = [0.5 + 0.4 * np.cos(angle), 0.5 + 0.4 * np.sin(angle)]
            self.add_city(name, coord)

    def plot_result(self, best_individual, title="Best route found", display=True):
        """Show a 2D map of the best route found.

        Draws the cities as red points, the best route in blue, and marks the
        start city in green.

        Args:
            best_individual: Individual instance representing the best route.
            title: Title of the plot.
            display: If True, show the plot immediately.
        """
        plt.figure(figsize=(10, 8))

        city_coords = list(self.cities.values())
        x_coords = [coord[0] for coord in city_coords]
        y_coords = [coord[1] for coord in city_coords]

        plt.scatter(x_coords, y_coords, c='red', s=100, zorder=5)

        route_coords = [self.cities[city] for city in best_individual.route]
        route_x = [coord[0] for coord in route_coords]
        route_y = [coord[1] for coord in route_coords]

        route_x.append(route_x[0])
        route_y.append(route_y[0])

        plt.plot(route_x, route_y, 'b-', linewidth=2, alpha=0.7, label=f'Distance: {best_individual.total_distance:.3f}')

        start_coord = self.cities[best_individual.route[0]]
        plt.scatter(start_coord[0], start_coord[1], c='green', s=150, marker='s', zorder=6, label='Start')

        plt.title(title, fontsize=14, fontweight='bold')
        plt.xlabel('X', fontsize=12)
        plt.ylabel('Y', fontsize=12)
        plt.grid(True, alpha=0.3)
        plt.legend()
        plt.axis('equal')
        plt.tight_layout()
        if display:
            plt.show()

    def plot_evolution(self, fitness_history, display=True):
        """Show the evolution of the best fitness across generations.

        Args:
            fitness_history: List of fitness values per generation.
            display: If True, show the plot immediately.
        """
        plt.figure(figsize=(10, 6))
        plt.plot(fitness_history, 'b-', linewidth=2)
        plt.title('Fitness evolution across generations', fontsize=14, fontweight='bold')
        plt.xlabel('Generation', fontsize=12)
        plt.ylabel('Best fitness', fontsize=12)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        if display:
            plt.show()

    def animate_evolution(self, generations=5000, method='ox', mutation='2opt', interval=300, pause_ms=2000, repeat=True, duration=None):
        """Create an animated GIF of the genetic algorithm's evolution.

        Resets the population, runs the GA for a number of generations or a given
        duration, captures a frame at each improvement, then generates a GIF.

        Args:
            generations: Number of generations (ignored if duration is given).
            method: Crossover type to use.
            mutation: Mutation type to use.
            interval: Delay between frames in ms.
            pause_ms: Pause between animation loops in ms.
            repeat: If True, loop the animation.
            duration: Duration in seconds (overrides generations if provided).

        Returns:
            A matplotlib animation, or None if no improvement occurred.
        """
        self.individuals = []
        self.generate_individuals()

        frames_data = []
        previous_best = None
        if duration is not None:
            start = time.time()
            generation = 0
            while time.time() - start < duration:
                self.crossover(method=method, mutation=mutation)
                current_best = max(self.individuals, key=lambda individual: individual.fitness)

                if previous_best is None or not current_best.equals(previous_best):
                    print(f"Generation {generation}: New best! Distance = {current_best.total_distance:.3f}")

                    frames_data.append({
                        'generation': generation,
                        'route': current_best.route.copy(),
                        'distance': current_best.total_distance
                    })
                    previous_best = current_best
                generation += 1
        else:
            for generation in range(generations):
                self.crossover(method=method, mutation=mutation)
                current_best = max(self.individuals, key=lambda individual: individual.fitness)

                if previous_best is None or not current_best.equals(previous_best):
                    print(f"Generation {generation}: New best! Distance = {current_best.total_distance:.3f}")

                    frames_data.append({
                        'generation': generation,
                        'route': current_best.route.copy(),
                        'distance': current_best.total_distance
                    })
                    previous_best = current_best

        if frames_data:
            nb_images = len(frames_data)
            print(f"Animation with {nb_images} frames, interval: {interval}ms, pause: {pause_ms}ms")

            fig, ax = plt.subplots(figsize=(10, 8))
            city_coords = list(self.cities.values())
            x_coords = [coord[0] for coord in city_coords]
            y_coords = [coord[1] for coord in city_coords]

            xmin, xmax = min(x_coords), max(x_coords)
            ymin, ymax = min(y_coords), max(y_coords)
            dx, dy = xmax - xmin, ymax - ymin
            pad_x = max(dx * 0.1, 0.05)
            pad_y = max(dy * 0.1, 0.05)
            ax.set_xlim(xmin - pad_x, xmax + pad_x)
            ax.set_ylim(ymin - pad_y, ymax + pad_y)
            ax.set_aspect('equal')
            ax.grid(True, alpha=0.3)

            ax.scatter(x_coords, y_coords, c='red', s=100, zorder=5)

            first = frames_data[0]
            route_coords = [self.cities[city] for city in first['route']]
            route_x = [c[0] for c in route_coords] + [route_coords[0][0]]
            route_y = [c[1] for c in route_coords] + [route_coords[0][1]]
            (line_path,) = ax.plot(route_x, route_y, 'b-', linewidth=2, alpha=0.7)

            start_coord = self.cities[first['route'][0]]
            start_scatter = ax.scatter([start_coord[0]], [start_coord[1]], c='green', s=150, marker='s', zorder=6)

            info_text = ax.text(0.02, 0.98, f"Gen {first['generation']} - Distance: {first['distance']:.3f}",
                                transform=ax.transAxes, va='top', ha='left', fontsize=12,
                                bbox=dict(boxstyle='round,pad=0.3', fc='white', ec='gray', alpha=0.7))

            def get_path_xy(route):
                coords = [self.cities[city] for city in route]
                xs = [c[0] for c in coords] + [coords[0][0]]
                ys = [c[1] for c in coords] + [coords[0][1]]
                return xs, ys

            def update(frame_idx):
                frame = frames_data[frame_idx]
                xs, ys = get_path_xy(frame['route'])
                line_path.set_data(xs, ys)

                sc_x, sc_y = self.cities[frame['route'][0]]
                start_scatter.set_offsets([[sc_x, sc_y]])

                info_text.set_text(f"Gen {frame['generation']} - Distance: {frame['distance']:.3f}")

                return line_path, start_scatter, info_text

            anim = animation.FuncAnimation(
                fig,
                update,
                frames=nb_images,
                interval=interval,
                blit=True,
                repeat=repeat,
                repeat_delay=pause_ms,
            )

            try:
                os.makedirs("outputs", exist_ok=True)
                output_path = "outputs/evolution_tsp.gif"
                anim.save(output_path, writer='pillow')
                print(f"Animation saved: {output_path}")
            except Exception as e:
                print(f"Error while saving: {e}")

            return anim
        else:
            print("No improvement detected!")
            return None

    def __str__(self):
        """Text representation of the population (cities and individuals).

        Returns:
            str: Description of the population with cities and individual routes.
        """
        return f"Cities: {self.cities}\nIndividuals: {[individual.route for individual in self.individuals]}"