namespace AntColony.Algorithm;

/// <summary>
/// Алгоритм построения кратчайшего пути с помощью симуляции муравьиной колонии
/// </summary>
public class AntColonyOptimizer
{
    private readonly int numberOfCities;
    private readonly double[][] distances;
    private readonly double[][] pheromones;
    private readonly double alpha; // Важность феромонов
    private readonly double beta;  // Важность расстояния 
    private readonly double evaporationRate;
    private readonly double Q;     // Множитель депозита феромонов
    private readonly int numberOfAnts;

    /// <summary>
    /// Инициализировать алгоритм
    /// </summary>
    /// <param name="distances">Матрица расстояний</param>
    /// <param name="numberOfAnts">Количество муравьев</param>
    /// <param name="alpha">Значение α (определяет вес феромонов)</param>
    /// <param name="beta">Значение β (определяет вес расстояния)</param>
    /// <param name="evaporationRate">Скорость испарения</param>
    /// <param name="Q">Значение Q</param>
    public AntColonyOptimizer(double[][] distances, int numberOfAnts = 20, double alpha = 1, double beta = 2,
        double evaporationRate = 0.1, double Q = 100)
    {
        numberOfCities = distances.GetLength(0);
        this.distances = distances;
        this.numberOfAnts = numberOfAnts;
        this.alpha = alpha;
        this.beta = beta;
        this.evaporationRate = evaporationRate;
        this.Q = Q;

        // Строим стартовую матрицу феромонов
        pheromones = new double[numberOfCities][];
        for (int i = 0; i < numberOfCities; i++)
        {
            var r = new double[numberOfCities];
            for (int j = 0; j < numberOfCities; j++)
            {
                if (i != j)
                    r[j] = 1.0; // Начальное значение
            }
            pheromones[i] = r;
        }
    }

    /// <summary>
    /// Решить задачу комивояжера (построить Гамильтонов граф)
    /// </summary>
    /// <param name="maxIterations">Максимальное число итераций</param>
    /// <returns>Лучший путь и его расстояние</returns>
    public (List<int> bestTour, double bestDistance) Solve(int maxIterations)
    {
        List<int> bestTourEver = [];
        double bestDistanceEver = double.MaxValue;
        Lock distanceLock = new();
        Lock pheromoneLock = new();

        // for (int iteration = 0; iteration < maxIterations; iteration++)
        Parallel.For(0, maxIterations, i =>
        {
            List<List<int>> antTours = [];
            List<double> antDistances = [];

            // Каждый муравей строит путь
            for (int ant = 0; ant < numberOfAnts; ant++)
            {
                var ct = ConstructTour();
                if (ct is null) continue;
                var (tour, distance) = ct.Value;
                antTours.Add(tour);
                antDistances.Add(distance);

                // Если текущий путь лучше - обновляем
                lock (distanceLock)
                {
                    if (distance < bestDistanceEver)
                    {
                        bestDistanceEver = distance;
                        bestTourEver = [.. tour];
                    }
                }
            }

            // Обновляем феромоны для всех путей
            lock (pheromoneLock)
            {
                UpdatePheromones(antTours, antDistances);
            }
            // File.AppendAllText("./iteration-results", bestDistanceEver.ToString() + '\n');
            // }
        });

        return (bestTourEver, bestDistanceEver);
    }

    private (List<int> tour, double distance)? ConstructTour()
    {
        List<int> tour = [];
        var visited = new bool[numberOfCities];
        var currentCity = Random.Shared.Next(numberOfCities);

        tour.Add(currentCity);
        visited[currentCity] = true;

        // Строим путь
        while (tour.Count < numberOfCities)
        {
            var sel = SelectNextCity(currentCity, visited);
            if (sel is not { } nextCity) return null;
            tour.Add(nextCity);
            visited[nextCity] = true;
            currentCity = nextCity;
        }

        if (distances[currentCity][tour[0]] == 0) return null;
        // Добавляем возврат на начальную точку
        tour.Add(tour[0]);

        return (tour, CalculateTourDistance(tour));
    }

    /// <summary>
    /// Выбрать следующий город
    /// </summary>
    /// <param name="currentCity">В каком городе муравей находится сейчас</param>
    /// <param name="visited">Какие города муравей уже посетил</param>
    /// <returns>Номер следующего города</returns>
    private int? SelectNextCity(int currentCity, bool[] visited)
    {
        var probabilities = new double[numberOfCities];
        double probabilitiesSum = 0;

        // Вероятности для каждого непосещенного города
        for (int i = 0; i < numberOfCities; i++)
        {
            if (!visited[i])
            {
                if (distances[currentCity][i] == 0) continue;
                probabilities[i] = Math.Pow(pheromones[currentCity][i], alpha) *
                                 Math.Pow(1.0 / distances[currentCity][i], beta);
                probabilitiesSum += probabilities[i];
            }
        }

        // Случайно выбираем следующий город (по вероятностям)
        double randomValue = Random.Shared.NextDouble() * probabilitiesSum;
        //double sum = 0;

        for (int i = 0; i < numberOfCities; i++)
        {
            if (!visited[i])
            {
                //sum += probabilities[i];
                if (probabilities[i] >= randomValue)
                    return i;
            }
        }

        return null;
    }

    /// <summary>
    /// Обновить матрицу феромонов
    /// </summary>
    /// <param name="antTours">Пути муравьев</param>
    /// <param name="antDistances">Их расстояния</param>
    private void UpdatePheromones(List<List<int>> antTours, List<double> antDistances)
    {
        // Испарение
        for (int i = 0; i < numberOfCities; i++)
        {
            for (int j = 0; j < numberOfCities; j++)
            {
                if (i != j)
                    pheromones[i][j] *= 1.0 - evaporationRate;
            }
        }

        // Добавляем новые феромоны
        for (int ant = 0; ant < antTours.Count; ant++)
        {
            double pheromoneDeposit = Q / antDistances[ant];
            var tour = antTours[ant];

            for (int i = 0; i < tour.Count - 1; i++)
            {
                int city1 = tour[i];
                int city2 = tour[i + 1];
                pheromones[city1][city2] += pheromoneDeposit;
                pheromones[city2][city1] += pheromoneDeposit; // Граф ненаправленный
            }
        }
    }

    /// <summary>
    /// Рассчитать длину пути
    /// </summary>
    /// <param name="tour">Путь</param>
    /// <returns>Длина пути</returns>
    private double CalculateTourDistance(List<int> tour)
    {
        double distance = 0;
        for (int i = 0; i < tour.Count - 1; i++)
        {
            distance += distances[tour[i]][tour[i + 1]];
        }
        return distance;
    }
}