import copy
from typing import Tuple

import numpy as np
import numpy.typing as npt

from common import indexes_to_cities, plot_plan, read_data


def hill_climbing(distances: npt.NDArray) -> Tuple[float, npt.NDArray]:
    """A simple hill climbing algorithm.

    The algorithm starts on a random permutation and attempts to improve
    the circuit by trying to switch neighboring elements. Each iteration
    tries to switch adjacent neighbors and sees which one yields the largest
    improvement.

    Args:
        distances npt.NDArray: A matrix containing the distances between cities.

    Returns:
        Tuple[float, npt.NDArray] A tuple containing the distance of the
        solution and the solution itself.

    """

    size: int = len(distances)  # The size of the permutation array
    perm: npt.NDArray = np.arange(size)  # Create an array from 0..size

    np.random.shuffle(perm)  # Get random permutation

    # Get the distance of the random permutation
    current_distance: float = np.sum(
        [distances[perm[i - 1], perm[i]] for i in range(size)]
    )

    found_improvement: bool = True

    while found_improvement:

        found_improvement = False  # Assume we haven't found an improvement
        tmp_distance: float = current_distance

        # Try to find an improvement
        for i in range(size):
            perm[[i - 1, i]] = perm[[i, i - 1]]  # Swap i - 1 and i

            tmp_distance: float = np.sum(
                [distances[perm[i - 1], perm[i]] for i in range(size)]
            )

            if tmp_distance < current_distance:
                current_distance = tmp_distance
                found_improvement = True
                break

            perm[[i - 1, i]] = perm[[i, i - 1]]  # Swap back i - 1 and i

    return (current_distance, perm)


def test_hill_climbing(data: npt.NDArray, cities: npt.NDArray, runs: int):
    res = [hill_climbing(data) for _ in range(runs)]
    res.sort(key=lambda n: n[0])

    distances = list(map(lambda n: n[0], res))
    best = res[0][0]
    worst = res[-1][0]
    avg = sum(distances) / runs
    standard_deviation = np.sqrt(sum([(i - avg)**2 for i in distances]) / runs)

    print(f"Hill climbing for {len(data)} cities.")
    print(f"best distance     : {best:>12.6f}km")
    print(f"worst distance    : {worst:>12.6f}km")
    print(f"average distance  : {avg:>12.6f}km")
    print(f"standard deviation: {standard_deviation:>12.6f}km\n")

    plot_plan(indexes_to_cities(res[0][1], cities)) # Plot the best one

if __name__ == "__main__":
    np.random.seed(1987)
    cities, data = read_data("./european_cities.csv")
    distance, perm = hill_climbing(data[:10, :10])

    # plot_plan(indexes_to_cities(perm, cities))
    test_hill_climbing(data[:10,:10], cities, 20)
    test_hill_climbing(data, cities, 20)