import math
import random
from dataclasses import dataclass
from typing import List, Sequence, Tuple

import matplotlib.pyplot as plt

City = Tuple[float, float]
Tour = List[int]


def euclidean_distance(c1: City, c2: City) -> float:
    return math.hypot(c1[0] - c2[0], c1[1] - c2[1])


def build_distance_matrix(cities: Sequence[City]) -> list[list[float]]:
    size = len(cities)
    matrix = [[0.0 for _ in range(size)] for _ in range(size)]
    for i in range(size):
        for j in range(i + 1, size):
            dist = euclidean_distance(cities[i], cities[j])
            matrix[i][j] = matrix[j][i] = dist
    return matrix


def plot_tour(cities: Sequence[City], tour: Sequence[int], save_path: str) -> None:
    x = [cities[i][0] for i in tour]
    y = [cities[i][1] for i in tour]

    fig, ax = plt.subplots(figsize=(7, 7))
    ax.plot(x + [x[0]], y + [y[0]], "k-", linewidth=1)
    ax.plot(x, y, "ro", markersize=4)

    ax.axis("equal")
    fig.tight_layout()
    fig.savefig(save_path, dpi=220)
    plt.close(fig)


def plot_history(best_lengths: Sequence[float], save_path: str) -> None:
    if not best_lengths:
        return

    iterations = list(range(len(best_lengths)))

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(iterations, best_lengths, linewidth=2, color="blue")

    ax.set_xlabel("Итерация", fontsize=12)
    ax.set_ylabel("Длина лучшего тура", fontsize=12)
    ax.grid(True, alpha=0.3)

    fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


@dataclass
class ACOConfig:
    cities: Sequence[City]
    n_ants: int
    n_iterations: int
    alpha: float = 1.0
    beta: float = 5.0
    rho: float = 0.5
    q: float = 1.0
    seed: int | None = None


@dataclass
class ACOResult:
    best_tour: Tour
    best_length: float
    history: List[float]


class AntColonyOptimizer:
    def __init__(self, config: ACOConfig):
        self.config = config
        if config.seed is not None:
            random.seed(config.seed)

        self.cities = config.cities
        self.dist_matrix = build_distance_matrix(config.cities)
        n = len(config.cities)
        self.pheromone = [[1.0 if i != j else 0.0 for j in range(n)] for i in range(n)]

    def _choose_next_city(self, current: int, unvisited: set[int]) -> int:
        candidates = list(unvisited)
        weights = []
        for nxt in candidates:
            tau = self.pheromone[current][nxt] ** self.config.alpha
            eta = (1.0 / (self.dist_matrix[current][nxt] + 1e-12)) ** self.config.beta
            weights.append(tau * eta)

        total = sum(weights)
        probs = [w / total for w in weights]
        return random.choices(candidates, weights=probs, k=1)[0]

    def _build_tour(self, start: int) -> Tour:
        n = len(self.cities)
        tour = [start]
        unvisited = set(range(n))
        unvisited.remove(start)

        current = start
        while unvisited:
            nxt = self._choose_next_city(current, unvisited)
            tour.append(nxt)
            unvisited.remove(nxt)
            current = nxt

        return tour

    def _tour_length(self, tour: Sequence[int]) -> float:
        return sum(
            self.dist_matrix[tour[i]][tour[(i + 1) % len(tour)]]
            for i in range(len(tour))
        )

    def run(self) -> ACOResult:
        best_tour: Tour = []
        best_length = float("inf")
        best_history: list[float] = []

        for _ in range(self.config.n_iterations):
            tours: list[Tour] = []
            lengths: list[float] = []

            for _ in range(self.config.n_ants):
                start_city = random.randrange(len(self.cities))
                tour = self._build_tour(start_city)
                length = self._tour_length(tour)
                tours.append(tour)
                lengths.append(length)

                if length < best_length:
                    best_length = length
                    best_tour = tour

            for i in range(len(self.pheromone)):
                for j in range(len(self.pheromone)):
                    self.pheromone[i][j] *= 1 - self.config.rho

            for tour, length in zip(tours, lengths):
                deposit = self.config.q / length
                for i in range(len(tour)):
                    a, b = tour[i], tour[(i + 1) % len(tour)]
                    self.pheromone[a][b] += deposit
                    self.pheromone[b][a] += deposit

            best_history.append(best_length)

        return ACOResult(
            best_tour=best_tour, best_length=best_length, history=best_history
        )


def run_aco(config: ACOConfig) -> ACOResult:
    optimizer = AntColonyOptimizer(config)
    return optimizer.run()