#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# IMPORTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

from __future__ import annotations;
from src.thirdparty.typing import *;
from src.thirdparty.maths import *;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# EXPORTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

__all__ = [
    'tsp_naive_algorithm',
];

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# METHOD tsp_naive_algorithm
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def tsp_naive_algorithm(
    dist: NDArray[(Any, Any), float],
    optimise = min,
    verbose: bool = False,
) -> tuple[float, list[list[int]]]:
    m, n = dist.shape[:2];
    assert m == n;
    memory: Dict[tuple[int, tuple], tuple[float, list[list[int]]]] = dict();

    def g(i: int, S: list[int]) -> tuple[float, list[list[int]]]:
        # wenn g bereits für den Input definiert ist, gib diesen Wert zurück:
        if (i, tuple(S)) not in memory.keys():
            if len(S) == 0:
                paths = [[i]] if i == 0 else [[i, 0]];
                memory[(i, tuple(S))] = (dist[i,0], paths);
            else:
                values_and_paths = [ (j, *g(j, (*S[:index], *S[(index+1):]))) for index, j in enumerate(S) ];
                # berechne d(i,j) + g(j, S \ {i}) for each j in S:
                values_and_paths = [ (j, dist[i,j] + value, paths) for  j, value, paths in values_and_paths];
                value = optimise([value for _, value, _  in values_and_paths]);
                paths = [];
                for j, value_, paths_ in values_and_paths:
                    if value_ == value:
                        paths += [ [i, *path] for path in paths_ ];
                memory[(i, tuple(S))] = (value, paths);
        return memory[(i, tuple(S))];

    # berechne g(0, {1,2,...,n-1}):
    optimal_wert = g(0, [i for i in range(1,n)]);

    if verbose:
        display_computation(n, memory);

    return optimal_wert, [];

def display_computation(n: int, memory: Dict[tuple[int, tuple], tuple[float, list[list[int]]]]):
    keys = sorted(memory.keys(), key=lambda key: (len(key[1]), key[0], key[1]));
    addone = lambda x: x + 1;
    for k in range(0,n):
        print(f'\x1b[4;1m|S| = {k}:\x1b[0m');
        for (i, S) in keys:
            if len(S) != k:
                continue;
            value, paths = memory[(i, S)];
            print(f'g({addone(i)}, {list(map(addone, S))}) = {value}');
            if len(paths) == 1:
                print(f'optimal way: {" -> ".join(map(str, map(addone, paths[0])))}');
            else:
                print('optimal ways:');
                for path in paths:
                    print(f'* {" -> ".join(map(str, map(addone, path)))}');
            print('');
    return;