import time
from itertools import permutations
from typing import Tuple
from math import factorial

import numpy as np
import numpy.typing as npt

from common import indexes_to_cities, plot_plan, read_data


def exhaustive_search(distances: npt.NDArray) -> Tuple[float, npt.NDArray]:
    """An implementation of exhaustive search.

        This implementation takes a permutation iterator, then maps each
        element to a tuple containing the travel distance of that permutation
        and the permutation tuple itself. Finally the function returns the
        tuple that contains the smallest travel distance.

    Args:
        distances (npt.NDArray): An array containing distances to travel.

    Returns:
        Tuple[float, npt.NDArray] A tuple containing the shortest travel 
        distance and its corresponding permutation.
    """

    size = len(distances)

    return min(  # Find the smallest travel distance from the array
        map(  # Map the permutation array to contain tuples (distance, permutation)
            lambda perm: (
                sum([distances[perm[i - 1], perm[i]] for i in range(size)]),
                np.array(perm),
            ),
            permutations(range(size)),
        ),
        key=lambda x: x[0] # Make sure that it finds the minimal distance
    )


if __name__ == "__main__":
    cities, data = read_data("./european_cities.csv")

    times = {}

    # A loop timing finding the optimal solution for different n
    for n in range(6,11):
        # Time exhaustive search
        t0 = time.time_ns()
        distance, perm = exhaustive_search(data[:n, :n])
        t1 = time.time_ns()

        time_elapsed_ms = (t1 - t0) / 1_000_000.0

        times[n] = time_elapsed_ms, distance

        if n in (6,10):
            city_seq = indexes_to_cities(perm, cities)
            plot_plan(city_seq)
            print(f"Sequence for {n} cities: {city_seq}")

    print("")
    for n, (time, distance) in times.items():
        print(f"Exhaustive search for the {n} first cities:")
        print(f"{'distance':<25}: {distance:>12.6f}km")
        print(f"{'time to find solution':<25}: {time:>12.6f}ms")
        print(f"{f'time / {n}!':<25}: {time / factorial(n):>12.6f}\n")


"""Running example

oblig1 on  main [!] via 🐍 v3.12.6 took 14s
❯ python exhaustive_search.py
Exhaustive search for the 6 first cities:
distance                 :  5018.810000km
time to find solution    :     1.485208ms
time / 6!                :     0.002063

Exhaustive search for the 10 first cities:
distance                 :  7486.310000km
time to find solution    : 10980.900480ms
time / 10!               :     0.003026
"""